From 905a4a04ad51b0922d36d01fd5c6fd509b2ff7af Mon Sep 17 00:00:00 2001 From: bxb1337 Date: Fri, 9 Oct 2026 19:46:26 +0800 Subject: [PATCH 1/4] =?UTF-8?q?feat(settings):=20=E9=85=8D=E7=BD=AE?= =?UTF-8?q?=E8=81=8A=E5=A4=A9=E7=BB=AD=E9=93=BE=E4=B8=8E=E7=A9=BA=E7=99=BD?= =?UTF-8?q?=E9=A2=84=E7=83=AD=E8=BF=9E=E6=8E=A5=E7=8B=AC=E7=AB=8B=E4=BF=9D?= =?UTF-8?q?=E7=95=99=E6=95=B0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- admin/handler.go | 12 ++++++ auth/codex_ws_settings.go | 18 +++++++++ auth/store.go | 35 +++++++++-------- database/codex_ws_settings.go | 17 +++++++++ database/postgres.go | 66 +++++++++++++++++++-------------- database/runtime_cache_scope.go | 20 ++++++++++ database/sqlite.go | 2 + frontend/src/locales/en.json | 6 ++- frontend/src/locales/zh.json | 6 ++- frontend/src/pages/Settings.tsx | 18 ++++++++- frontend/src/types.ts | 1 + proxy/runtime_config.go | 11 ++++-- 12 files changed, 160 insertions(+), 52 deletions(-) create mode 100644 auth/codex_ws_settings.go create mode 100644 database/codex_ws_settings.go create mode 100644 database/runtime_cache_scope.go diff --git a/admin/handler.go b/admin/handler.go index 5a53f69c0..74138f80d 100644 --- a/admin/handler.go +++ b/admin/handler.go @@ -9262,6 +9262,7 @@ type settingsResponse struct { CodexWSBusyOverflowEnabled bool `json:"codex_ws_busy_overflow_enabled"` CodexWSBusyPatienceSec int `json:"codex_ws_busy_patience_sec"` CodexWSStatelessSlots int `json:"codex_ws_stateless_slots"` + CodexWSDownstreamKeepaliveSlots int `json:"codex_ws_downstream_keepalive_slots"` GithubTokenConfigured bool `json:"github_token_configured"` GithubProxyURL string `json:"github_proxy_url"` CodexOverloadPauseEnabled bool `json:"codex_overload_pause_enabled"` @@ -9457,6 +9458,7 @@ type updateSettingsReq struct { CodexWSBusyOverflowEnabled *bool `json:"codex_ws_busy_overflow_enabled"` CodexWSBusyPatienceSec *int `json:"codex_ws_busy_patience_sec"` CodexWSStatelessSlots *int `json:"codex_ws_stateless_slots"` + CodexWSDownstreamKeepaliveSlots *int `json:"codex_ws_downstream_keepalive_slots"` GithubToken *string `json:"github_token"` GithubProxyURL *string `json:"github_proxy_url"` CodexOverloadPauseEnabled *bool `json:"codex_overload_pause_enabled"` @@ -10294,6 +10296,7 @@ func (h *Handler) GetSettings(c *gin.Context) { CodexWSBusyOverflowEnabled: h.store.CodexWSBusyOverflowEnabled(), CodexWSBusyPatienceSec: h.store.CodexWSBusyPatienceSec(), CodexWSStatelessSlots: h.store.CodexWSStatelessSlots(), + CodexWSDownstreamKeepaliveSlots: h.store.CodexWSDownstreamKeepaliveSlots(), GithubTokenConfigured: h.store.GithubToken() != "", GithubProxyURL: h.store.GithubProxyURL(), CodexOverloadPauseEnabled: runtimeCfg.CodexOverloadPauseEnabled, @@ -11101,6 +11104,13 @@ func (h *Handler) UpdateSettings(c *gin.Context) { log.Printf("设置已更新: codex_ws_stateless_slots = %d", v) } + if req.CodexWSDownstreamKeepaliveSlots != nil { + v := database.NormalizeCodexWSDownstreamKeepaliveSlots(*req.CodexWSDownstreamKeepaliveSlots) + h.store.SetCodexWSDownstreamKeepaliveSlots(v) + runtimeCfg.CodexWSDownstreamKeepaliveSlots = v + log.Printf("设置已更新: codex_ws_downstream_keepalive_slots = %d", v) + } + // GitHub 访问设置（issue #522）。token 不回显也不落日志，空串表示清除。 if req.GithubToken != nil { v := strings.TrimSpace(*req.GithubToken) @@ -11840,6 +11850,7 @@ func (h *Handler) UpdateSettings(c *gin.Context) { CodexWSBusyOverflowEnabled: h.store.CodexWSBusyOverflowEnabled(), CodexWSBusyPatienceSec: h.store.CodexWSBusyPatienceSec(), CodexWSStatelessSlots: h.store.CodexWSStatelessSlots(), + CodexWSDownstreamKeepaliveSlots: h.store.CodexWSDownstreamKeepaliveSlots(), GithubToken: h.store.GithubToken(), GithubProxyURL: h.store.GithubProxyURL(), CodexOverloadPauseEnabled: runtimeCfg.CodexOverloadPauseEnabled, @@ -12177,6 +12188,7 @@ func (h *Handler) UpdateSettings(c *gin.Context) { CodexWSBusyOverflowEnabled: h.store.CodexWSBusyOverflowEnabled(), CodexWSBusyPatienceSec: h.store.CodexWSBusyPatienceSec(), CodexWSStatelessSlots: h.store.CodexWSStatelessSlots(), + CodexWSDownstreamKeepaliveSlots: h.store.CodexWSDownstreamKeepaliveSlots(), GithubTokenConfigured: h.store.GithubToken() != "", GithubProxyURL: h.store.GithubProxyURL(), CodexOverloadPauseEnabled: runtimeCfg.CodexOverloadPauseEnabled, diff --git a/auth/codex_ws_settings.go b/auth/codex_ws_settings.go new file mode 100644 index 000000000..a66a5af0a --- /dev/null +++ b/auth/codex_ws_settings.go @@ -0,0 +1,18 @@ +package auth + +import "github.com/codex2api/database" + +// SetCodexWSDownstreamKeepaliveSlots 设置每账号独立的空闲聊天连接保留上限。 +func (s *Store) SetCodexWSDownstreamKeepaliveSlots(slots int) { + if s != nil { + s.codexWSDownstreamKeepaliveSlots.Store(int64(database.NormalizeCodexWSDownstreamKeepaliveSlots(slots))) + } +} + +// CodexWSDownstreamKeepaliveSlots 返回每账号独立的空闲聊天连接保留上限。 +func (s *Store) CodexWSDownstreamKeepaliveSlots() int { + if s == nil { + return database.DefaultCodexWSDownstreamKeepaliveSlots + } + return database.NormalizeCodexWSDownstreamKeepaliveSlots(int(s.codexWSDownstreamKeepaliveSlots.Load())) +} diff --git a/auth/store.go b/auth/store.go index 6507b9df0..58c950c89 100644 --- a/auth/store.go +++ b/auth/store.go @@ -3611,20 +3611,21 @@ type Store struct { codexForceWebsocket atomic.Bool // 强制 Codex 上游走 WebSocket（复用连接池） // codexRequestCompression HTTP /responses 请求体 zstd 压缩，默认开启（对齐真实客户端）。 // 与上面几项 WS 设置正交：WS 走 permessage-deflate，本项只作用于 HTTP 路径。 - codexRequestCompression atomic.Bool - codexWSKeepaliveEnabled atomic.Bool // 启用上游 WS 空闲连接保活（仅 Ping） - codexWSKeepaliveIntervalSec atomic.Int64 // WS 保活 Ping 间隔（秒），默认 60 - codexWSHideUpstreamErrors atomic.Bool // 隐藏上游 WS 原始错误，默认开启 - codexWSSilentRetryEnabled atomic.Bool // 首包前上游 WS 错误静默换号重试，默认开启 - codexWSSilentMaxRetries atomic.Int64 // WS 静默换号最大重试次数，默认 2 - codexWSSizeRouterEnabled atomic.Bool // 1009 自学习体积路由，默认开启 - codexWSBusyMaxWaitSec atomic.Int64 // busy session 等待上限（秒），默认 30（issue #413） - codexWSBusyOverflowEnabled atomic.Bool // busy session 溢出到同账号兄弟连接，默认关闭 - codexWSBusyPatienceSec atomic.Int64 // 触发溢出前的短等待（秒），默认 2 - codexWSStatelessSlots atomic.Int64 // 无状态请求每 (账号, cacheKey) 的连接槽位数，默认 8（issue #522） - overflowAutoCompactEnabled atomic.Bool // 上下文超窗自动摘要重试（实验性，默认关闭，issue #415） - compactViaResponsesEnabled atomic.Bool // /v1/responses/compact 改写为 /responses body-signal 压缩，默认关闭 - firstTokenExcludesWsAcquire atomic.Bool // 落库 first_token_ms 扣除 WS 取连耗时，默认关闭 + codexRequestCompression atomic.Bool + codexWSKeepaliveEnabled atomic.Bool // 启用上游 WS 空闲连接保活（仅 Ping） + codexWSKeepaliveIntervalSec atomic.Int64 // WS 保活 Ping 间隔（秒），默认 60 + codexWSHideUpstreamErrors atomic.Bool // 隐藏上游 WS 原始错误，默认开启 + codexWSSilentRetryEnabled atomic.Bool // 首包前上游 WS 错误静默换号重试，默认开启 + codexWSSilentMaxRetries atomic.Int64 // WS 静默换号最大重试次数，默认 2 + codexWSSizeRouterEnabled atomic.Bool // 1009 自学习体积路由，默认开启 + codexWSBusyMaxWaitSec atomic.Int64 // busy session 等待上限（秒），默认 30（issue #413） + codexWSBusyOverflowEnabled atomic.Bool // busy session 溢出到同账号兄弟连接，默认关闭 + codexWSBusyPatienceSec atomic.Int64 // 触发溢出前的短等待（秒），默认 2 + codexWSStatelessSlots atomic.Int64 // 每账号空白上游 WS 槽位预算，默认 8 + codexWSDownstreamKeepaliveSlots atomic.Int64 + overflowAutoCompactEnabled atomic.Bool // 上下文超窗自动摘要重试（实验性，默认关闭，issue #415） + compactViaResponsesEnabled atomic.Bool // /v1/responses/compact 改写为 /responses body-signal 压缩，默认关闭 + firstTokenExcludesWsAcquire atomic.Bool // 落库 first_token_ms 扣除 WS 取连耗时，默认关闭 // 前置元数据 SSE 事件立即透传下游（旧版兼容，默认关闭，issue #425） codexPreflightSSEPassthroughEnabled atomic.Bool @@ -4212,6 +4213,7 @@ func NewStore(db *database.DB, tc cache.TokenCache, settings *database.SystemSet CodexWSBusyAcquireMaxWaitSec: 30, CodexWSBusyPatienceSec: 2, CodexWSStatelessSlots: 8, + CodexWSDownstreamKeepaliveSlots: database.DefaultCodexWSDownstreamKeepaliveSlots, CodexContinueMaxRounds: 8, AutoPause5hGuardBandPercent: defaultAutoPause5hGuardBandPercent, AutoPause5hGuardConcurrency: defaultAutoPause5hGuardConcurrency, @@ -4332,6 +4334,7 @@ func NewStore(db *database.DB, tc cache.TokenCache, settings *database.SystemSet s.codexWSBusyOverflowEnabled.Store(settings.CodexWSBusyOverflowEnabled) s.codexWSBusyPatienceSec.Store(int64(database.NormalizeCodexWSBusyPatienceSec(settings.CodexWSBusyPatienceSec))) s.codexWSStatelessSlots.Store(int64(database.NormalizeCodexWSStatelessSlots(settings.CodexWSStatelessSlots))) + s.SetCodexWSDownstreamKeepaliveSlots(settings.CodexWSDownstreamKeepaliveSlots) s.overflowAutoCompactEnabled.Store(settings.OverflowAutoCompactEnabled) s.compactViaResponsesEnabled.Store(settings.CompactViaResponsesEnabled) s.codexPreflightSSEPassthroughEnabled.Store(settings.CodexPreflightSSEPassthroughEnabled) @@ -4780,7 +4783,7 @@ func (s *Store) GithubProxyURL() string { return "" } -// SetCodexWSStatelessSlots 设置无状态请求每 (账号, cacheKey) 的连接槽位数。 +// SetCodexWSStatelessSlots 设置每个账号的空白上游 WS 连接槽位预算。 func (s *Store) SetCodexWSStatelessSlots(slots int) { if s == nil { return @@ -4788,7 +4791,7 @@ func (s *Store) SetCodexWSStatelessSlots(slots int) { s.codexWSStatelessSlots.Store(int64(database.NormalizeCodexWSStatelessSlots(slots))) } -// CodexWSStatelessSlots 返回无状态请求每 (账号, cacheKey) 的连接槽位数。 +// CodexWSStatelessSlots 返回每个账号的空白上游 WS 连接槽位预算。 func (s *Store) CodexWSStatelessSlots() int { if s == nil { return 8 diff --git a/database/codex_ws_settings.go b/database/codex_ws_settings.go new file mode 100644 index 000000000..2eb07b395 --- /dev/null +++ b/database/codex_ws_settings.go @@ -0,0 +1,17 @@ +package database + +const ( + DefaultCodexWSDownstreamKeepaliveSlots = 8 + MaxCodexWSDownstreamKeepaliveSlots = 32 +) + +// NormalizeCodexWSDownstreamKeepaliveSlots 归一化空闲聊天连接保留上限，0 表示不保留。 +func NormalizeCodexWSDownstreamKeepaliveSlots(slots int) int { + if slots < 0 { + return DefaultCodexWSDownstreamKeepaliveSlots + } + if slots > MaxCodexWSDownstreamKeepaliveSlots { + return MaxCodexWSDownstreamKeepaliveSlots + } + return slots +} diff --git a/database/postgres.go b/database/postgres.go index 20456988a..55cc90589 100644 --- a/database/postgres.go +++ b/database/postgres.go @@ -200,9 +200,10 @@ type sqlExecer interface { // DB PostgreSQL 数据库操作 type DB struct { - conn *sql.DB - driver string - authCacheScope string + conn *sql.DB + driver string + authCacheScope string + runtimeCacheScope string promptFilterAudit *promptFilterAuditQueue @@ -513,6 +514,11 @@ func New(driver string, dsn string, schema ...string) (*DB, error) { _ = conn.Close() return nil, fmt.Errorf("初始化鉴权缓存修订表失败: %w", err) } + if err := db.initializeRuntimeCacheScope(ctx); err != nil { + backgroundTaskCancel() + _ = conn.Close() + return nil, fmt.Errorf("初始化共享运行态作用域失败: %w", err) + } // 启动批量写入后台协程 db.startLogFlusher() @@ -1505,6 +1511,7 @@ func (db *DB) migrate(ctx context.Context) error { ALTER TABLE system_settings ADD COLUMN IF NOT EXISTS codex_ws_busy_overflow_enabled BOOLEAN DEFAULT FALSE; ALTER TABLE system_settings ADD COLUMN IF NOT EXISTS codex_ws_busy_patience_sec INT DEFAULT 2; ALTER TABLE system_settings ADD COLUMN IF NOT EXISTS codex_ws_stateless_slots INT DEFAULT 8; + ALTER TABLE system_settings ADD COLUMN IF NOT EXISTS codex_ws_downstream_keepalive_slots INT DEFAULT 8; ALTER TABLE system_settings ADD COLUMN IF NOT EXISTS github_token TEXT DEFAULT ''; ALTER TABLE system_settings ADD COLUMN IF NOT EXISTS github_proxy_url TEXT DEFAULT ''; ALTER TABLE system_settings ADD COLUMN IF NOT EXISTS codex_overload_pause_enabled BOOLEAN DEFAULT FALSE; @@ -2437,23 +2444,24 @@ type SystemSettings struct { ShowFullUsageNumbers bool // ShowUpstreamModelMismatch 控制用量页是否标出上游响应模型与请求模型不一致。 // 只影响展示和筛选，不停止记录。默认开启。 - ShowUpstreamModelMismatch bool - PublicKeyUsagePageEnabled bool - PublicImageStudioPageEnabled bool - PublicAccountPortalPageEnabled bool // 账号自助添加公开门户开关，默认 false - CodexForceWebsocket bool // 强制 Codex 上游走 WebSocket（复用连接池），默认 false - CodexRequestCompression bool // HTTP /responses 请求体 zstd 压缩（对齐真实客户端），默认 true - CodexWSWeakNetworkMode bool // WS 弱网保守复用模式，默认 false - CodexWSKeepaliveEnabled bool // 启用上游 WS 空闲连接保活（仅 Ping，不发业务帧），默认 false - CodexWSKeepaliveIntervalSec int // WS 保活 Ping 间隔（秒），默认 60 - CodexWSHideUpstreamErrors bool // 隐藏上游 WS 原始错误，默认 true - CodexWSSilentRetryEnabled bool // 首包前 WS 上游错误静默换号重试，默认 true - CodexWSSilentMaxRetries int // WS 静默换号最大重试次数，默认 2 - CodexWSSizeRouterEnabled bool // 1009 自学习体积路由：超大请求直接首发 HTTP，默认 true - CodexWSBusyAcquireMaxWaitSec int // busy session/容量等待的累计上限（秒），默认 30（issue #413） - CodexWSBusyOverflowEnabled bool // busy session 溢出到同账号兄弟连接，默认 false（issue #413） - CodexWSBusyPatienceSec int // 触发溢出前的短等待（秒），默认 2（issue #413） - CodexWSStatelessSlots int // 无状态请求每 (账号, cacheKey) 的持久连接槽位数，默认 8，范围 1-32（issue #522） + ShowUpstreamModelMismatch bool + PublicKeyUsagePageEnabled bool + PublicImageStudioPageEnabled bool + PublicAccountPortalPageEnabled bool // 账号自助添加公开门户开关，默认 false + CodexForceWebsocket bool // 强制 Codex 上游走 WebSocket（复用连接池），默认 false + CodexRequestCompression bool // HTTP /responses 请求体 zstd 压缩（对齐真实客户端），默认 true + CodexWSWeakNetworkMode bool // WS 弱网保守复用模式，默认 false + CodexWSKeepaliveEnabled bool // 启用上游 WS 空闲连接保活（仅 Ping，不发业务帧），默认 false + CodexWSKeepaliveIntervalSec int // WS 保活 Ping 间隔（秒），默认 60 + CodexWSHideUpstreamErrors bool // 隐藏上游 WS 原始错误，默认 true + CodexWSSilentRetryEnabled bool // 首包前 WS 上游错误静默换号重试，默认 true + CodexWSSilentMaxRetries int // WS 静默换号最大重试次数，默认 2 + CodexWSSizeRouterEnabled bool // 1009 自学习体积路由：超大请求直接首发 HTTP，默认 true + CodexWSBusyAcquireMaxWaitSec int // busy session/容量等待的累计上限（秒），默认 30（issue #413） + CodexWSBusyOverflowEnabled bool // busy session 溢出到同账号兄弟连接，默认 false（issue #413） + CodexWSBusyPatienceSec int // 触发溢出前的短等待（秒），默认 2（issue #413） + CodexWSStatelessSlots int // 每账号空白上游 WS 槽位预算，默认 8，范围 0-32 + CodexWSDownstreamKeepaliveSlots int // 每账号空闲聊天续链连接的独立保留上限，默认 8，范围 0-32 // GithubToken 用于 api.github.com 请求的 Personal Access Token（提升限流配额， // 只发给 api.github.com，绝不发给镜像/其他主机；空表示未配置，issue #522）。 GithubToken string @@ -2741,7 +2749,8 @@ func (db *DB) GetSystemSettings(ctx context.Context) (*SystemSettings, error) { COALESCE(codex_oauth_keepalive_enabled, false), COALESCE(codex_telemetry_timing_debug, false), COALESCE(auto_reset_credits_on_exhaustion_enabled, false), - COALESCE(codex_unified_client_identity_enabled, false) + COALESCE(codex_unified_client_identity_enabled, false), + COALESCE(codex_ws_downstream_keepalive_slots, 8) FROM system_settings WHERE id = 1 `).Scan( &s.SiteName, &s.SiteLogo, @@ -2831,6 +2840,7 @@ func (db *DB) GetSystemSettings(ctx context.Context) (*SystemSettings, error) { &s.CodexTelemetryTimingDebug, &s.AutoResetCreditsOnExhaustionEnabled, &s.CodexUnifiedClientIdentityEnabled, + &s.CodexWSDownstreamKeepaliveSlots, ) if errors.Is(err, sql.ErrNoRows) { return nil, nil @@ -3096,9 +3106,10 @@ func (db *DB) UpdateSystemSettings(ctx context.Context, s *SystemSettings) error codex_oauth_keepalive_enabled, codex_telemetry_timing_debug, auto_reset_credits_on_exhaustion_enabled, - codex_unified_client_identity_enabled + codex_unified_client_identity_enabled, + codex_ws_downstream_keepalive_slots ) - VALUES (1, $1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23, $24, $25, $26, $27, $28, $29, $30, $31, $32, $33, $34, $35, $36, $37, $38, $39, $40, $41, $42, $43, $44, $45, $46, $47, $48, $49, $50, $51, $52, $53, $54, $55, $56, $57, $58, $59, $60, $61, $62, $63, $64, $65, $66, $67, $68, $69, $70, $71, $72, $73, $74, $75, $76, $77, $78, $79, $80, $81, $82, $83, $84, $85, $86, $87, $88, $89, $90, $91, $92, $93, $94, $95, $96, $97, $98, $99, $100, $101, $102, $103, $104, $105, $106, $107, $108, $109, $110, $111, $112, $113, $114, $115, $116, $117, $118, $119, $120, $121, $122, $123, $124, $125, $126, $127, $128) + VALUES (1, $1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23, $24, $25, $26, $27, $28, $29, $30, $31, $32, $33, $34, $35, $36, $37, $38, $39, $40, $41, $42, $43, $44, $45, $46, $47, $48, $49, $50, $51, $52, $53, $54, $55, $56, $57, $58, $59, $60, $61, $62, $63, $64, $65, $66, $67, $68, $69, $70, $71, $72, $73, $74, $75, $76, $77, $78, $79, $80, $81, $82, $83, $84, $85, $86, $87, $88, $89, $90, $91, $92, $93, $94, $95, $96, $97, $98, $99, $100, $101, $102, $103, $104, $105, $106, $107, $108, $109, $110, $111, $112, $113, $114, $115, $116, $117, $118, $119, $120, $121, $122, $123, $124, $125, $126, $127, $128, $131) ON CONFLICT (id) DO UPDATE SET site_name = EXCLUDED.site_name, site_logo = EXCLUDED.site_logo, @@ -3224,7 +3235,8 @@ func (db *DB) UpdateSystemSettings(ctx context.Context, s *SystemSettings) error codex_oauth_keepalive_enabled = EXCLUDED.codex_oauth_keepalive_enabled, codex_telemetry_timing_debug = EXCLUDED.codex_telemetry_timing_debug, auto_reset_credits_on_exhaustion_enabled = EXCLUDED.auto_reset_credits_on_exhaustion_enabled, - codex_unified_client_identity_enabled = EXCLUDED.codex_unified_client_identity_enabled + codex_unified_client_identity_enabled = EXCLUDED.codex_unified_client_identity_enabled, + codex_ws_downstream_keepalive_slots = EXCLUDED.codex_ws_downstream_keepalive_slots `, NormalizeSiteName(s.SiteName), strings.TrimSpace(s.SiteLogo), s.MaxConcurrency, s.GlobalRPM, s.TestModel, testContent, s.TestConcurrency, s.ProxyURL, s.PgMaxConns, s.RedisPoolSize, s.AutoCleanUnauthorized, s.AutoCleanRateLimited, s.AdminSecret, s.AutoCleanFullUsage, s.ProxyPoolEnabled, @@ -3283,7 +3295,7 @@ func (db *DB) UpdateSystemSettings(ctx context.Context, s *SystemSettings) error s.AutoResetCreditsOnExhaustionEnabled, s.CodexUnifiedClientIdentityEnabled, s.PreservePromptFilterCustomPatterns, - s.PreservePromptFilterReviewAPIKey) + s.PreservePromptFilterReviewAPIKey, NormalizeCodexWSDownstreamKeepaliveSlots(s.CodexWSDownstreamKeepaliveSlots)) return err } @@ -3449,9 +3461,9 @@ func NormalizeCodexOverloadWindowMinutes(minutes int) int { return minutes } -// NormalizeCodexWSStatelessSlots 把无状态 WS 连接槽位数限制在 1-32，非正值回落默认 8（issue #522）。 +// NormalizeCodexWSStatelessSlots 把账号共享的 WS 槽位预算限制在 0-32，负值回落默认 8。 func NormalizeCodexWSStatelessSlots(slots int) int { - if slots <= 0 { + if slots < 0 { return 8 } if slots > 32 { diff --git a/database/runtime_cache_scope.go b/database/runtime_cache_scope.go new file mode 100644 index 000000000..7f44a84c5 --- /dev/null +++ b/database/runtime_cache_scope.go @@ -0,0 +1,20 @@ +package database + +import "context" + +func (db *DB) initializeRuntimeCacheScope(ctx context.Context) error { + revision, err := db.GetAPIKeyAuthRevision(ctx) + if err != nil { + return err + } + db.runtimeCacheScope = revision.Namespace + return nil +} + +// RuntimeCacheScope 隔离使用同一个 Redis 的不同数据库，不包含连接凭证原文。 +func (db *DB) RuntimeCacheScope() string { + if db == nil { + return "" + } + return db.runtimeCacheScope +} diff --git a/database/sqlite.go b/database/sqlite.go index 6ad90f0b3..14aabf10b 100644 --- a/database/sqlite.go +++ b/database/sqlite.go @@ -354,6 +354,7 @@ func (db *DB) migrateSQLite(ctx context.Context) error { codex_ws_busy_overflow_enabled INTEGER DEFAULT 0, codex_ws_busy_patience_sec INTEGER DEFAULT 2, codex_ws_stateless_slots INTEGER DEFAULT 8, + codex_ws_downstream_keepalive_slots INTEGER DEFAULT 8, github_token TEXT DEFAULT '', github_proxy_url TEXT DEFAULT '', codex_overload_pause_enabled INTEGER DEFAULT 0, @@ -680,6 +681,7 @@ func (db *DB) migrateSQLite(ctx context.Context) error { {"system_settings", "codex_ws_busy_overflow_enabled", "INTEGER DEFAULT 0"}, {"system_settings", "codex_ws_busy_patience_sec", "INTEGER DEFAULT 2"}, {"system_settings", "codex_ws_stateless_slots", "INTEGER DEFAULT 8"}, + {"system_settings", "codex_ws_downstream_keepalive_slots", "INTEGER DEFAULT 8"}, {"system_settings", "github_token", "TEXT DEFAULT ''"}, {"system_settings", "github_proxy_url", "TEXT DEFAULT ''"}, {"system_settings", "codex_overload_pause_enabled", "INTEGER DEFAULT 0"}, diff --git a/frontend/src/locales/en.json b/frontend/src/locales/en.json index 78a398136..f105d55a0 100644 --- a/frontend/src/locales/en.json +++ b/frontend/src/locales/en.json @@ -4527,8 +4527,10 @@ "githubTokenClear": "Clear", "githubProxy": "GitHub Proxy", "githubProxyDesc": "Outbound proxy applied only to github.com / api.github.com / *.githubusercontent.com, decoupled from the global proxy. Empty falls back to the global proxy (online update falls back to env proxy).", - "codexWSStatelessSlots": "Stateless Connection Slots", - "codexWSStatelessSlotsDesc": "Persistent WS connection slots per (account, cache key) for stateless requests (1-32, default 8). Higher: more idle connections per account; lower: more handshakes, which may trip upstream handshake rate limits at high RPM. Effective value is still capped by the account's dynamic concurrency limit. Hot-reloads.", + "codexWSStatelessSlots": "Blank Prewarm Connections to Keep", + "codexWSStatelessSlotsDesc": "Maximum upstream connections kept ready per account before they contain chat content. Set to 0 to close them after each request. Once chat starts, retention follows Chat Continuation Connections to Keep. Default 8, range 0–32. Changes apply after saving.", + "codexWSDownstreamKeepaliveSlots": "Chat Continuation Connections to Keep", + "codexWSDownstreamKeepaliveSlotsDesc": "Maximum chat connections kept per account after a reply finishes, while waiting for the next input. Shared by all API keys. Close the least recently used connection when over the limit, without interrupting replies in progress. Set to 0 to close after each reply. Connections also close after 30 minutes of inactivity; later input can reconnect using context saved for 45 minutes by default. Default 8, range 0–32. Changes apply after saving.", "codexWSSilentMaxRetries": "Max WS Silent Retries", "codexWSSilentMaxRetriesDesc": "Maximum account-switch retries for upstream WS errors before the first token, range 0-10, default 2.", "schedulerMode": "Scheduler Mode", diff --git a/frontend/src/locales/zh.json b/frontend/src/locales/zh.json index 483c7233c..d0550bda6 100644 --- a/frontend/src/locales/zh.json +++ b/frontend/src/locales/zh.json @@ -4524,8 +4524,10 @@ "githubTokenClear": "清除", "githubProxy": "GitHub 专用代理", "githubProxyDesc": "仅对 github.com / api.github.com / *.githubusercontent.com 生效的出站代理，与全局代理解耦；留空则回落全局代理（在线更新回落环境变量代理）。", - "codexWSStatelessSlots": "无状态连接槽位数", - "codexWSStatelessSlotsDesc": "无状态请求每 (账号, 缓存键) 维度复用的持久 WS 连接槽位数（1-32，默认 8）。调大：单账号挂更多空闲连接；调小：握手更频繁，高 RPM 下可能触发上游握手限流。实际生效值仍受账号动态并发上限钳制，热更新生效。", + "codexWSStatelessSlots": "空白预热连接保留数", + "codexWSStatelessSlotsDesc": "每个账号最多保留多少条提前建立、还没有聊天内容的上游连接。设为 0 时不保留，请求结束后关闭。开始聊天后，连接按“聊天续链连接保留数”保留。默认 8，范围 0～32，保存后生效。", + "codexWSDownstreamKeepaliveSlots": "聊天续链连接保留数", + "codexWSDownstreamKeepaliveSlotsDesc": "每个账号在回复结束后，最多保留多少条等待用户继续输入的聊天连接，所有 API Key 共用。超过数量时关闭最久未使用的一条，不影响正在生成的回复。设为 0 时回复结束就关闭。闲置 30 分钟也会关闭，下次输入可通过默认保存 45 分钟的上下文重新连接。默认 8，范围 0～32，保存后生效。", "codexWSSilentMaxRetries": "WS 静默最大重试次数", "codexWSSilentMaxRetriesDesc": "首包前上游错误最多静默换号重试的次数，范围 0–10，默认 2。", "schedulerMode": "调度模式", diff --git a/frontend/src/pages/Settings.tsx b/frontend/src/pages/Settings.tsx index 7a387ab7a..0b9bfa6fd 100644 --- a/frontend/src/pages/Settings.tsx +++ b/frontend/src/pages/Settings.tsx @@ -2622,6 +2622,7 @@ export default function Settings() { codex_ws_busy_overflow_enabled: false, codex_ws_busy_patience_sec: 2, codex_ws_stateless_slots: 8, + codex_ws_downstream_keepalive_slots: 8, github_token_configured: false, github_proxy_url: '', codex_overload_pause_enabled: false, @@ -4242,12 +4243,27 @@ export default function Settings() { }} /> + + setSettingsForm(f => ({ ...f, codex_ws_downstream_keepalive_slots: value }))} + onValueCommit={(value) => { + void autoSaveSettingsPatch({ codex_ws_downstream_keepalive_slots: value }) + }} + /> + maxCodexWSBusyWaitSec { settings.CodexWSBusyPatienceSec = maxCodexWSBusyWaitSec } - if settings.CodexWSStatelessSlots <= 0 { + if settings.CodexWSStatelessSlots < 0 { settings.CodexWSStatelessSlots = defaultCodexWSStatelessSlots } if settings.CodexWSStatelessSlots > maxCodexWSStatelessSlots { settings.CodexWSStatelessSlots = maxCodexWSStatelessSlots } + settings.CodexWSDownstreamKeepaliveSlots = database.NormalizeCodexWSDownstreamKeepaliveSlots(settings.CodexWSDownstreamKeepaliveSlots) settings.CodexOverloadThresholdPercent = database.NormalizeCodexOverloadThresholdPercent(settings.CodexOverloadThresholdPercent) settings.CodexOverloadPauseMinutes = database.NormalizeCodexOverloadPauseMinutes(settings.CodexOverloadPauseMinutes) settings.CodexOverloadWindowMinutes = database.NormalizeCodexOverloadWindowMinutes(settings.CodexOverloadWindowMinutes) @@ -376,6 +378,7 @@ func ApplyRuntimeSettingsFromSystem(settings *database.SystemSettings) RuntimeSe next.CodexWSBusyOverflow = settings.CodexWSBusyOverflowEnabled next.CodexWSBusyPatienceSec = settings.CodexWSBusyPatienceSec next.CodexWSStatelessSlots = settings.CodexWSStatelessSlots + next.CodexWSDownstreamKeepaliveSlots = settings.CodexWSDownstreamKeepaliveSlots next.GithubToken = strings.TrimSpace(settings.GithubToken) next.GithubProxyURL = strings.TrimSpace(settings.GithubProxyURL) next.CodexOverloadPauseEnabled = settings.CodexOverloadPauseEnabled From c3635bb7b2a4801f0217f6f10509eca843033d80 Mon Sep 17 00:00:00 2001 From: bxb1337 Date: Fri, 9 Oct 2026 19:46:26 +0800 Subject: [PATCH 2/4] =?UTF-8?q?fix(proxy):=20=E6=81=A2=E5=A4=8D=E5=8E=9F?= =?UTF-8?q?=E7=94=9F=20WS=20=E5=B7=A5=E5=85=B7=E7=BB=AD=E9=93=BE=E5=B9=B6?= =?UTF-8?q?=E5=85=B1=E4=BA=AB=E4=B8=8A=E4=B8=8B=E6=96=87=E5=BF=AB=E7=85=A7?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- cache/auth_invalidation.go | 5 +- proxy/codex_ws_transport_lane.go | 5 + proxy/response_cache.go | 39 ++-- proxy/response_cache_budget_test.go | 4 +- proxy/responses_continuation.go | 41 ++++ proxy/responses_ws.go | 93 ++++++-- proxy/responses_ws_context.go | 152 +++++++++---- proxy/responses_ws_local_context.go | 235 ++++++++++++++++++++ proxy/responses_ws_native_context_test.go | 256 ++++++++++++++++++++++ proxy/responses_ws_recovery_error.go | 31 +++ proxy/responses_ws_shared_context.go | 253 +++++++++++++++++++++ proxy/review_regression_test.go | 13 +- proxy/scheduler_queue_error_test.go | 12 +- proxy/upstream_ws_capacity.go | 22 ++ 14 files changed, 1066 insertions(+), 95 deletions(-) create mode 100644 proxy/responses_continuation.go create mode 100644 proxy/responses_ws_local_context.go create mode 100644 proxy/responses_ws_native_context_test.go create mode 100644 proxy/responses_ws_recovery_error.go create mode 100644 proxy/responses_ws_shared_context.go create mode 100644 proxy/upstream_ws_capacity.go diff --git a/cache/auth_invalidation.go b/cache/auth_invalidation.go index 2b38a31ac..a04129695 100644 --- a/cache/auth_invalidation.go +++ b/cache/auth_invalidation.go @@ -19,6 +19,9 @@ type BoundedRuntimeReader interface { GetRuntimeBounded(context.Context, string, string, int64) ([]byte, bool, error) } +// ErrRuntimeSnapshotTooLarge 表示后端值在反序列化前超过读取上限。 +var ErrRuntimeSnapshotTooLarge = errors.New("runtime snapshot exceeds byte limit") + type AuthSnapshotWriter interface { SetAuthSnapshot(context.Context, string, string, json.RawMessage, time.Duration) error } @@ -73,7 +76,7 @@ func (tc *redisTokenCache) GetRuntimeBounded(ctx context.Context, namespace, key return nil, false, err } if int64(len(raw)) > maxBytes { - return nil, false, errors.New("runtime snapshot exceeds byte limit") + return nil, false, ErrRuntimeSnapshotTooLarge } return raw, len(raw) > 0, nil } diff --git a/proxy/codex_ws_transport_lane.go b/proxy/codex_ws_transport_lane.go index f1b4606ed..19ab86775 100644 --- a/proxy/codex_ws_transport_lane.go +++ b/proxy/codex_ws_transport_lane.go @@ -25,6 +25,11 @@ const ( codexRequestKindCompaction = "compaction" ) +// IsCodexWebsocketPrewarmRequest 识别客户端明确声明的预热请求。 +func IsCodexWebsocketPrewarmRequest(headers http.Header, body []byte) bool { + return extractCodexRequestKind(headers, body) == codexRequestKindPrewarm +} + // ResolveCodexWebsocketTransportSessionKey returns the local-only connection // pool lane for one explicit Codex session. A Codex session tree shares // session-id while child agents have independent thread-id values. The shared diff --git a/proxy/response_cache.go b/proxy/response_cache.go index 9317a78fc..c0522b98f 100644 --- a/proxy/response_cache.go +++ b/proxy/response_cache.go @@ -9,6 +9,7 @@ import ( "errors" "fmt" "log" + "strings" "sync" "time" @@ -29,7 +30,7 @@ import ( // 造成跨用户上下文泄露。owner 不匹配一律按缓存未命中处理。 const ( - responseCacheTTL = 10 * time.Minute + responseCacheTTL = 45 * time.Minute responseCacheMaxBytes = 64 << 20 responseCacheMaxEntry = 8 << 20 responseCacheMaxItems = 2000 // 缓存条目上限，防止内存膨胀 @@ -63,13 +64,14 @@ func responseCacheStoreKey(owner, responseID string) string { } type responseCacheEntry struct { - key string - serial uint64 - items []json.RawMessage - blobs []*sharedResponseContextItem - bytes int64 - expiresAt time.Time - element *list.Element + key string + serial uint64 + items []json.RawMessage + blobs []*sharedResponseContextItem + bytes int64 + expiresAt time.Time + element *list.Element + nativeProvenance *nativeWSReplayProvenance } type responseCacheConfig struct { @@ -178,11 +180,12 @@ const ( ) type responseCacheLookupResult struct { - Items []json.RawMessage - Kind responseCacheLookupKind - Source responseCacheLookupSource - Promoted bool - Err error + Items []json.RawMessage + Kind responseCacheLookupKind + Source responseCacheLookupSource + Promoted bool + Err error + nativeProvenance *nativeWSReplayProvenance // 记账辅助位：由 lookupResponseCacheResult 填写，getResponseCacheResult // 在单一临界区内据此更新全部计数器，保证快照不变量任意瞬间成立。 remoteMiss bool @@ -375,6 +378,16 @@ func admitResponseCacheWithTicket(storeKey string, items []json.RawMessage) ([]j respCache.entrySerial++ serial := respCache.entrySerial + native := strings.HasPrefix(storeKey, nativeWSCachePrefix) + if native && (len(items) > respCache.config.maxItems || responseContextLogicalBytes(items) > respCache.config.reconstructMaxBytes) { + if existing := respCache.store[storeKey]; existing != nil { + respCache.removeEntryLocked(existing, responseCacheRemovalReplace) + } + respCache.setMarkerLocked(storeKey, responseCacheLookupKnownOversize, time.Now().Add(respCache.config.ttl)) + respCache.stats.OversizeSkips++ + respCache.stats.OversizeRejections++ + return nil, false, true, serial + } items = trimResponseContextTail(items, respCache.config.maxItems) items, hashes, normalized := respCache.normalizeResponseContextItemsLocked(items) var entryBytes int64 diff --git a/proxy/response_cache_budget_test.go b/proxy/response_cache_budget_test.go index 242973270..c10d0f83c 100644 --- a/proxy/response_cache_budget_test.go +++ b/proxy/response_cache_budget_test.go @@ -53,8 +53,8 @@ func TestResponseCacheDefaultBudget(t *testing.T) { if config.maxEntries != 2000 { t.Fatalf("maxEntries = %d, want 2000", config.maxEntries) } - if config.ttl != 10*time.Minute { - t.Fatalf("ttl = %s, want 10m", config.ttl) + if config.ttl != 45*time.Minute { + t.Fatalf("ttl = %s, want 45m", config.ttl) } if config.maxItems != 200 { t.Fatalf("maxItems = %d, want 200", config.maxItems) diff --git a/proxy/responses_continuation.go b/proxy/responses_continuation.go new file mode 100644 index 000000000..1c4d9ba63 --- /dev/null +++ b/proxy/responses_continuation.go @@ -0,0 +1,41 @@ +package proxy + +import ( + "fmt" + "net/http" + + "github.com/codex2api/api" +) + +// ResponsesContinuationLostError 表示连接内续链状态已失效；不能携带旧 ID 重拨。 +type ResponsesContinuationLostError struct { + Reason string +} + +type ResponsesContinuationBusyError struct{} + +func (*ResponsesContinuationBusyError) Error() string { + return "original websocket connection remained busy until the wait deadline" +} + +func (*ResponsesContinuationBusyError) Unwrap() error { + return &Error{Code: ErrorCodeAccountPoolConcurrencySaturated, Type: ErrorTypeServerError, + HTTPStatus: http.StatusServiceUnavailable, Message: "Original WebSocket connection is busy, please retry later"} +} + +func (err *ResponsesContinuationLostError) Error() string { + return fmt.Sprintf("websocket continuation unavailable: %s", err.Reason) +} + +func (err *ResponsesContinuationLostError) Unwrap() error { + return &Error{Code: "previous_response_not_found", Type: ErrorTypeInvalidRequest, + HTTPStatus: http.StatusConflict, Message: "Previous response context is unavailable"} +} + +func nativeResponsesWSContextError(err *api.APIError) *api.APIError { + if err == nil || err.Code == api.ErrCodeServiceUnavailable { + return err + } + return api.NewAPIErrorWithDetails(api.ErrorCode("previous_response_not_found"), + err.Message, api.ErrorTypeInvalidRequest, err.Details) +} diff --git a/proxy/responses_ws.go b/proxy/responses_ws.go index ecf0d5438..4bdd2b611 100644 --- a/proxy/responses_ws.go +++ b/proxy/responses_ws.go @@ -82,7 +82,8 @@ func isPreviousResponseNotFoundBody(payload []byte) bool { for _, path := range []string{"error.message", "message"} { message := strings.ToLower(strings.TrimSpace(gjson.GetBytes(body, path).String())) if message == "previous_response_id is not available for this user" || - strings.Contains(message, "previous response with id") { + strings.Contains(message, "previous response with id") || + strings.TrimSuffix(strings.NewReplacer("`", "", "\"", "", "'", "").Replace(message), ".") == "invalid previous_response_id" { return true } } @@ -443,6 +444,9 @@ func (h *Handler) forwardResponsesWebSocketTurn(c *gin.Context, conn *websocket. turnContinuation := codexWSTurnContinuationToken(rawBody) != "" _, turnHasBinding := h.store.SessionAffinityAccountID(affinityKey) respCacheOwner := responseCacheOwner(apiKeyID) + if gjson.GetBytes(rawBody, "store").Type == gjson.False { + respCacheOwner = nativeWSTurnCacheOwner(c, respCacheOwner, nativeWSTurnScope{body: rawBody, identity: sessionIdentity, databaseScope: h.db.RuntimeCacheScope()}) + } markResponsesWSContinuationCapable(respCacheOwner, rawBody) ruleIdentity := h.payloadRuleIdentity(c) // 上下文压缩轮豁免首字超时看门狗（issue #381）：压缩首帧天然慢，超时换号无益。 @@ -454,8 +458,7 @@ func (h *Handler) forwardResponsesWebSocketTurn(c *gin.Context, conn *websocket. } codexBody, naturalImageIntent := prepareResponsesWebSocketTurnBody(rawBody) - // Pin an available L1 ancestor before upstream generation; defer backend - // lookup, merging and serialization until a snapshot is actually needed. + // 生成前固定祖先；原生 WS 必要时读共享快照，历史合并与序列化保持按需执行。 // strip 策略：剥离图片工具能力声明后作为普通文本请求继续（issue #411）。 codexBody = applyImageGenerationStripPolicy(c, codexBody) if err := validateResponsesImageGenerationSizes(codexBody); err != nil { @@ -503,7 +506,7 @@ func (h *Handler) forwardResponsesWebSocketTurn(c *gin.Context, conn *websocket. releaseAPIKeyConcurrency, concurrencyErr, ok := h.acquireAPIKeyConcurrencyForWebSocket(c) if !ok { _ = writeResponsesWSError(conn, concurrencyErr) - return newResponsesWSCloseError(websocket.CloseTryAgainLater, concurrencyErr.Message, concurrencyErr) + return nil } if releaseAPIKeyConcurrency != nil { defer releaseAPIKeyConcurrency() @@ -577,7 +580,11 @@ func (h *Handler) forwardResponsesWebSocketTurn(c *gin.Context, conn *websocket. // 忽略本地 WHAM 100% 快照；previous_response_id 本身不足以证明这是活跃 turn。 continuationPinned := turnContinuation && turnHasBinding continuationDegraded := false + if err := waitNativeWSCommit(c.Request.Context(), respCacheOwner, gjson.GetBytes(codexBody, "previous_response_id").String()); err != nil { + return writeResponsesWSError(conn, nativeResponsesWSContextError(responsesWSContextUnavailable(http.StatusServiceUnavailable, "context_commit_pending"))) + } turnReplay := newResponsesWSReplaySource(codexBody, respCacheOwner) + accountFilter = nativeWSReplayAccountFilter(accountFilter, turnReplay) degradeContinuation := func(reason string, attempt int) *api.APIError { expanded, lost, contextErr := degradeResponsesWSContinuationWithSource(codexBody, respCacheOwner, turnReplay) if contextErr != nil { @@ -586,11 +593,11 @@ func (h *Handler) forwardResponsesWebSocketTurn(c *gin.Context, conn *websocket. continuationDegraded = true continuationPinned = false codexBody = expanded - turnReplay = newResponsesWSReplaySourceFromInput(responsesInputRaw(expanded)) + turnReplay = newResponsesWSReplaySourceFromRecovery(responsesInputRaw(expanded), turnReplay) if lost { turnReplay = newResponsesWSReplaySourceFromInput("") } - log.Printf("Responses WebSocket continuation degraded: %s, stripped previous_response_id and retried once (attempt %d)", reason, attempt) + log.Printf("Responses WebSocket continuation degraded: reason=%s, stripped previous_response_id and retried once (attempt %d)", reason, attempt) return nil } preserveContinuationBinding := func() bool { @@ -628,8 +635,7 @@ func (h *Handler) forwardResponsesWebSocketTurn(c *gin.Context, conn *websocket. if boundID, bound := h.store.SessionAffinityAccountID(affinityKey); bound { if exclude := retryExclusions.ForSelection(); exclude[boundID] { if contextErr := degradeContinuation(fmt.Sprintf("bound account %d excluded by this request", boundID), attempt+1); contextErr != nil { - _ = writeResponsesWSError(conn, contextErr) - return newResponsesWSCloseError(responsesWSContextCloseCode(contextErr), contextErr.Message, contextErr) + return writeResponsesWSRecoveryError(conn, contextErr) } } } @@ -672,13 +678,18 @@ func (h *Handler) forwardResponsesWebSocketTurn(c *gin.Context, conn *websocket. return errResponsesWSClientGone } _ = writeResponsesWSError(conn, apiErr) - return newResponsesWSCloseError(websocket.CloseTryAgainLater, apiErr.Message, apiErr) + return nil } if attempt > 0 { clearNewAPIUpstreamCyberPolicyDecision(c) } h.AcquireAPIKeyScopeConcurrency(c, account) + turnReplay.setAccount(account) + if continuationDegraded && turnReplay.Input() == "" && turnReplay.err != nil { + h.store.Release(account) + return writeResponsesWSRecoveryError(conn, turnReplay.err) + } start := time.Now() proxyURL := h.resolveProxyForAttempt(account, stickyProxyURL) if !retainedHTTPFallback && !continuousRetryBuffersAttempts(continuousRetryPolicy) { @@ -738,18 +749,15 @@ func (h *Handler) forwardResponsesWebSocketTurn(c *gin.Context, conn *websocket. ttftGuard.Stop() upstreamCancel() h.store.Release(account) - _ = writeResponsesWSError(conn, contextErr) - return newResponsesWSCloseError(responsesWSContextCloseCode(contextErr), contextErr.Message, contextErr) + return writeResponsesWSRecoveryError(conn, contextErr) } // 降级时快照已经算过一次，直接复用展开后的 input，不再二次过滤。 - attemptReplay = newResponsesWSReplaySourceFromInput(responsesInputRaw(upstreamBody)) + attemptReplay = newResponsesWSReplaySourceFromRecovery(responsesInputRaw(upstreamBody), turnReplay) + attemptReplay.setAccount(account) if lost { attemptReplay = nil } } - if gjson.GetBytes(rawBody, "store").Type == gjson.False { - attemptReplay = nil - } // service_tier 记账按 payload 规则改写后的值归因（覆写 service_tier 的规则才生效）。 serviceTier = EffectiveRequestedServiceTier(upstreamBody, effectiveModel, downstreamHeaders, attemptIdentity) // 在 useWebsocket 最终确定后再派生上游身份键：与 handler.go 的 @@ -776,6 +784,33 @@ func (h *Handler) forwardResponsesWebSocketTurn(c *gin.Context, conn *websocket. } if reqErr != nil { + var capacityErr *UpstreamWSConnectionCapacityError + if errors.As(reqErr, &capacityErr) { + ttftGuard.Stop() + upstreamCancel() + h.store.Release(account) + return writeResponsesWSError(conn, api.NewAPIError(api.ErrorCode(ErrorCodeUpstreamWSConnectionCapacity), capacityErr.Error(), api.ErrorTypeServer)) + } + var continuationBusy *ResponsesContinuationBusyError + if errors.As(reqErr, &continuationBusy) { + ttftGuard.Stop() + upstreamCancel() + h.store.Release(account) + return writeResponsesWSError(conn, api.NewAPIError(api.ErrCodeServiceUnavailable, continuationBusy.Error(), api.ErrorTypeServer)) + } + var continuationLost *ResponsesContinuationLostError + if errors.As(reqErr, &continuationLost) { + ttftGuard.Stop() + upstreamCancel() + h.store.Release(account) + if canDegradeContinuation() { + if contextErr := degradeContinuation(continuationLost.Reason, attempt+1); contextErr != nil { + return writeResponsesWSRecoveryError(conn, contextErr) + } + continue + } + return writeResponsesWSRecoveryError(conn, nativeResponsesWSContextError(responsesWSContextUnavailable(http.StatusConflict, continuationLost.Reason))) + } if quotaErr := apiKeyModelRequestError(reqErr); quotaErr != nil { ttftGuard.Stop() h.store.Release(account) @@ -903,8 +938,7 @@ func (h *Handler) forwardResponsesWebSocketTurn(c *gin.Context, conn *websocket. if canDegradeContinuation() && isPreviousResponseNotFoundBody(errBody) { if contextErr := degradeContinuation(fmt.Sprintf("upstream rejected previous_response_id on account %d", account.ID()), attempt+1); contextErr != nil { h.store.Release(account) - _ = writeResponsesWSError(conn, contextErr) - return newResponsesWSCloseError(responsesWSContextCloseCode(contextErr), contextErr.Message, contextErr) + return writeResponsesWSRecoveryError(conn, contextErr) } SyncCodexUsageState(h.store, account, resp) h.store.Release(account) @@ -996,8 +1030,7 @@ func (h *Handler) forwardResponsesWebSocketTurn(c *gin.Context, conn *websocket. // 账号已在流内释放，未记失败也未解绑：剥离续链 id 后原地再试一次。 // turn-state 钉号同样走这条路：上游已经说找不到 id，换号无益，剥 id 才能继续。 if contextErr := degradeContinuation(fmt.Sprintf("upstream rejected previous_response_id on account %d", account.ID()), attempt+1); contextErr != nil { - _ = writeResponsesWSError(conn, contextErr) - return newResponsesWSCloseError(responsesWSContextCloseCode(contextErr), contextErr.Message, contextErr) + return writeResponsesWSRecoveryError(conn, contextErr) } continue } @@ -1156,7 +1189,12 @@ func (h *Handler) streamResponsesWSUpstream( if continuousRetryBuffersAttempts(continuousRetryPolicy) { wsReplay = h.newContinuousRetryWSReplay() } + finishContextCommit := func() {} + defer func() { finishContextCommit() }() writeClientMessage := func(payload []byte) error { + if isResponsesSuccessTerminalEvent(gjson.GetBytes(payload, "type").String()) { + finishContextCommit = beginNativeWSCommit(respCacheOwner, gjson.GetBytes(payload, "response.id").String()) + } return writeResponsesWSMessage(conn, payload) } // 首 token 前收到不可重试的 response.failed 时置位:不把原始失败帧透传给客户端, @@ -1492,10 +1530,12 @@ func (h *Handler) streamResponsesWSUpstream( // Committed completed/incomplete responses can be continuation history. // Failed attempts and locally blocked/unwritable replays leave no cache. if !outputCollector.overflow { - cacheResponsesWSCompletedResponse(respCacheOwner, replayInput.Input(), completedResponsePayload, outputCollector.Items()) + cacheCommittedResponsesWSContext(nativeWSCompletedContext{owner: respCacheOwner, input: replayInput.Input(), + completed: completedResponsePayload, outputs: outputCollector.Items(), provenance: replayInput.committedProvenance()}) } } } + finishContextCommit() _ = wsReplay.Close() if continuousRetryBufferedAttemptCommitted(continuousRetryPolicy, outcome) { h.store.BindSessionAffinityWithGuard(affinityKey, account, proxyURL, affinityGuard) @@ -1612,6 +1652,9 @@ func (h *Handler) streamResponsesWSUpstream( return errResponsesWSClientGone } _ = writeResponsesWSError(conn, clientErr) + if preserveErrorCode { + return nil + } return newResponsesWSCloseError(responsesWSCloseCodeForStatus(outcome.logStatusCode), clientErr.Message, apiErr) } if outcome.logStatusCode != http.StatusOK && !hideUpstreamErrors && len(terminalFailureClientPayload) > 0 && !downstreamWrote { @@ -1764,11 +1807,13 @@ func writeResponsesWSErrorWithin(conn *websocket.Conn, apiErr *api.APIError, tim apiErr = api.NewAPIError(api.ErrCodeServerError, "Internal server error", api.ErrorTypeServer) } payload, err := json.Marshal(struct { - Type string `json:"type"` - Error *api.APIError `json:"error"` + Type string `json:"type"` + Status int `json:"status"` + Error *api.APIError `json:"error"` }{ - Type: "error", - Error: apiErr, + Type: "error", + Status: responsesWSErrorStatus(apiErr), + Error: apiErr, }) if err != nil { return err diff --git a/proxy/responses_ws_context.go b/proxy/responses_ws_context.go index b7f9b8305..fbdfda798 100644 --- a/proxy/responses_ws_context.go +++ b/proxy/responses_ws_context.go @@ -10,7 +10,6 @@ import ( "time" "github.com/codex2api/api" - "github.com/gorilla/websocket" "github.com/tidwall/gjson" "github.com/tidwall/sjson" ) @@ -19,15 +18,28 @@ import ( // incremental request sent over a healthy upstream WebSocket. Missing ancestry // must not turn a partial input into a supposedly complete cached conversation. func responsesWSReplayInput(body []byte, owner string) (string, *api.APIError) { + if strings.HasPrefix(owner, nativeWSCachePrefix) { + return responsesWSReplayInputWithLookup(body, owner, nativeWSContextLookup) + } return responsesWSReplayInputWithLookup(body, owner, getResponseCacheForReplay) } func responsesWSReplayInputWithLookup(body []byte, owner string, lookupFn func(string, string) responseCacheLookupResult) (string, *api.APIError) { - current := gjson.GetBytes(body, "input") + return responsesWSReplayInputChecked(body, responsesWSReplayOptions{owner: owner, lookup: lookupFn}) +} + +type responsesWSReplayOptions struct { + owner string + lookup func(string, string) responseCacheLookupResult + allowOpaque bool +} + +func responsesWSReplayInputChecked(body []byte, options responsesWSReplayOptions) (string, *api.APIError) { + owner := options.owner var items []json.RawMessage previousID := strings.TrimSpace(gjson.GetBytes(body, "previous_response_id").String()) if previousID != "" { - lookup := lookupFn(owner, previousID) + lookup := options.lookup(owner, previousID) if lookup.Kind != responseCacheLookupHit { prepared := responsesBodyPreparation{PreviousResponseID: previousID, CacheLookup: lookup, RequiresLocalContext: true} status, reason, _ := responseCachePreparationFailure(prepared) @@ -35,18 +47,13 @@ func responsesWSReplayInputWithLookup(body []byte, owner string, lookupFn func(s } items = append(items, lookup.Items...) } - var input []json.RawMessage - current.ForEach(func(_, item gjson.Result) bool { - if raw, ok := replayableCachedInputItem(item); ok { - input = append(input, raw) - } - return true - }) - items = mergeResponsesWSContext(items, input) - for _, item := range items { - if gjson.GetBytes(item, "encrypted_content").String() != "" { - return "", responsesWSContextUnavailable(http.StatusConflict, "nonportable_encrypted_context") - } + input, itemErr := responsesWSReplayItems(gjson.GetBytes(body, "input"), options) + if itemErr != nil { + return "", itemErr + } + items, itemErr = validateResponsesWSReplayItems(mergeResponsesWSContext(items, input), options) + if itemErr != nil { + return "", itemErr } respCache.mu.RLock() config := respCache.config @@ -61,6 +68,40 @@ func responsesWSReplayInputWithLookup(body []byte, owner string, lookupFn func(s return string(raw), nil } +func responsesWSReplayItems(current gjson.Result, options responsesWSReplayOptions) ([]json.RawMessage, *api.APIError) { + var input []json.RawMessage + var itemErr *api.APIError + current.ForEach(func(_, item gjson.Result) bool { + if strings.HasPrefix(options.owner, nativeWSCachePrefix) { + raw, err := nativeWSReplayItemChecked(item, options.allowOpaque) + if err != nil { + itemErr = err + return false + } + input = append(input, raw) + } else if raw, ok := replayableCachedInputItem(item); ok { + input = append(input, raw) + } + return true + }) + return input, itemErr +} + +func validateResponsesWSReplayItems(items []json.RawMessage, options responsesWSReplayOptions) ([]json.RawMessage, *api.APIError) { + for index, item := range items { + if strings.HasPrefix(options.owner, nativeWSCachePrefix) { + raw, err := nativeWSReplayItemChecked(gjson.ParseBytes(item), options.allowOpaque) + if err != nil { + return nil, err + } + items[index] = raw + } else if gjson.GetBytes(item, "encrypted_content").String() != "" { + return nil, responsesWSContextUnavailable(http.StatusConflict, "nonportable_encrypted_context") + } + } + return items, nil +} + func responsesWSContextUnavailable(status int, reason string) *api.APIError { code, kind, message := api.ErrCodeResponseContextUnavailable, api.ErrorTypeInvalidRequest, "Previous response context is unavailable" if status == http.StatusServiceUnavailable { @@ -109,6 +150,10 @@ func mergeResponsesWSContext(previous, current []json.RawMessage) []json.RawMess // replay sanitizer: this snapshot may be replayed on a different account. // Call only after the attempt and downstream output have committed successfully. func cacheResponsesWSCompletedResponse(owner, input string, completed []byte, outputItems []json.RawMessage) { + if strings.HasPrefix(owner, nativeWSCachePrefix) { + cacheNativeWSContext(nativeWSCompletedContext{owner: owner, input: input, completed: completed, outputs: outputItems}) + return + } responseID := gjson.GetBytes(completed, "response.id").String() if responseID == "" || input == "" { return @@ -165,14 +210,19 @@ func degradeResponsesWSContinuationWithSource(body []byte, owner string, source } var input string var apiErr *api.APIError - if source != nil && source.previous != nil { + if source != nil { input = source.Input() - recordResponseCacheLookup(owner, *source.previous) + if source.previous != nil { + recordResponseCacheLookup(owner, *source.previous) + } apiErr = source.err } else { input, apiErr = responsesWSReplayInput(body, owner) } if apiErr != nil { + if strings.HasPrefix(owner, nativeWSCachePrefix) { + return nil, false, nativeResponsesWSContextError(apiErr) + } if !responsesWSContinuationFailOpen() { return nil, false, apiErr } @@ -225,10 +275,8 @@ func responsesWSContextReason(apiErr *api.APIError) string { return apiErr.Message } -// markResponsesWSContinuationCapable 给可能续链的 WS 会话在根轮就授予 on_demand -// 写入资格。Codex CLI 全量上下文每轮 store:false，永远不会带 previous_response_id -// 回来，不能为它开闸；store 非 false 的会话才有机会续链，若根轮不入缓存， -// 之后每轮都因祖先缺失无法成快照，降级必然失败。 +// markResponsesWSContinuationCapable 为允许持久存储的会话授予 on_demand 写入资格。 +// store:false 的原生 WS 续链单独使用完整快照，并复用共享后端。 func markResponsesWSContinuationCapable(owner string, rawBody []byte) { if store := gjson.GetBytes(rawBody, "store"); store.Exists() && store.Type == gjson.False { return @@ -236,17 +284,18 @@ func markResponsesWSContinuationCapable(owner string, rawBody []byte) { markResponseCacheChainOwnerIfOnDemand(owner) } -// responsesWSReplaySource pins a live L1 ancestor at turn admission. Its immutable -// bodies survive expiry/eviction during generation; merging and backend fallback -// remain lazy and do not count as client replay hits or misses. +// 入场时固定不可变祖先；原生 WS 缺少 L1 时读取共享快照以保留加密上下文的账号。 +// 旧续链后端查询、历史合并与序列化仍按需执行，不提前记为客户端回放命中。 type responsesWSReplaySource struct { - body []byte - owner string - precomputed bool - once sync.Once - input string - previous *responseCacheLookupResult - err *api.APIError + body []byte + owner string + precomputed bool + once sync.Once + input string + previous *responseCacheLookupResult + err *api.APIError + attemptProvenance *nativeWSReplayProvenance + recoveryProvenance *nativeWSReplayProvenance } func newResponsesWSReplaySource(body []byte, owner string) *responsesWSReplaySource { @@ -254,9 +303,13 @@ func newResponsesWSReplaySource(body []byte, owner string) *responsesWSReplaySou if previousID := gjson.GetBytes(body, "previous_response_id").String(); previousID != "" { respCache.mu.RLock() if entry := respCache.store[responseCacheStoreKey(owner, previousID)]; entry != nil && time.Now().Before(entry.expiresAt) { - source.previous = &responseCacheLookupResult{Kind: responseCacheLookupHit, Source: responseCacheSourceLocal, Items: append([]json.RawMessage(nil), entry.items...)} + source.previous = &responseCacheLookupResult{Kind: responseCacheLookupHit, Source: responseCacheSourceLocal, Items: append([]json.RawMessage(nil), entry.items...), nativeProvenance: entry.nativeProvenance} } respCache.mu.RUnlock() + if source.previous == nil && strings.HasPrefix(owner, nativeWSCachePrefix) { + lookup := nativeWSContextLookup(owner, previousID) + source.previous = &lookup + } } return source } @@ -271,25 +324,30 @@ func (s *responsesWSReplaySource) Input() string { return "" } if s.precomputed { - return s.input + return s.nativeValidatedInput() } s.once.Do(func() { - s.input, s.err = responsesWSReplayInputWithLookup(s.body, s.owner, func(owner, id string) responseCacheLookupResult { + if strings.HasPrefix(s.owner, nativeWSCachePrefix) && s.previous == nil { + if id := gjson.GetBytes(s.body, "previous_response_id").String(); id != "" { + lookup := nativeWSContextLookup(s.owner, id) + s.previous = &lookup + } + } + options := responsesWSReplayOptions{owner: s.owner, allowOpaque: s.nativeReplayProvenance().matches(s.attemptProvenance)} + options.lookup = func(owner, id string) responseCacheLookupResult { if s.previous != nil { return *s.previous } - return lookupResponseCacheResultWithOwnership(owner, id, true) - }) + lookup := nativeWSLocalLookup(owner, id) + if !strings.HasPrefix(owner, nativeWSCachePrefix) { + lookup = lookupResponseCacheResultWithOwnership(owner, id, true) + } else { + lookup = nativeWSContextLookup(owner, id) + } + s.previous = &lookup + return lookup + } + s.input, s.err = responsesWSReplayInputChecked(s.body, options) }) - return s.input -} - -// Record failures only when returning them to a client, not while considering a -// best-effort snapshot after an otherwise successful upstream turn. -func responsesWSContextCloseCode(apiErr *api.APIError) int { - if apiErr.Code == api.ErrCodeServiceUnavailable { - return websocket.CloseInternalServerErr - } - recordResponseCacheKnownUnavailableError() - return websocket.ClosePolicyViolation + return s.nativeValidatedInput() } diff --git a/proxy/responses_ws_local_context.go b/proxy/responses_ws_local_context.go new file mode 100644 index 000000000..8ef2b7909 --- /dev/null +++ b/proxy/responses_ws_local_context.go @@ -0,0 +1,235 @@ +package proxy + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "net/http" + "strings" + "sync" + "time" + + "github.com/codex2api/api" + "github.com/codex2api/auth" + "github.com/gin-gonic/gin" + "github.com/tidwall/gjson" +) + +const nativeWSCachePrefix = "native-ws:" +const nativeWSCommitMaxEntries = 2000 +const nativeWSCommitWait = 5 * time.Second + +var nativeWSCommits = struct { + sync.Mutex + pending map[string]chan struct{} +}{pending: make(map[string]chan struct{})} + +func nativeWSCacheOwner(owner string, identity requestSessionIdentity) string { + if identity.explicitUpstreamID == "" && !identity.hasDownstreamAffinity { + return nativeWSCachePrefix + owner + } + hash := sha256.Sum256([]byte(identity.affinityID)) + return nativeWSCachePrefix + owner + ":" + hex.EncodeToString(hash[:]) +} + +type nativeWSTurnScope struct { + body []byte + identity requestSessionIdentity + databaseScope string +} + +func nativeWSTurnCacheOwner(c *gin.Context, owner string, scope nativeWSTurnScope) string { + keyIdentity := deterministicPromptCacheKey(strings.TrimPrefix(strings.TrimSpace(downstreamAuthorizationHeader(c.Request)), "Bearer "), nil) + parts := []string{nativeWSCacheOwner(owner, scope.identity), responsesWSSessionPreemptScopeHash(c, scope.identity), + responsesWSTransportLane(c, scope.body, scope.identity), gjson.GetBytes(scope.body, "stream_id").String(), keyIdentity, scope.databaseScope} + hash := sha256.Sum256([]byte(strings.Join(parts, "\x00"))) + return nativeWSCachePrefix + hex.EncodeToString(hash[:]) +} + +func nativeWSLocalLookup(owner, responseID string) responseCacheLookupResult { + key := responseCacheStoreKey(owner, responseID) + respCache.mu.Lock() + defer respCache.mu.Unlock() + if entry := respCache.store[key]; entry != nil && time.Now().Before(entry.expiresAt) { + respCache.lru.MoveToFront(entry.element) + return responseCacheLookupResult{Kind: responseCacheLookupHit, Source: responseCacheSourceLocal, + Items: append([]json.RawMessage(nil), entry.items...), nativeProvenance: entry.nativeProvenance} + } + if marker := respCache.markers[key]; marker != nil && time.Now().Before(marker.expiresAt) { + return responseCacheLookupResult{Kind: marker.kind} + } + return responseCacheLookupResult{Kind: responseCacheLookupMiss} +} + +// 终态暴露前建立屏障，跨下游重连须等提交完成再读取快照。 +func beginNativeWSCommit(owner, responseID string) func() { + if !strings.HasPrefix(owner, nativeWSCachePrefix) || responseID == "" { + return func() {} + } + key := responseCacheStoreKey(owner, responseID) + nativeWSCommits.Lock() + if nativeWSCommits.pending[key] != nil || len(nativeWSCommits.pending) >= nativeWSCommitMaxEntries { + nativeWSCommits.Unlock() + return func() {} + } + done := make(chan struct{}) + nativeWSCommits.pending[key] = done + nativeWSCommits.Unlock() + finishShared := beginNativeWSSharedCommit(key) + var once sync.Once + return func() { + once.Do(func() { + finishShared() + nativeWSCommits.Lock() + delete(nativeWSCommits.pending, key) + close(done) + nativeWSCommits.Unlock() + }) + } +} + +func waitNativeWSCommit(ctx context.Context, owner, responseID string) error { + nativeWSCommits.Lock() + done := nativeWSCommits.pending[responseCacheStoreKey(owner, responseID)] + nativeWSCommits.Unlock() + if done == nil { + return waitNativeWSSharedCommit(ctx, owner, responseID) + } + timer := time.NewTimer(nativeWSCommitWait) + defer timer.Stop() + select { + case <-done: + return waitNativeWSSharedCommit(ctx, owner, responseID) + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: + return &ResponsesContinuationLostError{Reason: "context_commit_pending"} + } +} + +func nativeWSReplayItem(item gjson.Result) (json.RawMessage, *api.APIError) { + return nativeWSReplayItemChecked(item, false) +} + +func nativeWSReplayItemChecked(item gjson.Result, allowOpaque bool) (json.RawMessage, *api.APIError) { + if (item.Get("encrypted_content").String() != "" && !allowOpaque) || item.Get("type").String() == "item_reference" { + return nil, nativeResponsesWSContextError(responsesWSContextUnavailable(http.StatusConflict, "nonportable_context")) + } + raw, ok := normalizeReplayableCachedFunctionCall(json.RawMessage(item.Raw)) + if !ok { + return nil, nativeResponsesWSContextError(responsesWSContextUnavailable(http.StatusConflict, "invalid_context_item")) + } + raw, ok = stripResponseItemID(raw) + if !ok { + return nil, nativeResponsesWSContextError(responsesWSContextUnavailable(http.StatusConflict, "invalid_context_item")) + } + return raw, nil +} + +type nativeWSCompletedContext struct { + owner string + input string + completed []byte + outputs []json.RawMessage + provenance *nativeWSReplayProvenance +} + +type nativeWSReplayProvenance struct { + accountID int64 + generation int64 + domain string +} + +func nativeWSAccountProvenance(account *auth.Account) *nativeWSReplayProvenance { + if account == nil { + return nil + } + return &nativeWSReplayProvenance{accountID: account.ID(), generation: account.GetCredentialGeneration(), domain: accountCompactionDomain(account)} +} + +func (provenance *nativeWSReplayProvenance) matches(current *nativeWSReplayProvenance) bool { + return provenance != nil && current != nil && *provenance == *current +} + +func cacheNativeWSContext(snapshot nativeWSCompletedContext) { + owner, input, completed, outputs := snapshot.owner, snapshot.input, snapshot.completed, snapshot.outputs + responseID := gjson.GetBytes(completed, "response.id").String() + if input == "" || responseID == "" { + return + } + var items []json.RawMessage + if json.Unmarshal([]byte(input), &items) != nil { + return + } + terminal := gjson.GetBytes(completed, "response.output").Array() + if len(terminal) > len(outputs) { + outputs = nil + for _, item := range terminal { + outputs = append(outputs, json.RawMessage(item.Raw)) + } + } + // 保留不透明输出，恢复时核对账号兼容性；不能丢掉必要推理后声称上下文完整。 + items = mergeResponsesWSContext(items, outputs) + key := responseCacheStoreKey(owner, responseID) + items, admitted, _, serial := admitResponseCacheWithTicket(key, items) + respCache.mu.RLock() + expiresAt := time.Now().Add(respCache.config.ttl) + respCache.mu.RUnlock() + setNativeWSLocalMetadata(key, nativeWSLocalMetadata{serial: serial, provenance: snapshot.provenance, expiresAt: expiresAt}) + record := nativeWSSharedContext{Version: nativeWSContextVersion, Items: items, ExpiresAt: expiresAt} + if snapshot.provenance != nil { + record.AccountID, record.Generation, record.Domain = snapshot.provenance.accountID, snapshot.provenance.generation, snapshot.provenance.domain + } + if !admitted && items == nil || len(record.Domain) > nativeWSDomainMaxBytes { + record.Items, record.Domain, record.Unavailable = nil, "", "oversize" + } + persistNativeWSContext(nativeWSSharedWrite{key: key, response: responseID, serial: serial, record: record}) +} + +func cacheCommittedResponsesWSContext(snapshot nativeWSCompletedContext) { + if strings.HasPrefix(snapshot.owner, nativeWSCachePrefix) { + cacheNativeWSContext(snapshot) + return + } + cacheResponsesWSCompletedResponse(snapshot.owner, snapshot.input, snapshot.completed, snapshot.outputs) +} + +func (s *responsesWSReplaySource) setAccount(account *auth.Account) { + if s != nil { + s.attemptProvenance = nativeWSAccountProvenance(account) + } +} + +func (s *responsesWSReplaySource) nativeReplayProvenance() *nativeWSReplayProvenance { + if s.previous != nil { + return s.previous.nativeProvenance + } + if s.recoveryProvenance != nil { + return s.recoveryProvenance + } + return s.attemptProvenance +} + +func (s *responsesWSReplaySource) nativeValidatedInput() string { + if strings.Contains(s.input, `"encrypted_content"`) && !s.nativeReplayProvenance().matches(s.attemptProvenance) { + s.err = nativeResponsesWSContextError(responsesWSContextUnavailable(http.StatusConflict, "incompatible_encrypted_context")) + return "" + } + return s.input +} + +func (s *responsesWSReplaySource) committedProvenance() *nativeWSReplayProvenance { + if s == nil { + return nil + } + return s.attemptProvenance +} + +func newResponsesWSReplaySourceFromRecovery(input string, source *responsesWSReplaySource) *responsesWSReplaySource { + replay := newResponsesWSReplaySourceFromInput(input) + if source != nil && strings.Contains(input, `"encrypted_content"`) { + replay.recoveryProvenance = source.nativeReplayProvenance() + } + return replay +} diff --git a/proxy/responses_ws_native_context_test.go b/proxy/responses_ws_native_context_test.go new file mode 100644 index 000000000..c6c3a80fb --- /dev/null +++ b/proxy/responses_ws_native_context_test.go @@ -0,0 +1,256 @@ +package proxy + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/codex2api/auth" + "github.com/codex2api/config" + "github.com/codex2api/database" + "github.com/gin-gonic/gin" + "github.com/gorilla/websocket" + "github.com/tidwall/gjson" +) + +type nativeWSFixture struct { + url string + account *auth.Account +} + +func newNativeWSFixture(t *testing.T, execute func(context.Context, []byte) (*http.Response, error)) nativeWSFixture { + t.Helper() + gin.SetMode(gin.TestMode) + resetResponseCacheForTest() + oldExecutor := WebsocketExecuteFunc + t.Cleanup(func() { + WebsocketExecuteFunc = oldExecutor + resetResponseCacheForTest() + }) + WebsocketExecuteFunc = func(ctx context.Context, account *auth.Account, body []byte, _, _, _ string, _ *DeviceProfileConfig, _ http.Header, _ string) (*http.Response, error) { + if int64(1) != account.GetActiveRequests() { + t.Error("upstream request did not occupy account capacity") + } + return execute(ctx, body) + } + store := auth.NewStore(nil, nil, &database.SystemSettings{MaxConcurrency: 1}) + t.Cleanup(store.Stop) + account := &auth.Account{DBID: 1, AccessToken: "fixture", PlanType: "plus"} + store.AddAccount(account) + handler := NewHandler(store, nil, &config.Config{AllowAnonymousV1: true}, nil) + router := gin.New() + handler.RegisterRoutes(router) + server := httptest.NewServer(router) + t.Cleanup(server.Close) + return nativeWSFixture{url: "ws" + strings.TrimPrefix(server.URL, "http") + "/v1/responses", account: account} +} + +func (fixture nativeWSFixture) connect(t *testing.T) *websocket.Conn { + t.Helper() + conn, _, err := websocket.DefaultDialer.Dial(fixture.url, http.Header{"Session-Id": {"native-context"}}) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { conn.Close() }) + return conn +} + +func nativeWSTurn(t *testing.T, conn *websocket.Conn, input string) []byte { + t.Helper() + var request map[string]any + if err := json.Unmarshal([]byte(input), &request); err != nil { + t.Fatal(err) + } + request["type"], request["model"], request["store"] = "response.create", "gpt-5.5", false + if err := conn.WriteJSON(request); err != nil { + t.Fatal(err) + } + conn.SetReadDeadline(time.Now().Add(3 * time.Second)) + return readResponsesWSTerminalEvent(t, conn) +} + +func nativeWSResponse(id string, outputs ...string) *http.Response { + return &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(wsContextTestSSE(id, outputs...)))} +} + +func TestNativeWSStoreFalseRecoversToolChainAfterReconnect(t *testing.T) { + for _, failure := range []string{"binding_lost", "invalid_id"} { + t.Run(failure, func(t *testing.T) { + var calls atomic.Int32 + fixture := newNativeWSFixture(t, func(_ context.Context, body []byte) (*http.Response, error) { + switch calls.Add(1) { + case 1: + return nativeWSResponse("one", `{"type":"reasoning","encrypted_content":"fixture-opaque"}`, `{"type":"function_call","call_id":"a","name":"exec","arguments":"{}"}`), nil + case 2: + if int64(1) != gjson.GetBytes(body, "input.#").Int() || "one" != gjson.GetBytes(body, "previous_response_id").String() { + t.Error("healthy continuation repeated history") + } + return nativeWSResponse("two", `{"type":"function_call","call_id":"b","name":"exec","arguments":"{}"}`), nil + case 3: + return nativeWSContinuationFailure(failure), nativeWSContinuationError(failure) + default: + if gjson.GetBytes(body, "previous_response_id").Exists() || int64(7) != gjson.GetBytes(body, "input.#").Int() || "fixture-opaque" != gjson.GetBytes(body, "input.2.encrypted_content").String() { + t.Errorf("incomplete or duplicate recovery: %s", body) + } + return nativeWSResponse("recovered"), nil + } + }) + conn := fixture.connect(t) + assertNativeWSSuccess(t, nativeWSTurn(t, conn, `{"input":[{"type":"additional_tools","tools":[]},{"role":"user","content":"root"}]}`)) + assertNativeWSSuccess(t, nativeWSTurn(t, conn, `{"previous_response_id":"one","input":[{"type":"function_call_output","call_id":"a","output":"1"}]}`)) + conn.Close() + conn = fixture.connect(t) + assertNativeWSSuccess(t, nativeWSTurn(t, conn, `{"previous_response_id":"two","input":[{"type":"function_call_output","call_id":"b","output":"2"}]}`)) + if int32(4) != calls.Load() { + t.Fatal("recovery repeated") + } + }) + } +} + +func TestNativeWSEncryptedRecoveryRequiresCompatibleAccount(t *testing.T) { + resetResponseCacheForTest() + t.Cleanup(resetResponseCacheForTest) + owner := nativeWSCachePrefix + "fixture" + account := &auth.Account{DBID: 1, CredentialGeneration: 1} + cacheNativeWSContext(nativeWSCompletedContext{owner: owner, input: `[{"role":"user","content":"root"}]`, + completed: []byte(`{"response":{"id":"encrypted","output":[{"type":"reasoning","encrypted_content":"fixture"}]}}`), + provenance: nativeWSAccountProvenance(account)}) + body := []byte(`{"previous_response_id":"encrypted","input":[{"role":"user","content":"next"}]}`) + source := newResponsesWSReplaySource(body, owner) + source.setAccount(account) + input := source.Input() + if "fixture" != gjson.Get(input, "1.encrypted_content").String() { + t.Fatal("compatible recovery discarded encrypted reasoning") + } + recovered := newResponsesWSReplaySourceFromRecovery(input, source) + for _, changed := range []*auth.Account{{DBID: 2, CredentialGeneration: 1}, {DBID: 1, CredentialGeneration: 2}} { + for _, snapshot := range []*responsesWSReplaySource{source, recovered} { + snapshot.setAccount(changed) + if "" != snapshot.Input() || snapshot.err == nil { + t.Fatal("encrypted context replayed on incompatible account") + } + } + } +} + +func TestNativeWSCacheSeparatesStreamAndExecutionLane(t *testing.T) { + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil) + identity := requestSessionIdentity{affinityID: "thread", explicitUpstreamID: "thread"} + root := nativeWSTurnCacheOwner(c, "key:1", nativeWSTurnScope{body: []byte(`{"stream_id":"main"}`), identity: identity}) + for _, body := range []string{`{"stream_id":"other"}`, `{"stream_id":"main","client_metadata":{"x-codex-turn-metadata":"{\"request_kind\":\"memory\"}"}}`} { + if root == nativeWSTurnCacheOwner(c, "key:1", nativeWSTurnScope{body: []byte(body), identity: identity}) { + t.Fatalf("execution lane shared native context: %s", body) + } + } + for _, key := range []string{"fixture-key-one", "fixture-key-two"} { + c.Request.Header.Set("Authorization", "Bearer "+key) + other := nativeWSTurnCacheOwner(c, "key:1", nativeWSTurnScope{body: []byte(`{"stream_id":"main"}`), identity: identity}) + if root == other { + t.Fatal("static API keys shared native context") + } + root = other + } +} + +func nativeWSContinuationFailure(failure string) *http.Response { + if failure == "binding_lost" { + return nil + } + return &http.Response{StatusCode: http.StatusBadRequest, Header: make(http.Header), Body: io.NopCloser(strings.NewReader("{\"error\":{\"type\":\"invalid_request_error\",\"message\":\"Invalid `previous_response_id`.\"}}"))} +} + +func nativeWSContinuationError(failure string) error { + if failure == "binding_lost" { + return &ResponsesContinuationLostError{Reason: "original_connection_unavailable"} + } + return nil +} + +func assertNativeWSSuccess(t *testing.T, terminal []byte) { + t.Helper() + if "response.completed" != gjson.GetBytes(terminal, "type").String() { + t.Fatalf("unexpected terminal: %s", terminal) + } +} + +func TestNativeWSMissingContextKeepsDownstreamUsable(t *testing.T) { + fixture := newNativeWSFixture(t, func(_ context.Context, body []byte) (*http.Response, error) { + if gjson.GetBytes(body, "previous_response_id").Exists() { + return nil, &ResponsesContinuationLostError{Reason: "original_connection_unavailable"} + } + return nativeWSResponse("fresh"), nil + }) + conn := fixture.connect(t) + terminal := nativeWSTurn(t, conn, `{"previous_response_id":"missing","input":[{"role":"user","content":"next"}]}`) + if "previous_response_not_found" != gjson.GetBytes(terminal, "error.code").String() || int64(409) != gjson.GetBytes(terminal, "status").Int() { + t.Fatalf("missing continuation classification: %s", terminal) + } + assertNativeWSSuccess(t, nativeWSTurn(t, conn, `{"input":[{"role":"user","content":"full history"}]}`)) +} + +func TestNativeWSCacheRejectsIncompleteSnapshotsAndIsolatesLegacyNamespace(t *testing.T) { + resetResponseCacheForTest() + t.Cleanup(resetResponseCacheForTest) + backend := newRecordingResponseContextBackend(true) + SetResponseContextCache(backend) + t.Cleanup(func() { backend.TokenCache.Close() }) + owner := nativeWSCachePrefix + "key:1" + input := `[{"role":"user","content":"root"}]` + completed := []byte(`{"response":{"id":"one","output":[]}}`) + cacheResponsesWSCompletedResponse(owner, input, completed, []json.RawMessage{json.RawMessage(`{"type":"reasoning","encrypted_content":"opaque"}`)}) + next := []byte(`{"previous_response_id":"one","input":[]}`) + if _, _, err := degradeResponsesWSContinuationWithSource(next, owner, newResponsesWSReplaySource(next, owner)); err == nil || err.Code != "previous_response_not_found" { + t.Fatal("opaque reasoning was silently removed") + } + missing := []byte(`{"previous_response_id":"missing","input":[]}`) + _, _, _ = degradeResponsesWSContinuationWithSource(missing, owner, newResponsesWSReplaySource(missing, owner)) + drainResponseCacheBackendWrites() + if sets, gets := backend.counts(); sets != 0 || gets != 0 { + t.Fatalf("native context accessed legacy namespace: writes=%d reads=%d", sets, gets) + } + config := defaultResponseCacheConfig() + config.maxItems = 1 + configureResponseCacheForTest(config) + cacheResponsesWSCompletedResponse(owner, input, completed, []json.RawMessage{json.RawMessage(`{"role":"assistant","content":"reply"}`)}) + if responseCacheLookupHit == nativeWSLocalLookup(owner, "one").Kind { + t.Fatal("tail-trimmed context admitted") + } +} + +func TestNativeWSCommitBarrierAndScope(t *testing.T) { + owner := nativeWSCacheOwner("key:1", requestSessionIdentity{affinityID: "thread", explicitUpstreamID: "thread"}) + other := nativeWSCacheOwner("key:1", requestSessionIdentity{affinityID: "other", explicitUpstreamID: "other"}) + if owner == other || owner == nativeWSCacheOwner("key:2", requestSessionIdentity{affinityID: "thread", explicitUpstreamID: "thread"}) { + t.Fatal("cache scope is not isolated") + } + finish := beginNativeWSCommit(owner, "pending") + defer finish() + ctx, cancel := context.WithCancel(context.Background()) + cancel() + if err := waitNativeWSCommit(ctx, owner, "pending"); err != context.Canceled { + t.Fatalf("commit barrier ignored cancellation: %v", err) + } + finish() + if err := waitNativeWSCommit(context.Background(), owner, "pending"); err != nil { + t.Fatal(err) + } +} + +func TestPreviousIDErrorRecognitionIsNarrow(t *testing.T) { + for _, message := range []string{"Invalid `previous_response_id`.", "Invalid previous_response_id", "Invalid model", "Invalid request", "previous_response_id format unsupported"} { + payload, _ := json.Marshal(map[string]any{"error": map[string]string{"type": "invalid_request_error", "message": message}}) + want := strings.HasPrefix(message, "Invalid ") && strings.Contains(message, "previous_response_id") + if want != isPreviousResponseNotFoundBody(payload) { + t.Fatal(fmt.Sprintf("wrong classification: %s", payload)) + } + } +} diff --git a/proxy/responses_ws_recovery_error.go b/proxy/responses_ws_recovery_error.go new file mode 100644 index 000000000..4b42fcc54 --- /dev/null +++ b/proxy/responses_ws_recovery_error.go @@ -0,0 +1,31 @@ +package proxy + +import ( + "log" + "net/http" + + "github.com/codex2api/api" + "github.com/gorilla/websocket" +) + +func responsesWSErrorStatus(err *api.APIError) int { + switch string(err.Code) { + case "previous_response_not_found", string(api.ErrCodeResponseContextUnavailable): + return http.StatusConflict + case ErrorCodeAccountPoolConcurrencySaturated, ErrorCodeUpstreamWSConnectionCapacity: + return http.StatusServiceUnavailable + } + if err.Type == api.ErrorTypeRateLimit { + return http.StatusTooManyRequests + } + if err.Type == api.ErrorTypeInvalidRequest && api.HTTPStatusCode(err.Code) == http.StatusInternalServerError { + return http.StatusBadRequest + } + return api.HTTPStatusCode(err.Code) +} + +func writeResponsesWSRecoveryError(conn *websocket.Conn, err *api.APIError) error { + recordResponseCacheKnownUnavailableError() + log.Printf("Responses WebSocket context recovery unavailable: code=%s reason=%s", err.Code, responsesWSContextReason(err)) + return writeResponsesWSError(conn, err) +} diff --git a/proxy/responses_ws_shared_context.go b/proxy/responses_ws_shared_context.go new file mode 100644 index 000000000..604f1d5c4 --- /dev/null +++ b/proxy/responses_ws_shared_context.go @@ -0,0 +1,253 @@ +package proxy + +import ( + "context" + "encoding/json" + "errors" + "log" + "strings" + "time" + + "github.com/codex2api/auth" + "github.com/codex2api/cache" + "github.com/google/uuid" +) + +const ( + nativeWSContextNamespace = "native-ws-context" + nativeWSPendingNamespace = "native-ws-context-pending" + nativeWSContextVersion = 1 + nativeWSMetadataBytes = int64(4096) + nativeWSDomainMaxBytes = 1024 + nativeWSPendingMaxBytes = int64(64) + nativeWSCommitPoll = 100 * time.Millisecond +) + +type nativeWSSharedContext struct { + Version int `json:"version"` + Items []json.RawMessage `json:"items"` + AccountID int64 `json:"account_id,omitempty"` + Generation int64 `json:"generation,omitempty"` + Domain string `json:"domain,omitempty"` + ExpiresAt time.Time `json:"expires_at"` + Unavailable string `json:"unavailable,omitempty"` +} + +type nativeWSSharedWrite struct { + key string + response string + serial uint64 + record nativeWSSharedContext +} + +func nativeWSReplayAccountFilter(filter auth.AccountFilter, source *responsesWSReplaySource) auth.AccountFilter { + if source == nil || source.previous == nil || source.previous.nativeProvenance == nil { + return filter + } + previous := source.previous + hasOpaque := false + for _, item := range previous.Items { + if strings.Contains(string(item), `"encrypted_content"`) { + hasOpaque = true + break + } + } + if !hasOpaque || previous.nativeProvenance.accountID <= 0 { + return filter + } + return func(account *auth.Account) bool { + return account.ID() == previous.nativeProvenance.accountID && (filter == nil || filter(account)) + } +} + +func nativeWSSharedBackend() cache.TokenCache { + respCache.mu.RLock() + defer respCache.mu.RUnlock() + return respCache.runtimeCache +} + +func nativeWSContextLookup(owner, responseID string) responseCacheLookupResult { + local := nativeWSLocalLookup(owner, responseID) + if local.Kind == responseCacheLookupHit || responseID == "" { + return local + } + backend := nativeWSSharedBackend() + if backend == nil { + return local + } + ctx, cancel := context.WithTimeout(context.Background(), responseCacheBackendSyncTimeout) + defer cancel() + respCache.mu.RLock() + config := respCache.config + respCache.mu.RUnlock() + key := responseCacheStoreKey(owner, responseID) + raw, found, err := readNativeWSRuntime(ctx, backend, nativeWSRuntimeRead{namespace: nativeWSContextNamespace, key: key, + maxBytes: responseContextWireLimit(config) + nativeWSMetadataBytes}) + if errors.Is(err, cache.ErrRuntimeSnapshotTooLarge) { + return responseCacheLookupResult{Kind: responseCacheLookupReconstructionTooLarge} + } + if err != nil { + return responseCacheLookupResult{Kind: responseCacheLookupBackendError, Err: err} + } + if !found { + if local.Kind == responseCacheLookupBackendError { + return local + } + return responseCacheLookupResult{Kind: responseCacheLookupMiss, remoteMiss: true} + } + return hydrateNativeWSContext(key, raw, config) +} + +type nativeWSRuntimeRead struct { + namespace string + key string + maxBytes int64 +} + +func readNativeWSRuntime(ctx context.Context, backend cache.TokenCache, read nativeWSRuntimeRead) ([]byte, bool, error) { + if bounded, ok := backend.(cache.BoundedRuntimeReader); ok { + return bounded.GetRuntimeBounded(ctx, read.namespace, read.key, read.maxBytes) + } + raw, found, err := backend.GetRuntime(ctx, read.namespace, read.key) + if int64(len(raw)) > read.maxBytes { + return nil, false, cache.ErrRuntimeSnapshotTooLarge + } + return raw, found, err +} + +func hydrateNativeWSContext(key string, raw []byte, config responseCacheConfig) responseCacheLookupResult { + var record nativeWSSharedContext + if json.Unmarshal(raw, &record) != nil || record.Version != nativeWSContextVersion || record.ExpiresAt.IsZero() || len(record.Domain) > nativeWSDomainMaxBytes { + return responseCacheLookupResult{Kind: responseCacheLookupBackendCorrupt} + } + if !time.Now().Before(record.ExpiresAt) { + return responseCacheLookupResult{Kind: responseCacheLookupExpired} + } + if record.Unavailable != "" { + kind := responseCacheLookupBackendError + if record.Unavailable == "oversize" { + kind = responseCacheLookupReconstructionTooLarge + } + return responseCacheLookupResult{Kind: kind} + } + items, err := cache.NormalizeResponseContextItems(record.Items) + if err != nil { + return responseCacheLookupResult{Kind: responseCacheLookupBackendCorrupt} + } + if len(items) > config.maxItems || responseContextLogicalBytes(items) > config.reconstructMaxBytes { + return responseCacheLookupResult{Kind: responseCacheLookupReconstructionTooLarge} + } + provenance := &nativeWSReplayProvenance{accountID: record.AccountID, generation: record.Generation, domain: record.Domain} + items, admitted, oversize, serial := admitResponseCacheWithTicket(key, items) + setNativeWSLocalMetadata(key, nativeWSLocalMetadata{serial: serial, provenance: provenance, expiresAt: record.ExpiresAt}) + return responseCacheLookupResult{Kind: responseCacheLookupHit, Source: responseCacheSourceBackend, + Items: items, nativeProvenance: provenance, Promoted: admitted, oversizeBypass: !admitted && oversize} +} + +type nativeWSLocalMetadata struct { + serial uint64 + provenance *nativeWSReplayProvenance + expiresAt time.Time +} + +func setNativeWSLocalMetadata(key string, metadata nativeWSLocalMetadata) { + respCache.mu.Lock() + defer respCache.mu.Unlock() + if entry := respCache.store[key]; entry != nil && entry.serial == metadata.serial { + entry.nativeProvenance = metadata.provenance + entry.expiresAt = metadata.expiresAt + } +} + +func persistNativeWSContext(write nativeWSSharedWrite) { + backend := nativeWSSharedBackend() + if backend == nil { + return + } + if err := writeNativeWSContext(backend, write); err != nil { + write.record.Items = nil + write.record.Unavailable = "backend_error" + ctx, cancel := context.WithTimeout(context.Background(), responseCacheBackendSyncTimeout) + defer cancel() + _ = storeNativeWSRecord(ctx, backend, write) + } +} + +func writeNativeWSContext(backend cache.TokenCache, write nativeWSSharedWrite) (err error) { + writer := responseCacheBackendWriter + writer.start() + defer writer.finish() + defer func() { recordResponseCacheBackendWriteResult(write.key, write.response, write.serial, err) }() + bytes := responseContextLogicalBytes(write.record.Items) + waitCtx, cancel := context.WithTimeout(context.Background(), responseCacheBackendWaitTimeout) + _, err = writer.acquire(waitCtx, bytes) + cancel() + if err != nil { + return err + } + defer writer.release(bytes) + ctx, cancel := context.WithTimeout(context.Background(), responseCacheBackendWriteTimeout) + defer cancel() + return storeNativeWSRecord(ctx, backend, write) +} + +func storeNativeWSRecord(ctx context.Context, backend cache.TokenCache, write nativeWSSharedWrite) error { + raw, err := json.Marshal(write.record) + if err != nil { + return err + } + ttl := time.Until(write.record.ExpiresAt) + if ttl <= 0 { + return nil + } + return backend.SetRuntime(ctx, nativeWSContextNamespace, write.key, raw, ttl) +} + +func beginNativeWSSharedCommit(key string) func() { + backend := nativeWSSharedBackend() + owners, ok := backend.(cache.RuntimeOwnerStore) + if !ok { + return func() {} + } + owner, _ := json.Marshal(uuid.NewString()) + ctx, cancel := context.WithTimeout(context.Background(), responseCacheBackendSyncTimeout) + _, err := owners.ClaimRuntimeOwner(ctx, nativeWSPendingNamespace, key, owner, nativeWSCommitWait+responseCacheBackendWaitTimeout+responseCacheBackendWriteTimeout) + cancel() + if err != nil { + log.Printf("登记 Redis 原生 WS 提交屏障失败: %v", err) + return func() {} + } + return func() { + ctx, cancel := context.WithTimeout(context.Background(), responseCacheBackendSyncTimeout) + defer cancel() + _, _ = owners.CompareAndDeleteRuntimeOwner(ctx, nativeWSPendingNamespace, key, owner) + } +} + +func waitNativeWSSharedCommit(ctx context.Context, owner, responseID string) error { + backend := nativeWSSharedBackend() + if backend == nil || responseID == "" || !strings.HasPrefix(owner, nativeWSCachePrefix) { + return nil + } + if nativeWSLocalLookup(owner, responseID).Kind == responseCacheLookupHit { + return nil + } + waitCtx, cancel := context.WithTimeout(ctx, nativeWSCommitWait) + defer cancel() + ticker := time.NewTicker(nativeWSCommitPoll) + defer ticker.Stop() + for { + readCtx, readCancel := context.WithTimeout(waitCtx, responseCacheBackendSyncTimeout) + _, pending, err := readNativeWSRuntime(readCtx, backend, + nativeWSRuntimeRead{namespace: nativeWSPendingNamespace, key: responseCacheStoreKey(owner, responseID), maxBytes: nativeWSPendingMaxBytes}) + readCancel() + if err != nil || !pending { + return err + } + select { + case <-waitCtx.Done(): + return waitCtx.Err() + case <-ticker.C: + } + } +} diff --git a/proxy/review_regression_test.go b/proxy/review_regression_test.go index 035e21081..a1ed54190 100644 --- a/proxy/review_regression_test.go +++ b/proxy/review_regression_test.go @@ -277,8 +277,11 @@ func newReviewWSClient(t *testing.T, execute func(context.Context, *auth.Account func TestReviewWSHTTPFallbackMissingContextReleasesLease(t *testing.T) { resetResponseCacheForTest() - conn, account := newReviewWSClient(t, func(context.Context, *auth.Account, []byte, string, string, string, *DeviceProfileConfig, http.Header, string) (*http.Response, error) { - return nil, errors.New("websocket: close 1009 (message too big)") + conn, account := newReviewWSClient(t, func(_ context.Context, _ *auth.Account, body []byte, _, _, _ string, _ *DeviceProfileConfig, _ http.Header, _ string) (*http.Response, error) { + if gjson.GetBytes(body, "previous_response_id").Exists() { + return nil, errors.New("websocket: close 1009 (message too big)") + } + return nativeWSResponse("fresh"), nil }) if err := conn.WriteMessage(websocket.TextMessage, []byte(`{"type":"response.create","model":"gpt-5.5","previous_response_id":"missing","input":[{"type":"message","role":"user","content":"next"}]}`)); err != nil { t.Fatal(err) @@ -288,13 +291,13 @@ func TestReviewWSHTTPFallbackMissingContextReleasesLease(t *testing.T) { if err != nil || gjson.GetBytes(body, "error.code").String() != "response_context_unavailable" { t.Fatalf("missing fail-closed error: %s %v", body, err) } - _, _, err = conn.ReadMessage() - if !websocket.IsCloseError(err, websocket.ClosePolicyViolation) || account.GetActiveRequests() != 0 { - t.Fatalf("fallback lease leaked: %d, %v", account.GetActiveRequests(), err) + if int64(http.StatusConflict) != gjson.GetBytes(body, "status").Int() || int64(0) != account.GetActiveRequests() { + t.Fatalf("fallback status or lease invalid: %s, active=%d", body, account.GetActiveRequests()) } if GetResponseCacheStats().KnownUnavailableErrors != 1 { t.Fatal("WS failure was not counted") } + assertNativeWSSuccess(t, nativeWSTurn(t, conn, `{"input":[{"role":"user","content":"fresh history"}]}`)) } func TestReviewWSStoreFalseSkipsAlwaysCache(t *testing.T) { diff --git a/proxy/scheduler_queue_error_test.go b/proxy/scheduler_queue_error_test.go index 27863f812..08c9f3ba0 100644 --- a/proxy/scheduler_queue_error_test.go +++ b/proxy/scheduler_queue_error_test.go @@ -220,9 +220,15 @@ func TestSchedulerQueueOverloadResponsesWebSocket(t *testing.T) { if json.Unmarshal(payload, &event) != nil || event["type"] != "error" || !strings.Contains(string(payload), schedulerQueueFullMessage) { t.Fatalf("WS overload frame = %s", payload) } - _, _, err = conn.ReadMessage() - if !websocket.IsCloseError(err, websocket.CloseTryAgainLater) { - t.Fatalf("WS overload close = %v", err) + if float64(http.StatusServiceUnavailable) != event["status"] { + t.Fatalf("WS overload status = %v", event["status"]) + } + if err := conn.WriteJSON(map[string]any{"type": "response.create", "model": "gpt-5.5", "input": "retry"}); err != nil { + t.Fatal(err) + } + _, payload, err = conn.ReadMessage() + if err != nil || !strings.Contains(string(payload), schedulerQueueFullMessage) { + t.Fatalf("WS overload prevented another turn: %s %v", payload, err) } } diff --git a/proxy/upstream_ws_capacity.go b/proxy/upstream_ws_capacity.go new file mode 100644 index 000000000..8528a095f --- /dev/null +++ b/proxy/upstream_ws_capacity.go @@ -0,0 +1,22 @@ +package proxy + +import ( + "fmt" + "net/http" +) + +const ErrorCodeUpstreamWSConnectionCapacity = "upstream_ws_connection_capacity_exceeded" + +// UpstreamWSConnectionCapacityError 是本地连接容量拒绝，不是上游账号故障。 +type UpstreamWSConnectionCapacityError struct { + Limit int +} + +func (err *UpstreamWSConnectionCapacityError) Error() string { + return fmt.Sprintf("Upstream WebSocket connection capacity reached (max %d); existing continuations are protected, please retry later", err.Limit) +} + +func (err *UpstreamWSConnectionCapacityError) Unwrap() error { + return &Error{Code: ErrorCodeUpstreamWSConnectionCapacity, Type: ErrorTypeServerError, + HTTPStatus: http.StatusServiceUnavailable, Message: err.Error()} +} From 7c210fa54c394caf3c378c05913c27f788163c2f Mon Sep 17 00:00:00 2001 From: bxb1337 Date: Fri, 9 Oct 2026 19:46:26 +0800 Subject: [PATCH 3/4] =?UTF-8?q?fix(wsrelay):=20=E5=88=86=E7=A6=BB=E8=AF=B7?= =?UTF-8?q?=E6=B1=82=E5=B9=B6=E5=8F=91=E4=B8=8E=E8=BF=9E=E6=8E=A5=E4=BF=9D?= =?UTF-8?q?=E7=95=99=E5=B9=B6=E4=BF=9D=E6=8A=A4=E5=9C=A8=E9=80=94=E7=BB=AD?= =?UTF-8?q?=E9=93=BE?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- proxy/wsrelay/acquire_wakeup_test.go | 8 +- proxy/wsrelay/busy_policy.go | 9 +- proxy/wsrelay/chat_retention.go | 79 +++ proxy/wsrelay/client_identity.go | 3 + proxy/wsrelay/conn_lifetime_test.go | 5 +- proxy/wsrelay/connection_budget.go | 137 ++++++ proxy/wsrelay/connection_capacity.go | 242 ++++++++++ proxy/wsrelay/connection_close.go | 114 +++++ proxy/wsrelay/connection_dial.go | 97 ++++ proxy/wsrelay/continuation.go | 79 +++ proxy/wsrelay/continuation_test.go | 198 ++++++++ proxy/wsrelay/executor.go | 70 +-- .../wsrelay/issue436_pending_protect_test.go | 1 + proxy/wsrelay/keepalive.go | 5 +- proxy/wsrelay/manager.go | 450 +++++------------- proxy/wsrelay/manager_test.go | 43 +- proxy/wsrelay/oneshot_conn_test.go | 23 +- proxy/wsrelay/read_pump.go | 11 +- proxy/wsrelay/session.go | 26 +- proxy/wsrelay/session_activity.go | 37 ++ proxy/wsrelay/session_test.go | 7 +- proxy/wsrelay/subagent_transport_lane_test.go | 15 +- 22 files changed, 1228 insertions(+), 431 deletions(-) create mode 100644 proxy/wsrelay/chat_retention.go create mode 100644 proxy/wsrelay/connection_budget.go create mode 100644 proxy/wsrelay/connection_capacity.go create mode 100644 proxy/wsrelay/connection_close.go create mode 100644 proxy/wsrelay/connection_dial.go create mode 100644 proxy/wsrelay/continuation.go create mode 100644 proxy/wsrelay/continuation_test.go create mode 100644 proxy/wsrelay/session_activity.go diff --git a/proxy/wsrelay/acquire_wakeup_test.go b/proxy/wsrelay/acquire_wakeup_test.go index 07b0334f0..99f3e65b1 100644 --- a/proxy/wsrelay/acquire_wakeup_test.go +++ b/proxy/wsrelay/acquire_wakeup_test.go @@ -58,9 +58,13 @@ func TestAccountWaitSignalClosedByPoolChanges(t *testing.T) { // 拨号占位归还 → 唤醒 third := manager.accountWaitSignal(accountID) - manager.releaseAccountConnectionCapacity(accountID) + capacity := manager.reserveAccountConnectionCapacity(connectionCapacityRequest{accountID: accountID, limit: connectionLimit(capacityChat)}) + if capacity == nil { + t.Fatal("dial capacity reservation failed") + } + capacity.release() if !signalClosed(third) { - t.Fatal("releaseAccountConnectionCapacity did not wake account waiters") + t.Fatal("capacity release did not wake account waiters") } // 没有等待者时 notify 是空操作，不应留下条目。 diff --git a/proxy/wsrelay/busy_policy.go b/proxy/wsrelay/busy_policy.go index 832dfbc64..047dbdf79 100644 --- a/proxy/wsrelay/busy_policy.go +++ b/proxy/wsrelay/busy_policy.go @@ -55,11 +55,8 @@ func isBusyOverflowSessionKey(sessionKey string) bool { return strings.Contains(sessionKey, busyOverflowKeyInfix) } -// statelessConnectionSlots 返回无状态连接槽位数（系统设置 codex_ws_stateless_slots， -// 热更新生效，RuntimeSettings 已钳 1~32；0 值兜底回内置常量，issue #522）。 +// statelessConnectionSlots 返回账号空白连接预算及通用池扫描窗口（系统设置 codex_ws_stateless_slots， +// 热更新生效，范围 0~32；0 表示不保留空白连接、不扫描通用池）。 func statelessConnectionSlots() int { - if v := proxy.CurrentRuntimeSettings().CodexWSStatelessSlots; v > 0 { - return v - } - return StatelessConnectionSlots + return proxy.CurrentRuntimeSettings().CodexWSStatelessSlots } diff --git a/proxy/wsrelay/chat_retention.go b/proxy/wsrelay/chat_retention.go new file mode 100644 index 000000000..736646ab3 --- /dev/null +++ b/proxy/wsrelay/chat_retention.go @@ -0,0 +1,79 @@ +package wsrelay + +import "sort" + +// 调用方持有账号锁。推理、拨号及尚未承载聊天的连接不占空闲保留额度。 +func (m *Manager) trimIdleChatConnections(input connectionCapacityRequest) { + input.kind = capacityChat + idle := m.idleCapacityConnections(input) + sort.Slice(idle, func(i, j int) bool { + return idle[i].lastUsed < idle[j].lastUsed + }) + excess := len(idle) - input.limit + for _, candidate := range idle { + if excess <= 0 { + return + } + if m.discardIdleChatConnection(candidate.wc) { + excess-- + } + } +} + +// 即使存在续链绑定也可按 LRU 回收；账号锁阻止候选被新推理同时取走。 +func (m *Manager) discardIdleChatConnection(wc *WsConnection) bool { + m.respConnMu.Lock() + current, exists := m.connections.Load(wc.PoolKey) + if !exists || current != wc || wc.session.PendingCount() != 0 { + m.respConnMu.Unlock() + return false + } + wc.responseBindingsDisabled = true + m.respConnMu.Unlock() + m.discardConnectionFor(wc, closeIdleChatLRU) + return true +} + +func (m *Manager) retainReleasedConnection(wc *WsConnection) { + lock, release := m.accountLock(wc.session.AccountID) + defer release() + lock.Lock() + defer lock.Unlock() + current, exists := m.connections.Load(wc.PoolKey) + if !exists || current != wc || !wc.IsConnected() { + return + } + if connectionBudgetKind(wc) == capacityBlank && connectionLimit(capacityBlank) == 0 && wc.session.PendingCount() == 0 { + m.discardConnectionFor(wc, closeCapacityReclaim) + return + } + wc.Touch() + input := connectionCapacityRequest{accountID: wc.session.AccountID, limit: connectionLimit(capacityChat)} + m.trimIdleChatConnections(input) +} + +// 保存较小的设置后，即使没有新请求，后台巡检也会收缩空闲保留池。 +func (m *Manager) trimAllIdleChatConnections() { + accounts := make(map[int64]struct{}) + m.connections.Range(func(_, value any) bool { + wc := value.(*WsConnection) + if wc.session != nil { + accounts[wc.session.AccountID] = struct{}{} + } + return true + }) + for accountID := range accounts { + m.trimAccountIdleChatConnections(accountID) + } +} + +func (m *Manager) trimAccountIdleChatConnections(accountID int64) { + lock, release := m.accountLock(accountID) + defer release() + lock.Lock() + defer lock.Unlock() + m.trimIdleChatConnections(connectionCapacityRequest{accountID: accountID, limit: connectionLimit(capacityChat)}) + if connectionLimit(capacityBlank) == 0 { + m.reclaimConnectionCapacity(connectionCapacityRequest{accountID: accountID, kind: capacityBlank}) + } +} diff --git a/proxy/wsrelay/client_identity.go b/proxy/wsrelay/client_identity.go index 50bc47918..b9aca7cd6 100644 --- a/proxy/wsrelay/client_identity.go +++ b/proxy/wsrelay/client_identity.go @@ -26,6 +26,9 @@ type websocketContinuation struct { accountID int64 apiKey string identity string + model string + url string + proxyURL string } func (e *Executor) acquireClientContinuation(input websocketContinuation) (*WsConnection, *PendingRequest, string) { diff --git a/proxy/wsrelay/conn_lifetime_test.go b/proxy/wsrelay/conn_lifetime_test.go index 9d01fd249..542b3b497 100644 --- a/proxy/wsrelay/conn_lifetime_test.go +++ b/proxy/wsrelay/conn_lifetime_test.go @@ -77,9 +77,8 @@ func TestEvictExpiredRotatesIdleOverAgeConnection(t *testing.T) { } } -// TestPingIdleConnectionsSkipsOverAgeConnection 验证保活不给到龄连接续命： -// Pong 会刷新 lastUsed 让连接永不空闲过期，若继续 Ping，开保活的实例上所有 -// 连接最终都会撞上游 60 分钟寿命上限。 +// TestPingIdleConnectionsSkipsOverAgeConnection 验证到龄连接不再保活， +// 由清理器轮转，避免撞上游 60 分钟寿命上限。 func TestPingIdleConnectionsSkipsOverAgeConnection(t *testing.T) { m := NewManager() t.Cleanup(m.Stop) diff --git a/proxy/wsrelay/connection_budget.go b/proxy/wsrelay/connection_budget.go new file mode 100644 index 000000000..b17a413ff --- /dev/null +++ b/proxy/wsrelay/connection_budget.go @@ -0,0 +1,137 @@ +package wsrelay + +import ( + "context" + "errors" + "net/http" + + "github.com/codex2api/proxy" + "github.com/tidwall/gjson" +) + +type connectionCapacityKind int32 + +const ( + capacityChat connectionCapacityKind = iota + capacityBlank +) + +func (kind connectionCapacityKind) String() string { + if kind == capacityBlank { + return "blank" + } + return "chat" +} + +type connectionCapacityContextKey struct{} + +func requestConnectionCapacityKind(ctx context.Context) connectionCapacityKind { + if kind, ok := ctx.Value(connectionCapacityContextKey{}).(connectionCapacityKind); ok { + return kind + } + return capacityChat +} + +// 只有没有聊天内容的明确预热使用空白预算；携带上下文仍属于聊天连接。 +func withConnectionCapacityKind(ctx context.Context, headers http.Header, body []byte) context.Context { + kind := capacityChat + prewarm := gjson.GetBytes(body, "generate").Type == gjson.False || proxy.IsCodexWebsocketPrewarmRequest(headers, body) + if prewarm && !prewarmHasChatContext(body) { + kind = capacityBlank + } + return context.WithValue(ctx, connectionCapacityContextKey{}, kind) +} + +func prewarmHasChatContext(body []byte) bool { + for _, key := range []string{"input", "instructions", "tools", "prompt", "previous_response_id"} { + value := gjson.GetBytes(body, key) + if !value.Exists() || value.Type == gjson.Null { + continue + } + if value.IsArray() && len(value.Array()) == 0 || value.Type == gjson.String && value.String() == "" { + continue + } + return true + } + return false +} + +func connectionLimit(kind connectionCapacityKind) int { + if kind == capacityBlank { + return statelessConnectionSlots() + } + return proxy.CurrentRuntimeSettings().CodexWSDownstreamKeepaliveSlots +} + +func connectionBudgetKind(wc *WsConnection) connectionCapacityKind { + if slot := wc.capacitySlot.Load(); slot != nil { + return connectionCapacityKind(slot.kind.Load()) + } + if wc.session != nil && wc.session.hasUserContext() { + return capacityChat + } + return capacityBlank +} + +func (m *Manager) budgetConnectionCountLocked(accountID int64, kind connectionCapacityKind) int { + count := 0 + for slot := range m.connectionSlots[accountID] { + if connectionCapacityKind(slot.kind.Load()) == kind { + count++ + } + } + return count +} + +func (m *Manager) budgetConnectionCount(input connectionCapacityRequest) int { + m.capacityMu.Lock() + defer m.capacityMu.Unlock() + return m.budgetConnectionCountLocked(input.accountID, input.kind) +} + +// 首次用户使用空白连接前转为聊天连接；推理不受空闲保留额度限制。 +func (m *Manager) ensureConnectionBudget(wc *WsConnection, kind connectionCapacityKind) error { + m.adoptConnectionCapacity(wc) + if kind == capacityBlank || connectionBudgetKind(wc) == capacityChat { + return nil + } + m.capacityMu.Lock() + defer m.capacityMu.Unlock() + slot := wc.capacitySlot.Load() + if slot.released.Load() || !wc.IsConnected() { + return errors.New("websocket disconnected before budget transfer") + } + slot.kind.Store(int32(capacityChat)) + return nil +} + +// 未发送请求就被本地额度拒绝的新连接没有用户上下文，可退回空白预算。 +func (m *Manager) releaseUnusedChatCapacity(wc *WsConnection) { + if wc.session.hasUserContext() || wc.session.PendingCount() != 0 { + return + } + m.capacityMu.Lock() + slot := wc.capacitySlot.Load() + keep := slot != nil && !slot.released.Load() && wc.IsConnected() + if keep && connectionCapacityKind(slot.kind.Load()) == capacityChat { + keep = m.budgetConnectionCountLocked(slot.accountID, capacityBlank) < connectionLimit(capacityBlank) + if keep { + slot.kind.Store(int32(capacityBlank)) + } + } + m.capacityMu.Unlock() + if !keep { + m.discardConnectionFor(wc, closeCapacityReclaim) + } +} + +func isConnectionCapacityError(err error) bool { + var capacityErr *proxy.UpstreamWSConnectionCapacityError + return errors.As(err, &capacityErr) +} + +func (m *Manager) discardAcquisitionFailure(wc *WsConnection, err error) { + if !isConnectionCapacityError(err) { + m.DiscardConnection(wc) + } +} diff --git a/proxy/wsrelay/connection_capacity.go b/proxy/wsrelay/connection_capacity.go new file mode 100644 index 000000000..19b9efa83 --- /dev/null +++ b/proxy/wsrelay/connection_capacity.go @@ -0,0 +1,242 @@ +package wsrelay + +import ( + "sort" + "sync/atomic" + "time" + + "github.com/codex2api/proxy" +) + +type idleAccountConnection struct { + wc *WsConnection + lastUsed int64 +} + +func isOneShotPoolConn(wc *WsConnection) bool { + return wc != nil && wc.session != nil && proxy.IsStatelessWebsocketSessionID(wc.session.ID) +} + +// 无绑定的一次性连接先回收，其余按业务使用时间排序。 +func sortIdleForEviction(idle []idleAccountConnection) { + sort.Slice(idle, func(i, j int) bool { + one, other := isOneShotPoolConn(idle[i].wc), isOneShotPoolConn(idle[j].wc) + if one != other { + return one + } + return idle[i].lastUsed < idle[j].lastUsed + }) +} + +func (m *Manager) hasLiveResponseBinding(wc *WsConnection) bool { + if m == nil || wc == nil { + return false + } + m.respConnMu.Lock() + defer m.respConnMu.Unlock() + return m.hasLiveResponseBindingLocked(wc) +} + +func (m *Manager) hasLiveResponseBindingLocked(wc *WsConnection) bool { + if wc.session != nil && !wc.session.hasUserContext() { + return false + } + now := time.Now() + for _, binding := range m.respConnBindings { + if binding.conn == wc && now.Before(binding.expiresAt) { + return true + } + } + return false +} + +// 实体连接登记覆盖拨号、入池和摘池后的收尾，直到 socket 实际关闭。 +type connectionCapacitySlot struct { + manager *Manager + accountID int64 + kind atomic.Int32 + released atomic.Bool +} + +func (slot *connectionCapacitySlot) release() { + if slot == nil || !slot.released.CompareAndSwap(false, true) { + return + } + m := slot.manager + m.capacityMu.Lock() + delete(m.connectionSlots[slot.accountID], slot) + if len(m.connectionSlots[slot.accountID]) == 0 { + delete(m.connectionSlots, slot.accountID) + } + m.capacityMu.Unlock() + m.notifyAccountWaiters(slot.accountID) +} + +func (slot *connectionCapacitySlot) attach(wc *WsConnection) { + m := slot.manager + m.capacityMu.Lock() + wc.capacitySlot.Store(slot) + m.connectionSlots[slot.accountID][slot] = wc + m.capacityMu.Unlock() + if wc.state.Load() == int32(StateDisconnected) { + slot.release() + } +} + +func (m *Manager) newCapacitySlotLocked(accountID int64) *connectionCapacitySlot { + if m.connectionSlots == nil { + m.connectionSlots = make(map[int64]map[*connectionCapacitySlot]*WsConnection) + } + if m.connectionSlots[accountID] == nil { + m.connectionSlots[accountID] = make(map[*connectionCapacitySlot]*WsConnection) + } + slot := &connectionCapacitySlot{manager: m, accountID: accountID} + m.connectionSlots[accountID][slot] = nil + return slot +} + +// 兼容已存在于池中的连接；生产拨号在永久 reader 启动前完成登记。 +func (m *Manager) adoptConnectionCapacity(wc *WsConnection) { + kind := connectionBudgetKind(wc) + m.capacityMu.Lock() + if wc.capacitySlot.Load() != nil { + m.capacityMu.Unlock() + return + } + slot := m.newCapacitySlotLocked(wc.session.AccountID) + slot.kind.Store(int32(kind)) + wc.capacitySlot.Store(slot) + m.connectionSlots[slot.accountID][slot] = wc + m.capacityMu.Unlock() + if wc.state.Load() == int32(StateDisconnected) { + slot.release() + } +} + +func (m *Manager) accountConnectionCount(accountID int64) int { + m.capacityMu.Lock() + defer m.capacityMu.Unlock() + return len(m.connectionSlots[accountID]) +} + +type connectionCapacityRequest struct { + accountID int64 + limit int + protectedKey string + extraPending int + freeSlots int + kind connectionCapacityKind +} + +func (m *Manager) idleCapacityConnections(input connectionCapacityRequest) []idleAccountConnection { + var idle []idleAccountConnection + m.connections.Range(func(_, value any) bool { + wc, ok := value.(*WsConnection) + if !ok || wc == nil || wc.session == nil || wc.session.AccountID != input.accountID { + return true + } + m.adoptConnectionCapacity(wc) + if connectionBudgetKind(wc) != input.kind { + return true + } + if wc.session.PendingCount() != 0 { + return true + } + if input.kind == capacityChat && !wc.session.hasUserContext() { + return true + } + if !wc.IsConnected() || wc.IsExpired() || wc.IsOverAge() { + m.discardConnectionFor(wc, connectionExpiryReason(wc)) + return true + } + if input.kind == capacityChat || wc.PoolKey != input.protectedKey && !m.hasLiveResponseBinding(wc) { + idle = append(idle, idleAccountConnection{wc: wc, lastUsed: wc.lastUsed.Load()}) + } + return true + }) + return idle +} + +// 绑定检查与禁止晚到绑定原子完成，避免终态绑定和容量回收交错。 +func (m *Manager) discardUnboundIdleConnection(wc *WsConnection) { + m.respConnMu.Lock() + if wc.session.PendingCount() != 0 || m.hasLiveResponseBindingLocked(wc) { + m.respConnMu.Unlock() + return + } + current, exists := m.connections.Load(wc.PoolKey) + if !exists || current != wc { + m.respConnMu.Unlock() + return + } + wc.responseBindingsDisabled = true + m.respConnMu.Unlock() + m.discardConnectionFor(wc, closeCapacityReclaim) +} + +func (m *Manager) reclaimConnectionCapacity(input connectionCapacityRequest) bool { + if input.limit < 0 { + input.limit = 0 + } + if input.kind == capacityChat { + m.trimIdleChatConnections(input) + return true + } + idle := m.idleCapacityConnections(input) + sortIdleForEviction(idle) + for _, candidate := range idle { + if input.limit > 0 && m.budgetConnectionCount(input)+input.extraPending+input.freeSlots <= input.limit { + return true + } + m.discardUnboundIdleConnection(candidate.wc) + } + // 0 关闭空白保留池，但仍允许请求临时拨号；收尾时关闭连接。 + return input.limit == 0 || m.budgetConnectionCount(input)+input.extraPending+input.freeSlots <= input.limit +} + +func (m *Manager) ensureAccountConnectionCapacity(input connectionCapacityRequest) bool { + input.freeSlots = 1 + return m.reclaimConnectionCapacity(input) +} + +func (m *Manager) trimIdleAccountConnections(accountID int64, limit int, protected *WsConnection) { + input := connectionCapacityRequest{accountID: accountID, limit: limit} + if protected != nil { + input.protectedKey = protected.PoolKey + input.kind = connectionBudgetKind(protected) + } + m.reclaimConnectionCapacity(input) +} + +// 调用方持有账号锁；聊天拨号只登记实体连接，空白拨号仍检查原预算。 +func (m *Manager) reserveAccountConnectionCapacity(input connectionCapacityRequest) *connectionCapacitySlot { + if !m.ensureAccountConnectionCapacity(input) { + return nil + } + m.capacityMu.Lock() + defer m.capacityMu.Unlock() + select { + case <-m.stopCleanup: + return nil + default: + slot := m.newCapacitySlotLocked(input.accountID) + slot.kind.Store(int32(input.kind)) + return slot + } +} + +func (m *Manager) closeCapacityConnections() { + var connections []*WsConnection + m.capacityMu.Lock() + for _, slots := range m.connectionSlots { + for _, wc := range slots { + if wc != nil { + connections = append(connections, wc) + } + } + } + m.capacityMu.Unlock() + for _, wc := range connections { + m.discardConnectionFor(wc, closeManagerShutdown) + } +} diff --git a/proxy/wsrelay/connection_close.go b/proxy/wsrelay/connection_close.go new file mode 100644 index 000000000..681c9aff0 --- /dev/null +++ b/proxy/wsrelay/connection_close.go @@ -0,0 +1,114 @@ +package wsrelay + +import ( + "context" + "errors" + "log" + "time" + + "github.com/codex2api/proxy" + "github.com/gorilla/websocket" +) + +const ( + closeCapacityReclaim = "capacity_reclaim" + closeIdleChatLRU = "idle_chat_lru" + closeIdleTimeout = "idle_timeout" + closeLifetimeRotation = "lifetime_rotation" + closeDisconnected = "disconnected" + closeManagerShutdown = "manager_shutdown" + closeProbeFailed = "probe_failed" + closePoolRemoved = "pool_removed" + closeUnusable = "unusable_connection" + closeUnconsumed = "unconsumed_response" + closeOneShotComplete = "unbound_oneshot_complete" +) + +type connectionCloseCause struct { + source string + reason string + code int + bound bool +} + +func (wc *WsConnection) recordCloseCause(cause connectionCloseCause) { + wc.closeCause.CompareAndSwap(nil, &cause) +} + +func (wc *WsConnection) recordDirectClose() { + if wc.closeCause.Load() != nil { + return + } + cause := connectionCloseCause{source: "local", reason: "explicit_close"} + if slot := wc.capacitySlot.Load(); slot != nil { + cause.bound = slot.manager.hasLiveResponseBinding(wc) + } + wc.recordCloseCause(cause) +} + +func (m *Manager) discardConnectionFor(wc *WsConnection, reason string) { + if wc == nil { + return + } + wc.recordCloseCause(connectionCloseCause{source: "local", reason: reason, bound: m.hasLiveResponseBinding(wc)}) + m.DiscardConnection(wc) +} + +func connectionExpiryReason(wc *WsConnection) string { + if !wc.IsConnected() { + return closeDisconnected + } + if wc.IsExpired() { + return closeIdleTimeout + } + return closeLifetimeRotation +} + +func (wc *WsConnection) recordReadClose(err error) { + cause := connectionCloseCause{source: "network", reason: "read_failure"} + var peer *websocket.CloseError + if errors.As(err, &peer) { + cause.code = peer.Code + if peer.Code != websocket.CloseAbnormalClosure && peer.Code != websocket.CloseNoStatusReceived && peer.Code != websocket.CloseTLSHandshake { + cause.source, cause.reason = "upstream", "peer_close" + } + } + if slot := wc.capacitySlot.Load(); slot != nil { + cause.bound = slot.manager.hasLiveResponseBinding(wc) + } + wc.recordCloseCause(cause) +} + +func (wc *WsConnection) logClose() { + cause := wc.closeCause.Load() + if cause == nil { + cause = &connectionCloseCause{source: "local", reason: "explicit_close"} + } + accountID, pending, remaining := int64(0), 0, 0 + if wc.session != nil { + accountID, pending = wc.session.AccountID, wc.session.PendingCount() + } + if slot := wc.capacitySlot.Load(); slot != nil { + remaining = slot.manager.accountConnectionCount(accountID) + } + age, idle := time.Duration(0), time.Duration(0) + if wc.createdAt > 0 { + age = time.Since(time.Unix(0, wc.createdAt)) + } + if used := wc.lastUsed.Load(); used > 0 { + idle = time.Since(time.Unix(0, used)) + } + log.Printf("[WS] 连接关闭 account=%d connection=%p source=%s reason=%s code=%d bound=%t pending=%d age_ms=%d idle_ms=%d connections=%d chat_idle_limit=%d blank_limit=%d", + accountID, wc, cause.source, cause.reason, cause.code, cause.bound, pending, age.Milliseconds(), idle.Milliseconds(), remaining, connectionLimit(capacityChat), connectionLimit(capacityBlank)) +} + +func (m *Manager) connectionCapacityError(accountID int64, kind connectionCapacityKind) error { + select { + case <-m.stopCleanup: + return context.Canceled + default: + } + limit := connectionLimit(kind) + log.Printf("[WS] 新建连接容量拒绝 account=%d budget=%s connections=%d limit=%d", accountID, kind.String(), m.budgetConnectionCount(connectionCapacityRequest{accountID: accountID, kind: kind}), limit) + return &proxy.UpstreamWSConnectionCapacityError{Limit: limit} +} diff --git a/proxy/wsrelay/connection_dial.go b/proxy/wsrelay/connection_dial.go new file mode 100644 index 000000000..c04bf01cc --- /dev/null +++ b/proxy/wsrelay/connection_dial.go @@ -0,0 +1,97 @@ +package wsrelay + +import ( + "context" + "net/http" + "strings" + + "github.com/codex2api/auth" + "github.com/codex2api/proxy" + "github.com/gorilla/websocket" +) + +type connectionDialRequest struct { + account *auth.Account + url string + sessionKey string + headers http.Header + proxyOverride string + capacity *connectionCapacitySlot +} + +func (m *Manager) createConnection(ctx context.Context, input connectionDialRequest) (wc *WsConnection, err error) { + defer func() { + if err != nil { + input.capacity.release() + } + }() + dialer, err := m.connectionDialer(input) + if err != nil { + return nil, err + } + session, poolKey := m.newDialSession(input) + conn, resp, err := dialer.DialContext(ctx, input.url, input.headers) + if err != nil { + m.sessions.CompareAndDelete(poolKey, session) + session.Close() + return nil, formatDialHandshakeError(err, resp) + } + logCompressionNegotiation(resp, input.account.ID()) + wc = m.wrapDialedConnection(dialedConnection{conn: conn, response: resp, session: session, poolKey: poolKey, request: input}) + input.capacity.attach(wc) + select { + case <-m.stopCleanup: + m.DiscardConnection(wc) + return nil, context.Canceled + default: + wc.installControlHandlers() + wc.StartReadPump() + return wc, nil + } +} + +func (m *Manager) connectionDialer(input connectionDialRequest) (*websocket.Dialer, error) { + dialer := *m.dialer + proxyURL := effectiveProxyURL(input.account, input.proxyOverride) + if dialProxy := proxy.CodexDialProxyURL(input.account, proxyURL); dialProxy != "" { + if err := configureWebsocketDialerProxy(&dialer, dialProxy); err != nil { + return nil, err + } + } + return &dialer, nil +} + +func (m *Manager) newDialSession(input connectionDialRequest) (*Session, string) { + poolKey := m.poolKey(input.account.ID(), input.url, input.sessionKey, effectiveProxyURL(input.account, input.proxyOverride)) + if value, ok := m.sessions.Load(poolKey); ok { + value.(*Session).Close() + } + session := NewSession(input.account.ID(), m) + if key := strings.TrimSpace(input.sessionKey); key != "" { + session.ID = key + } + m.sessions.Store(poolKey, session) + return session, poolKey +} + +type dialedConnection struct { + conn *websocket.Conn + response *http.Response + session *Session + poolKey string + request connectionDialRequest +} + +func (m *Manager) wrapDialedConnection(input dialedConnection) *WsConnection { + wc := NewWsConnection(input.conn, input.session, input.request.url) + wc.account = input.request.account + wc.PoolKey = input.poolKey + wc.upstreamUserAgent = strings.TrimSpace(input.request.headers.Get("User-Agent")) + wc.upstreamUserAgentKnown = true + wc.upstreamClientIdentity = websocketClientIdentity(input.request.headers) + wc.httpResp = input.response + wc.onDisconnected = m.getOnDisconnected() + wc.onReadFailure = m.DiscardConnection + input.session.SetConnected(true) + return wc +} diff --git a/proxy/wsrelay/continuation.go b/proxy/wsrelay/continuation.go new file mode 100644 index 000000000..f5d5a75d3 --- /dev/null +++ b/proxy/wsrelay/continuation.go @@ -0,0 +1,79 @@ +package wsrelay + +import ( + "context" + "log" + "time" + + "github.com/codex2api/proxy" +) + +// acquireContinuation 只等待原连接；绑定失效后由调用方恢复完整上下文。 +func (e *Executor) acquireContinuation(ctx context.Context, input websocketContinuation) (*WsConnection, *PendingRequest, string, error) { + parent := ctx + ctx, cancel := context.WithTimeout(ctx, busyAcquireMaxWait()) + defer cancel() + start, waited := time.Now(), false + for { + if err := ctx.Err(); err != nil { + return nil, nil, "", err + } + wake := e.manager.accountWaitSignal(input.accountID) + wc, _ := e.manager.lookupContinuationConn(input) + if wc == nil { + return nil, nil, "", &proxy.ResponsesContinuationLostError{Reason: "original_connection_unavailable"} + } + connection, pending, key := e.manager.acquirePreferredConnection(input) + if connection != nil { + if waited { + log.Printf("[WS] 续链等待后复用原连接 account=%d waited_ms=%d", input.accountID, time.Since(start).Milliseconds()) + } + return connection, pending, key, nil + } + waited = true + if _, err := waitForAccountChange(ctx, wake, AcquireMaxBackoff); err != nil { + if parent.Err() == nil && ctx.Err() == context.DeadlineExceeded { + return nil, nil, "", &proxy.ResponsesContinuationBusyError{} + } + return nil, nil, "", err + } + } +} + +func (r *WsResponse) bindCompletedResponse(responseID string) { + if responseID == "" || r.manager == nil || r.conn == nil || r.conn.session == nil { + return + } + r.manager.bindResponseConn(responseID, responseConnBinding{conn: r.conn, sessionKey: r.sessionID, + accountID: r.conn.session.AccountID, apiKey: r.apiKey, model: r.model}) +} + +func (m *Manager) lookupContinuationConn(input websocketContinuation) (*WsConnection, string) { + wc, key := m.lookupResponseConn(input.responseID, input.accountID, input.apiKey) + if wc == nil || (input.identity != "" && wc.upstreamClientIdentity != input.identity) { + return nil, "" + } + if input.url != "" && wc.PoolKey != m.poolKey(input.accountID, input.url, key, input.proxyURL) { + return nil, "" + } + m.respConnMu.Lock() + binding, exists := m.respConnBindings[input.responseID] + m.respConnMu.Unlock() + if !exists || binding.conn != wc || (binding.model != "" && input.model != "" && binding.model != input.model) { + return nil, "" + } + return wc, key +} + +// 发送失败后可取消地退避，避免瞬间重连风暴。 +func waitWebsocketSendRetry(ctx context.Context, retry int) error { + const sendRetryBackoff = 200 * time.Millisecond + timer := time.NewTimer(time.Duration(retry+1) * sendRetryBackoff) + defer timer.Stop() + select { + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: + return nil + } +} diff --git a/proxy/wsrelay/continuation_test.go b/proxy/wsrelay/continuation_test.go new file mode 100644 index 000000000..df4cee2ca --- /dev/null +++ b/proxy/wsrelay/continuation_test.go @@ -0,0 +1,198 @@ +package wsrelay + +import ( + "context" + "errors" + "io" + "net/http" + "testing" + "time" + + "github.com/codex2api/auth" + "github.com/codex2api/proxy" +) + +func continuationFixture(t *testing.T) (*Executor, *WsConnection) { + t.Helper() + manager := NewManager() + t.Cleanup(manager.Stop) + manager.probeFunc = func(*WsConnection) bool { return true } + wc := newBoundTestConn(t, manager, 7, "continuation") + wc.session.ID = "continuation" + wc.upstreamClientIdentity = "client" + manager.BindResponseConn("previous", wc, wc.session.ID, 7, "key") + return NewExecutorWithManager(manager), wc +} + +func TestContinuationWaitsForOriginalConnection(t *testing.T) { + executor, original := continuationFixture(t) + busy, err := executor.manager.addPendingAndBeginReadLease(original, "busy", capacityChat) + if err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + done := make(chan *WsConnection, 1) + go func() { + wc, pending, _, err := executor.acquireContinuation(ctx, websocketContinuation{responseID: "previous", accountID: 7, apiKey: "key", identity: "client"}) + if err == nil { + original.cancelUnsentReadLease(pending.RequestID) + original.session.RemovePendingRequest(pending.RequestID) + } + done <- wc + }() + select { + case <-done: + t.Fatal("busy original connection fell through") + case <-time.After(20 * time.Millisecond): + } + original.cancelUnsentReadLease(busy.RequestID) + original.session.RemovePendingRequest(busy.RequestID) + select { + case got := <-done: + if original != got { + t.Fatal("continuation did not reuse original connection") + } + case <-ctx.Done(): + t.Fatal("release did not wake continuation") + } +} + +func TestContinuationCancellationAndLostBinding(t *testing.T) { + executor, original := continuationFixture(t) + busy, _ := executor.manager.addPendingAndBeginReadLease(original, "busy", capacityChat) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + _, _, _, err := executor.acquireContinuation(ctx, websocketContinuation{responseID: "previous", accountID: 7, apiKey: "key", identity: "client"}) + if !errors.Is(err, context.Canceled) || original.session.PendingCount() != 1 { + t.Fatalf("cancellation leaked a reservation: %v", err) + } + original.cancelUnsentReadLease(busy.RequestID) + original.session.RemovePendingRequest(busy.RequestID) + executor.manager.DiscardConnection(original) + _, _, _, err = executor.acquireContinuation(context.Background(), websocketContinuation{responseID: "previous", accountID: 7, apiKey: "key", identity: "client"}) + var lost *proxy.ResponsesContinuationLostError + if !errors.As(err, &lost) { + t.Fatalf("lost binding returned %v", err) + } +} + +func TestContinuationSendFailureDoesNotRedialWithOldID(t *testing.T) { + executor, original := continuationFixture(t) + account := &auth.Account{DBID: 7, AccessToken: "fixture-token"} + headers := executor.prepareWebsocketHeaders(context.Background(), account.AccessToken, account, "", "session", "key", nil, http.Header{}, nil, "") + original.upstreamClientIdentity = websocketClientIdentity(headers) + wsURL, err := buildWebsocketURL(proxy.CodexBaseURL + CodexWsEndpoint) + if err != nil { + t.Fatal(err) + } + executor.manager.connections.Delete(original.PoolKey) + executor.manager.sessions.Delete(original.PoolKey) + original.PoolKey = executor.manager.poolKey(account.ID(), wsURL, original.session.ID, "") + executor.manager.connections.Store(original.PoolKey, original) + executor.manager.sessions.Store(original.PoolKey, original.session) + original.conn = newClosedTestWebsocketConn(t) + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + _, err = executor.ExecuteRequestViaWebsocket(ctx, account, []byte(`{"model":"gpt-5.5","previous_response_id":"previous","input":[]}`), "session", "", "key", nil, http.Header{}, "") + var lost *proxy.ResponsesContinuationLostError + if !errors.As(err, &lost) || lost.Reason != "original_connection_send_failed" { + t.Fatalf("send retried stale ID: %v", err) + } + if executor.manager.ConnectionCount() != 0 || original.session.PendingCount() != 0 { + t.Fatal("failed original retained connection or reservation") + } +} + +func TestContinuationBindingPrecedesTerminalDelivery(t *testing.T) { + executor, original := continuationFixture(t) + response := &WsResponse{conn: original, manager: executor.manager, sessionID: original.session.ID, apiKey: "key"} + err := response.handleMessage([]byte(`{"type":"response.completed","response":{"id":"latest"}}`), func([]byte) bool { + got, _ := executor.manager.lookupResponseConn("latest", 7, "key") + if original != got { + t.Fatal("terminal exposed before binding") + } + if err := response.Close(); err != nil || !original.IsConnected() { + t.Fatal("terminal-triggered cancellation discarded completed connection") + } + return false + }) + if !errors.Is(err, io.EOF) { + t.Fatal(err) + } + if got, _ := executor.manager.lookupResponseConn("latest", 7, "key"); original != got { + t.Fatal("consumed terminal lost reusable connection") + } +} + +func TestContinuationBusyTimeoutPreservesOriginal(t *testing.T) { + previous := proxy.CurrentRuntimeSettings() + proxy.UpdateRuntimeSettings(func(settings proxy.RuntimeSettings) proxy.RuntimeSettings { + settings.CodexWSBusyMaxWaitSec = 1 + return settings + }) + t.Cleanup(func() { proxy.ApplyRuntimeSettings(previous) }) + executor, original := continuationFixture(t) + busy, err := executor.manager.addPendingAndBeginReadLease(original, "busy", capacityChat) + if err != nil { + t.Fatal(err) + } + defer original.session.RemovePendingRequest(busy.RequestID) + defer original.cancelUnsentReadLease(busy.RequestID) + _, _, _, err = executor.acquireContinuation(context.Background(), websocketContinuation{responseID: "previous", accountID: 7, apiKey: "key", identity: "client"}) + var full *proxy.ResponsesContinuationBusyError + if !errors.As(err, &full) || !original.IsConnected() || original.session.PendingCount() != 1 { + t.Fatalf("busy continuation changed original or misclassified timeout: %v", err) + } +} + +func TestContinuationChecksModelAndEgress(t *testing.T) { + executor, original := continuationFixture(t) + executor.manager.bindResponseConn("previous", responseConnBinding{conn: original, sessionKey: original.session.ID, + accountID: 7, apiKey: "key", model: "model-a"}) + base := websocketContinuation{responseID: "previous", accountID: 7, apiKey: "key", identity: "client", model: "model-a", url: original.URL} + for _, changed := range []websocketContinuation{ + {responseID: "previous", accountID: 7, apiKey: "key", identity: "client", model: "model-b"}, + {responseID: "previous", accountID: 7, apiKey: "key", identity: "client", url: "wss://other.test/responses"}, + {responseID: "previous", accountID: 7, apiKey: "key", identity: "client", url: original.URL, proxyURL: "http://other.test:8080"}, + } { + if got, _ := executor.manager.lookupContinuationConn(changed); got != nil { + t.Fatal("changed model or egress reused original") + } + } + if got, _ := executor.manager.lookupContinuationConn(base); original != got { + t.Fatal("matching model and egress lost original") + } +} + +func TestIncompleteResponseEndsReadLeaseAndBinds(t *testing.T) { + executor, original := continuationFixture(t) + response := &WsResponse{conn: original, manager: executor.manager, sessionID: original.session.ID, apiKey: "key"} + payload := []byte(`{"type":"response.incomplete","response":{"id":"partial"}}`) + if !isReadLeaseTerminal(payload) || response.handleMessage(payload, func([]byte) bool { return true }) != io.EOF { + t.Fatal("incomplete response retained active stream") + } + if got, _ := executor.manager.lookupResponseConn("partial", 7, "key"); original != got { + t.Fatal("incomplete response lost continuation binding") + } +} + +func TestIdleChatRetentionEvictsOldestBoundAndProtectsActive(t *testing.T) { + executor, bound := continuationFixture(t) + manager := executor.manager + active := newBoundTestConn(t, manager, 7, "active") + active.session.AddPendingRequest("active") + active.lastUsed.Store(time.Now().Add(-2 * connectionIdleTimeout()).UnixNano()) + bound.lastUsed.Store(time.Now().Add(-2 * time.Minute).UnixNano()) + unbound := newBoundTestConn(t, manager, 7, "unbound") + unbound.lastUsed.Store(time.Now().Add(-time.Minute).UnixNano()) + if !manager.ensureAccountConnectionCapacity(connectionCapacityRequest{accountID: 7, limit: 1}) { + t.Fatal("idle retention should not reject new inference") + } + if !active.IsConnected() || bound.IsConnected() || !unbound.IsConnected() { + t.Fatal("LRU should evict the oldest bound chat and preserve active requests") + } + if got, _ := manager.lookupResponseConn("previous", 7, "key"); got == bound { + t.Fatal("evicted connection retained its continuation binding") + } +} diff --git a/proxy/wsrelay/executor.go b/proxy/wsrelay/executor.go index c81ee7f64..f153707c3 100644 --- a/proxy/wsrelay/executor.go +++ b/proxy/wsrelay/executor.go @@ -183,13 +183,18 @@ func (e *Executor) ExecuteRequestViaWebsocket( // 同线程上的后台副请求（request_kind=memory、guardian 子代理）另成一道， // 不与用户在飞轮次同键排队；Desktop 走 HTTP 时元数据只在请求体里。 poolSessionID := proxy.ResolveCodexWebsocketTransportSessionKeyWithBody(sessionID, ginHeaders, wsBody) + ctx = withConnectionCapacityKind(ctx, ginHeaders, wsBody) var wc *WsConnection var pr *PendingRequest var err2 error acquireStart := time.Now() if prevRespID := strings.TrimSpace(gjson.GetBytes(wsBody, "previous_response_id").String()); prevRespID != "" { - if pwc, ppr, slotKey := e.acquireClientContinuation(websocketContinuation{responseID: prevRespID, accountID: account.ID(), apiKey: apiKey, identity: websocketClientIdentity(headers)}); pwc != nil { - wc, pr, poolSessionID = pwc, ppr, slotKey + wc, pr, poolSessionID, err2 = e.acquireContinuation(ctx, websocketContinuation{ + responseID: prevRespID, accountID: account.ID(), apiKey: apiKey, identity: websocketClientIdentity(headers), + model: gjson.GetBytes(wsBody, "model").String(), url: wsURL, proxyURL: effectiveProxyURL(account, proxyOverride), + }) + if err2 != nil { + return nil, err2 } } baseKey := strings.TrimSpace(poolRouteKey) @@ -236,18 +241,21 @@ func (e *Executor) ExecuteRequestViaWebsocket( e.manager.DiscardConnection(wc) } wc.session.RemovePendingRequest(pr.RequestID) + e.manager.releaseUnusedChatCapacity(wc) return nil, err } sendErr := e.sendRequest(wc, wsBody, pr.RequestID) + if sendErr != nil && gjson.GetBytes(wsBody, "previous_response_id").String() != "" { + wc.session.RemovePendingRequest(pr.RequestID) + e.manager.DiscardConnection(wc) + return nil, &proxy.ResponsesContinuationLostError{Reason: "original_connection_send_failed"} + } for retries := 0; shouldRetryWebsocketSendError(sendErr) && retries < 2; retries++ { wc.session.RemovePendingRequest(pr.RequestID) e.manager.DiscardConnection(wc) - // 短暂退避，避免瞬间重连风暴 - select { - case <-ctx.Done(): - return nil, ctx.Err() - case <-time.After(time.Duration(retries+1) * 200 * time.Millisecond): + if err := waitWebsocketSendRetry(ctx, retries); err != nil { + return nil, err } reacquireStart := time.Now() @@ -276,6 +284,7 @@ func (e *Executor) ExecuteRequestViaWebsocket( sessionID: poolSessionID, manager: e.manager, apiKey: apiKey, + model: gjson.GetBytes(wsBody, "model").String(), readErrChan: make(chan error, 1), }, nil } @@ -451,6 +460,7 @@ func (e *Executor) sendRequest(wc *WsConnection, body []byte, requestID string) // WsResponse WebSocket 响应包装器 type WsResponse struct { + model string conn *WsConnection pendingReq *PendingRequest sessionID string @@ -531,7 +541,9 @@ func (r *WsResponse) handleMessage(payload []byte, callback func(data []byte) bo r.markConnBroken() } // 把错误内容作为 SSE 数据写给下游，让客户端看到完整错误 JSON。 - callback(errEvent) + if !callback(errEvent) { + r.markConnBroken() + } // 错误即终止：结束流(等价于 response.failed)。 return io.EOF } @@ -539,29 +551,29 @@ func (r *WsResponse) handleMessage(payload []byte, callback func(data []byte) bo // 标准化完成事件类型 payload = normalizeCompletionEvent(payload) + // 终态先记录绑定与完成状态，再交给下游。 + eventType := gjson.GetBytes(payload, "type").String() + terminal := isReadLeaseTerminal(payload) + if eventType == "response.completed" || eventType == "response.incomplete" { + r.bindCompletedResponse(gjson.GetBytes(payload, "response.id").String()) + } + if terminal { + // 下游读到终态即可取消请求；提前标记，避免取消回调误销毁已完成连接。 + r.markStreamCompleted() + } // 调用回调 if !callback(payload) { // 下游写入失败(broken pipe / 客户端断开)：响应流在非终止边界被截断， // 上游仍会在这条连接上继续推送本响应的剩余帧。连接必须销毁， // 归还池中复用会把残留帧串给下一个请求(issue #308)。 - r.markConnBroken() + if !terminal { + r.markConnBroken() + } return io.EOF } // 检查是否是终止事件 - eventType := gjson.GetBytes(payload, "type").String() - if eventType == "response.completed" || eventType == "response.failed" { - // 续链亲和：记录本响应由哪条连接产出，后续带 previous_response_id 的 - // 请求可回到原连接（上游无服务端存储时上下文只存活在连接内）。 - if eventType == "response.completed" && r.manager != nil && r.conn != nil { - if respID := gjson.GetBytes(payload, "response.id").String(); respID != "" { - accountID := int64(0) - if r.conn.session != nil { - accountID = r.conn.session.AccountID - } - r.manager.BindResponseConn(respID, r.conn, r.sessionID, accountID, r.apiKey) - } - } + if terminal { return io.EOF } @@ -650,6 +662,9 @@ func (r *WsResponse) markConnBroken() { r.mu.Lock() r.connBroken = true r.mu.Unlock() + if r.manager != nil && r.conn != nil { + r.manager.DiscardConnection(r.conn) + } } // markStreamCompleted 标记读流已消费到明确的终止边界（幂等，受 mu 保护）。 @@ -685,7 +700,11 @@ func (r *WsResponse) Close() error { if !r.connBroken && r.streamCompleted && !r.shouldDiscardOneShotConn() { r.manager.ReleaseConnection(r.conn) } else { - r.manager.DiscardConnection(r.conn) + reason := closeUnconsumed + if !r.connBroken && r.streamCompleted { + reason = closeOneShotComplete + } + r.manager.discardConnectionFor(r.conn, reason) } } @@ -696,14 +715,11 @@ func (r *WsResponse) Close() error { // 这类连接的池键每请求唯一，归还池后不可能再被按键复用，只会占用账号连接名额 // 直到空闲超时；唯一的保留价值是 response_id 续链亲和（上游无服务端存储时， // previous_response_id 的上下文只存活在产出响应的那条连接里），因此有存活绑定时 -// 仍归还池。CODEX_WS_STATELESS_ONESHOT 模式显式承诺用完即毁，无条件销毁。 +// 仍归还池；一次性模式只禁止通用槽位复用，不应破坏已经产生的续链。 func (r *WsResponse) shouldDiscardOneShotConn() bool { if r.conn == nil || r.conn.session == nil || !proxy.IsStatelessWebsocketSessionID(r.conn.session.ID) { return false } - if statelessOneShotEnabled() { - return true - } return r.manager == nil || !r.manager.hasLiveResponseBinding(r.conn) } diff --git a/proxy/wsrelay/issue436_pending_protect_test.go b/proxy/wsrelay/issue436_pending_protect_test.go index 4ebfd60a7..bfc7667ae 100644 --- a/proxy/wsrelay/issue436_pending_protect_test.go +++ b/proxy/wsrelay/issue436_pending_protect_test.go @@ -165,6 +165,7 @@ func TestEvictExpiredSkipsConnectionsWithPendingRequests(t *testing.T) { wc := NewWsConnection(nil, session, "ws://example") wc.PoolKey = "key-busy" wc.SetState(StateConnected) + wc.Touch() // 伪造超龄空闲时间戳：lastUsed 早于 IdleTimeout，LastActiveAt 同步做旧 wc.lastUsed.Store(time.Now().Add(-2 * IdleTimeout).UnixNano()) session.mu.Lock() diff --git a/proxy/wsrelay/keepalive.go b/proxy/wsrelay/keepalive.go index 3b870f66b..773f19b70 100644 --- a/proxy/wsrelay/keepalive.go +++ b/proxy/wsrelay/keepalive.go @@ -42,9 +42,8 @@ func (m *Manager) PingIdleConnections() (pinged int, failed int) { if !wc.IsConnected() { return true } - // 到龄连接不续命：Pong 会刷新 lastUsed 让连接永不空闲过期，最终撞上游 - // 60 分钟寿命上限；跳过 Ping，交给 evictExpired 按到龄轮转销毁。 - if wc.IsOverAge() { + // 业务空闲过期或到龄的连接不再保活，交给 evictExpired 关闭并归还槽位。 + if wc.IsExpired() || wc.IsOverAge() { return true } // 只保活空闲连接：有 pending 请求时跳过（该连接本就活跃，无需额外 Ping）。 diff --git a/proxy/wsrelay/manager.go b/proxy/wsrelay/manager.go index ef9191c10..25532355e 100644 --- a/proxy/wsrelay/manager.go +++ b/proxy/wsrelay/manager.go @@ -6,14 +6,12 @@ import ( "log" "net" "net/http" - "sort" "strings" "sync" "sync/atomic" "time" "github.com/codex2api/auth" - "github.com/codex2api/proxy" "github.com/codex2api/security" "github.com/gorilla/websocket" ) @@ -41,8 +39,8 @@ type WsConnection struct { upstreamUserAgentKnown bool upstreamClientIdentity string - // 创建/复用该连接的账号。仅用于读取当前动态并发上限，让 response_id - // 续链复用路径也能在账号上限下调后收敛空闲连接数。 + // 创建/复用该连接的账号。请求并发由调度层独立管理， + // response_id 复用按共享连接槽位预算收敛无绑定的空闲连接。 account *auth.Account // 会话 @@ -57,7 +55,10 @@ type WsConnection struct { // 连接状态 state atomic.Int32 - // 最后使用时间 + capacitySlot atomic.Pointer[connectionCapacitySlot] + closeCause atomic.Pointer[connectionCloseCause] + + // 最后业务使用时间，心跳和空闲元数据不续期。 lastUsed atomic.Int64 // 最近入站活动时间（数据帧/对端 Ping/Pong 回执，UnixNano）。仅供 probe @@ -150,9 +151,13 @@ func NewWsConnection(conn *websocket.Conn, session *Session, wsURL string) *WsCo return wc } -// Touch 更新最后使用时间 +// Touch 同步更新连接和会话的最后业务使用时间。 func (wc *WsConnection) Touch() { wc.lastUsed.Store(time.Now().UnixNano()) + if wc.session != nil { + slot := wc.capacitySlot.Load() + wc.session.touchActivity(slot == nil || connectionCapacityKind(slot.kind.Load()) == capacityChat) + } } // touchInbound 记录一次入站活动（数据帧/对端 Ping/我方 Ping 的 Pong 回执）。 @@ -169,10 +174,13 @@ func (wc *WsConnection) recentInboundWithin(window time.Duration) bool { return time.Since(time.Unix(0, ts)) <= window } -// IsExpired 检查连接是否过期 +// IsExpired 只对曾用于聊天的连接执行业务空闲回收；空白槽位正常保活。 func (wc *WsConnection) IsExpired() bool { + if !weakNetworkModeEnabled() && (wc.session == nil || !wc.session.hasUserContext()) { + return false + } lastUsed := time.Unix(0, wc.lastUsed.Load()) - return time.Since(lastUsed) > connectionIdleTimeout() + return time.Since(lastUsed) >= connectionIdleTimeout() } // IsOverAge 检查连接是否超过当前模式的最大寿命。到龄连接不能再接新请求： @@ -195,11 +203,14 @@ func (wc *WsConnection) Close() error { return nil } wc.closeOnce.Do(func() { + wc.recordDirectClose() wc.state.Store(int32(StateClosing)) if wc.conn != nil { wc.closeErr = wc.conn.Close() } wc.state.Store(int32(StateDisconnected)) + wc.capacitySlot.Load().release() + wc.logClose() }) if wc.onDisconnected != nil && wc.session != nil && wc.disconnectNotified.CompareAndSwap(false, true) { wc.onDisconnected(wc.session.AccountID) @@ -229,6 +240,9 @@ func (wc *WsConnection) WriteMessage(messageType int, data []byte) error { defer wc.conn.SetWriteDeadline(time.Time{}) writeErr := wc.conn.WriteMessage(messageType, data) + if writeErr == nil && (messageType == websocket.TextMessage || messageType == websocket.BinaryMessage) { + wc.Touch() + } if tracksLease { return wc.completeReadLeaseWrite(leaseID, writeErr) } @@ -269,9 +283,9 @@ type Manager struct { keyLocks refCountedLockRegistry[string] // 账号级串行化连接获取，确保跨 session/pool key 创建连接时仍能严格执行 // 每账号连接上限，避免大量短会话各留一条空闲连接。 - accountLocks refCountedLockRegistry[int64] - capacityMu sync.Mutex - pendingCreates map[int64]int + accountLocks refCountedLockRegistry[int64] + capacityMu sync.Mutex + connectionSlots map[int64]map[*connectionCapacitySlot]*WsConnection // 账号级等待唤醒：busy/容量等待者在同账号连接释放、销毁、新建或在途请求结束时 // 立即被唤醒重新选连，不再空等整个退避间隔（最多 AcquireMaxBackoff）。 @@ -306,12 +320,13 @@ type responseConnBinding struct { sessionKey string accountID int64 apiKey string + model string expiresAt time.Time } const ( // responseConnBindingTTL 续链绑定的存活时间。上游空闲连接本身在 IdleTimeout - // (5min) 后被清理，绑定活得再久也无意义，与其对齐。 + // (30min) 后被清理，绑定活得再久也无意义，与其对齐。 responseConnBindingTTL = IdleTimeout // responseConnBindingMaxEntries 绑定表上限，防止内存膨胀。 responseConnBindingMaxEntries = 4096 @@ -370,9 +385,7 @@ func (m *Manager) evictExpired() { return true } if wc.IsExpired() || !wc.IsConnected() || isRotatableOverAge(wc) { - m.connections.CompareAndDelete(key, wc) - m.removeResponseConnBindings(wc) - wc.Close() + m.discardConnectionFor(wc, connectionExpiryReason(wc)) if wc.session != nil { m.notifyAccountWaiters(wc.session.AccountID) } @@ -392,6 +405,7 @@ func (m *Manager) evictExpired() { return true }) m.evictResponseConnBindings(time.Now()) + m.trimAllIdleChatConnections() } // Stop 停止管理器 @@ -404,11 +418,10 @@ func (m *Manager) Stop() { // closeAll 关闭所有连接 func (m *Manager) closeAll() { + m.closeCapacityConnections() m.connections.Range(func(key, value any) bool { wc := value.(*WsConnection) - m.connections.CompareAndDelete(key, wc) - m.removeResponseConnBindings(wc) - wc.Close() + m.discardConnectionFor(wc, closeManagerShutdown) return true }) m.respConnMu.Lock() @@ -494,178 +507,6 @@ func (m *Manager) accountLock(accountID int64) (*sync.Mutex, func()) { return m.accountLocks.retain(accountID) } -func accountConnectionLimit(account *auth.Account) int { - if account != nil { - if limit := account.GetDynamicConcurrencyLimit(); limit > 0 { - return int(limit) - } - } - // 尚未完成调度快照初始化的账号保留原有槽位上限，生产请求进入账号池后 - // DynamicConcurrencyLimit 会始终为正数。 - return statelessConnectionSlots() -} - -type idleAccountConnection struct { - wc *WsConnection - lastUsed int64 -} - -// isOneShotPoolConn 判断连接是否挂在每请求唯一的 stateless 池键下（8 槽全忙时的 -// fallback 一次性连接，或直连 stateless key 的路径）。这类连接的池键不可能被后续 -// 请求按键命中，仅 response_id 续链亲和可能定向取回。 -func isOneShotPoolConn(wc *WsConnection) bool { - return wc != nil && wc.session != nil && proxy.IsStatelessWebsocketSessionID(wc.session.ID) -} - -// sortIdleForEviction 决定容量裁剪的逐出顺序：一次性池键连接优先，组内按 LRU。 -// 纯 LRU 会出现倒挂——刚用完的一次性连接比几分钟前用过的可复用槽连接"更新"， -// 结果保住永远不会被按键复用的僵尸、逐出真正的热槽，拖垮复用率。 -func sortIdleForEviction(idle []idleAccountConnection) { - sort.Slice(idle, func(i, j int) bool { - oi, oj := isOneShotPoolConn(idle[i].wc), isOneShotPoolConn(idle[j].wc) - if oi != oj { - return oi - } - return idle[i].lastUsed < idle[j].lastUsed - }) -} - -// hasLiveResponseBinding 是否存在未过期的 response_id 续链绑定指向该连接。 -func (m *Manager) hasLiveResponseBinding(wc *WsConnection) bool { - if m == nil || wc == nil { - return false - } - now := time.Now() - m.respConnMu.Lock() - defer m.respConnMu.Unlock() - for _, b := range m.respConnBindings { - if b.conn == wc && now.Before(b.expiresAt) { - return true - } - } - return false -} - -// ensureAccountConnectionCapacity 为即将创建的新连接腾出一个账号级槽位。 -// 只淘汰没有在途请求的最久未使用连接，绝不打断活跃响应。调用方必须持有该账号的 -// accountLock，因此不同 session 同时建连也不会越过动态并发上限。 -func (m *Manager) ensureAccountConnectionCapacity(accountID int64, limit int, protectedKey string, pendingCreates int) bool { - if limit < 1 { - limit = 1 - } - count := 0 - stale := make([]*WsConnection, 0) - idle := make([]idleAccountConnection, 0) - m.connections.Range(func(_, value any) bool { - wc, ok := value.(*WsConnection) - if !ok || wc == nil || wc.session == nil || wc.session.AccountID != accountID { - return true - } - if !wc.IsConnected() || wc.IsExpired() || isRotatableOverAge(wc) { - stale = append(stale, wc) - return true - } - count++ - if wc.PoolKey != protectedKey && wc.session.PendingCount() == 0 { - idle = append(idle, idleAccountConnection{wc: wc, lastUsed: wc.lastUsed.Load()}) - } - return true - }) - for _, wc := range stale { - m.DiscardConnection(wc) - } - if count+pendingCreates < limit { - return true - } - sortIdleForEviction(idle) - for _, candidate := range idle { - if count+pendingCreates < limit { - break - } - wc := candidate.wc - if wc == nil || wc.session == nil || wc.session.PendingCount() != 0 || !wc.IsConnected() { - continue - } - if current, ok := m.connections.Load(wc.PoolKey); !ok || current != wc { - continue - } - m.DiscardConnection(wc) - count-- - } - return count+pendingCreates < limit -} - -// trimIdleAccountConnections 在复用现有连接时把账号连接数收敛到当前动态并发上限。 -// 当前要复用的连接始终受保护；其它仍有在途请求的连接也不会被打断。若活跃连接数 -// 本身已经超过上限，本次只清理能够安全回收的空闲连接，后续 acquire 再继续收敛。 -// 调用方必须持有该账号的 accountLock。 -func (m *Manager) trimIdleAccountConnections(accountID int64, limit int, protected *WsConnection) { - if limit < 1 { - limit = 1 - } - count := 0 - idle := make([]idleAccountConnection, 0) - m.connections.Range(func(_, value any) bool { - wc, ok := value.(*WsConnection) - if !ok || wc == nil || wc.session == nil || wc.session.AccountID != accountID { - return true - } - count++ - if wc != protected && wc.session.PendingCount() == 0 { - idle = append(idle, idleAccountConnection{wc: wc, lastUsed: wc.lastUsed.Load()}) - } - return true - }) - if count <= limit { - return - } - - sortIdleForEviction(idle) - for _, candidate := range idle { - if count <= limit { - break - } - wc := candidate.wc - if wc == nil || wc == protected || wc.session == nil || wc.session.PendingCount() != 0 { - continue - } - if current, ok := m.connections.Load(wc.PoolKey); !ok || current != wc { - continue - } - m.DiscardConnection(wc) - count-- - } -} - -func (m *Manager) reserveAccountConnectionCapacity(accountID int64, limit int, protectedKey string) bool { - m.capacityMu.Lock() - pending := m.pendingCreates[accountID] - m.capacityMu.Unlock() - if !m.ensureAccountConnectionCapacity(accountID, limit, protectedKey, pending) { - return false - } - m.capacityMu.Lock() - if m.pendingCreates == nil { - m.pendingCreates = make(map[int64]int) - } - m.pendingCreates[accountID]++ - m.capacityMu.Unlock() - return true -} - -func (m *Manager) releaseAccountConnectionCapacity(accountID int64) { - m.capacityMu.Lock() - if pending := m.pendingCreates[accountID]; pending <= 1 { - delete(m.pendingCreates, accountID) - } else { - m.pendingCreates[accountID] = pending - 1 - } - m.capacityMu.Unlock() - m.notifyAccountWaiters(accountID) -} - -// accountWaitSignal 返回账号当前的等待唤醒 channel：下一次 notifyAccountWaiters 会关闭它。 -// 等待者必须在检查池状态之前取到 channel，才不会漏掉检查与等待之间发生的变化。 func (m *Manager) accountWaitSignal(accountID int64) <-chan struct{} { m.waitNotifyMu.Lock() defer m.waitNotifyMu.Unlock() @@ -721,6 +562,7 @@ func (m *Manager) AcquireConnection( headers http.Header, proxyOverride string, ) (*WsConnection, *PendingRequest, error) { + kind := requestConnectionCapacityKind(ctx) key := m.poolKey(account.ID(), wsURL, sessionKey, effectiveProxyURL(account, proxyOverride)) lock, releaseKeyLock := m.keyLock(key) defer releaseKeyLock() @@ -749,26 +591,29 @@ func (m *Manager) AcquireConnection( lock.Unlock() continue } - pr, leaseErr := m.addPendingAndBeginReadLease(wc, sessionKey) + pr, leaseErr := m.addPendingAndBeginReadLease(wc, sessionKey, kind) if leaseErr == nil { wc.account = account wc.Touch() - m.trimIdleAccountConnections(account.ID(), accountConnectionLimit(account), wc) + m.trimIdleAccountConnections(account.ID(), connectionLimit(connectionBudgetKind(wc)), wc) accountLock.Unlock() lock.Unlock() return wc, pr, nil } - m.DiscardConnection(wc) + m.discardAcquisitionFailure(wc, leaseErr) accountLock.Unlock() lock.Unlock() + if isConnectionCapacityError(leaseErr) { + return nil, nil, leaseErr + } continue } // 探活失败，清理死连接 - m.DiscardConnection(wc) + m.discardConnectionFor(wc, closeProbeFailed) lock.Unlock() continue } - if wc.IsConnected() && !wc.IsExpired() && wc.session != nil && wc.session.PendingCount() > 0 && !isRotatableOverAge(wc) { + if wc.session != nil && wc.session.PendingCount() > 0 { lock.Unlock() // 连接被同 session 的前一个请求占用：指数退避轮询等待其空闲， // 累计等待超过上限则返回错误，避免无界阻塞与固定间隔空转抢锁。 @@ -804,33 +649,18 @@ func (m *Manager) AcquireConnection( m.DiscardConnection(wc) } accountLock.Lock() - if !m.reserveAccountConnectionCapacity(account.ID(), accountConnectionLimit(account), key) { + capacity := m.reserveAccountConnectionCapacity(connectionCapacityRequest{accountID: account.ID(), limit: connectionLimit(kind), protectedKey: key, kind: kind}) + if capacity == nil { accountLock.Unlock() lock.Unlock() - if maxWait := busyAcquireMaxWait(); waited >= maxWait { - return nil, nil, fmt.Errorf("acquire websocket connection timed out after %s waiting for account connection capacity", maxWait) - } - // 同账号连接销毁/拨号占位归还时立即唤醒重试容量预留，退避只是兜底。 - slept, waitErr := waitForAccountChange(ctx, wake, wait) - if waitErr != nil { - return nil, nil, waitErr - } - waited += slept - if wait < AcquireMaxBackoff { - wait *= 2 - if wait > AcquireMaxBackoff { - wait = AcquireMaxBackoff - } - } - continue + return nil, nil, m.connectionCapacityError(account.ID(), kind) } - // 容量已预留，拨号期间不再持有账号锁；其他 session 可以复用已有连接， - // 但会把本次 pending create 计入上限，避免并发握手越界。 + // 实体连接已登记，拨号期间释放账号锁，其他 session 可继续取连。 + // 空白连接的拨号登记占用预算；聊天推理不占空闲保留额度。 accountLock.Unlock() - wc, err := m.createConnection(ctx, account, wsURL, sessionKey, headers, proxyOverride) + wc, err := m.createConnection(ctx, connectionDialRequest{account: account, url: wsURL, sessionKey: sessionKey, headers: headers, proxyOverride: proxyOverride, capacity: capacity}) if err != nil { - m.releaseAccountConnectionCapacity(account.ID()) lock.Unlock() return nil, nil, err } @@ -840,14 +670,13 @@ func (m *Manager) AcquireConnection( if m.afterConnectionStored != nil { m.afterConnectionStored(wc) } - pr, leaseErr := m.addPendingAndBeginReadLease(wc, sessionKey) + pr, leaseErr := m.addPendingAndBeginReadLease(wc, sessionKey, kind) if leaseErr == nil { if earlyErr := wc.waitForEarlyReadFailure(ctx, newConnectionReadFailureGrace); earlyErr != nil { wc.session.RemovePendingRequest(pr.RequestID) leaseErr = earlyErr } } - m.releaseAccountConnectionCapacity(account.ID()) if leaseErr != nil { m.DiscardConnection(wc) lock.Unlock() @@ -889,7 +718,8 @@ func (m *Manager) tryAcquireBusyOverflow( proxyOverride string, ) (*WsConnection, *PendingRequest, bool) { proxyURL := effectiveProxyURL(account, proxyOverride) - accountLimit := accountConnectionLimit(account) + kind := requestConnectionCapacityKind(ctx) + accountLimit := connectionLimit(kind) accountLock, releaseAccountLock := m.accountLock(account.ID()) defer releaseAccountLock() for i := 1; i <= BusyOverflowSlots; i++ { @@ -901,33 +731,37 @@ func (m *Manager) tryAcquireBusyOverflow( if v, ok := m.connections.Load(key); ok { wc := v.(*WsConnection) if canReuseConnection(wc) { + if m.hasLiveResponseBinding(wc) { + lock.Unlock() + continue + } if m.probe(wc) { accountLock.Lock() current, exists := m.connections.Load(key) - if exists && current == wc && canReuseConnection(wc) { - pr, leaseErr := m.addPendingAndBeginReadLease(wc, slotSession) + if exists && current == wc && canReuseConnection(wc) && !m.hasLiveResponseBinding(wc) { + pr, leaseErr := m.addPendingAndBeginReadLease(wc, slotSession, kind) if leaseErr == nil { wc.account = account wc.Touch() - m.trimIdleAccountConnections(account.ID(), accountLimit, wc) + m.trimIdleAccountConnections(account.ID(), connectionLimit(connectionBudgetKind(wc)), wc) accountLock.Unlock() lock.Unlock() return wc, pr, true } - m.DiscardConnection(wc) + m.discardAcquisitionFailure(wc, leaseErr) } accountLock.Unlock() lock.Unlock() continue } - m.DiscardConnection(wc) - } else if wc.IsConnected() && !wc.IsExpired() && wc.session != nil && wc.session.PendingCount() > 0 && !isRotatableOverAge(wc) { + m.discardConnectionFor(wc, closeProbeFailed) + } else if wc.session != nil && wc.session.PendingCount() > 0 { // 兄弟槽位也在忙：换下一个槽位 lock.Unlock() continue } else { // 死/到龄/过期连接：清掉腾出槽位，下方直接新建 - m.DiscardConnection(wc) + m.discardConnectionFor(wc, connectionExpiryReason(wc)) } } accountLock.Lock() @@ -936,16 +770,16 @@ func (m *Manager) tryAcquireBusyOverflow( lock.Unlock() continue } - if !m.reserveAccountConnectionCapacity(account.ID(), accountLimit, key) { + capacity := m.reserveAccountConnectionCapacity(connectionCapacityRequest{accountID: account.ID(), limit: accountLimit, protectedKey: key, kind: kind}) + if capacity == nil { // 账号连接容量已满：不为 overflow 挤占更多连接，放弃降级回到等待 accountLock.Unlock() lock.Unlock() return nil, nil, false } accountLock.Unlock() - wc, err := m.createConnection(ctx, account, wsURL, slotSession, headers, proxyOverride) + wc, err := m.createConnection(ctx, connectionDialRequest{account: account, url: wsURL, sessionKey: slotSession, headers: headers, proxyOverride: proxyOverride, capacity: capacity}) if err != nil { - m.releaseAccountConnectionCapacity(account.ID()) lock.Unlock() log.Printf("[WS] busy overflow 新建连接失败，回落等待原连接 (account=%d): %v", account.ID(), err) return nil, nil, false @@ -954,14 +788,13 @@ func (m *Manager) tryAcquireBusyOverflow( if m.afterConnectionStored != nil { m.afterConnectionStored(wc) } - pr, leaseErr := m.addPendingAndBeginReadLease(wc, slotSession) + pr, leaseErr := m.addPendingAndBeginReadLease(wc, slotSession, kind) if leaseErr == nil { if earlyErr := wc.waitForEarlyReadFailure(ctx, newConnectionReadFailureGrace); earlyErr != nil { wc.session.RemovePendingRequest(pr.RequestID) leaseErr = earlyErr } } - m.releaseAccountConnectionCapacity(account.ID()) if leaseErr != nil { m.DiscardConnection(wc) lock.Unlock() @@ -992,7 +825,7 @@ const newConnectionReadFailureGrace = 5 * time.Millisecond // AcquireReusableConnection 在固定槽位内复用或创建连接，返回实际使用的 session key。 // 第一遍只复用已存在且空闲的连接；第二遍在空槽位新建持久连接；槽位全忙时回退到 -// fallbackKey 的临时连接。所有路径仍受账号动态并发对应的连接总数上限约束。 +// fallbackKey 的临时连接。空白预热共用账号空白预算；聊天收尾后按空闲上限保留。 func (m *Manager) AcquireReusableConnection( ctx context.Context, account *auth.Account, @@ -1004,9 +837,10 @@ func (m *Manager) AcquireReusableConnection( proxyOverride string, ) (*WsConnection, *PendingRequest, string, error) { proxyURL := effectiveProxyURL(account, proxyOverride) - accountLimit := accountConnectionLimit(account) - if slots < 1 || slots > accountLimit { - slots = accountLimit + kind := requestConnectionCapacityKind(ctx) + accountLimit := connectionLimit(kind) + if slots < 1 || slots > statelessConnectionSlots() { + slots = statelessConnectionSlots() } accountLock, releaseAccountLock := m.accountLock(account.ID()) defer releaseAccountLock() @@ -1020,31 +854,35 @@ func (m *Manager) AcquireReusableConnection( if v, ok := m.connections.Load(key); ok { wc := v.(*WsConnection) if canReuseConnection(wc) { + if m.hasLiveResponseBinding(wc) { + lock.Unlock() + continue + } if m.probe(wc) { accountLock.Lock() current, exists := m.connections.Load(key) - if !exists || current != wc || !canReuseConnection(wc) { + if !exists || current != wc || !canReuseConnection(wc) || m.hasLiveResponseBinding(wc) { accountLock.Unlock() lock.Unlock() continue } - pr, leaseErr := m.addPendingAndBeginReadLease(wc, slotSession) + pr, leaseErr := m.addPendingAndBeginReadLease(wc, slotSession, kind) if leaseErr == nil { wc.account = account wc.Touch() - m.trimIdleAccountConnections(account.ID(), accountLimit, wc) + m.trimIdleAccountConnections(account.ID(), connectionLimit(connectionBudgetKind(wc)), wc) accountLock.Unlock() lock.Unlock() return wc, pr, slotSession, nil } - m.DiscardConnection(wc) + m.discardAcquisitionFailure(wc, leaseErr) accountLock.Unlock() lock.Unlock() continue } - m.DiscardConnection(wc) - } else if !wc.IsConnected() || wc.IsExpired() || isRotatableOverAge(wc) || wc.session == nil || wc.session.PendingCount() == 0 { - m.DiscardConnection(wc) + m.discardConnectionFor(wc, closeProbeFailed) + } else if wc.session == nil || wc.session.PendingCount() == 0 { + m.discardConnectionFor(wc, connectionExpiryReason(wc)) } } lock.Unlock() @@ -1066,15 +904,15 @@ func (m *Manager) AcquireReusableConnection( lock.Unlock() continue } - if !m.reserveAccountConnectionCapacity(account.ID(), accountLimit, key) { + capacity := m.reserveAccountConnectionCapacity(connectionCapacityRequest{accountID: account.ID(), limit: accountLimit, protectedKey: key, kind: kind}) + if capacity == nil { accountLock.Unlock() lock.Unlock() continue } accountLock.Unlock() - wc, err := m.createConnection(ctx, account, wsURL, slotSession, headers, proxyOverride) + wc, err := m.createConnection(ctx, connectionDialRequest{account: account, url: wsURL, sessionKey: slotSession, headers: headers, proxyOverride: proxyOverride, capacity: capacity}) if err != nil { - m.releaseAccountConnectionCapacity(account.ID()) lock.Unlock() return nil, nil, "", err } @@ -1082,14 +920,13 @@ func (m *Manager) AcquireReusableConnection( if m.afterConnectionStored != nil { m.afterConnectionStored(wc) } - pr, leaseErr := m.addPendingAndBeginReadLease(wc, slotSession) + pr, leaseErr := m.addPendingAndBeginReadLease(wc, slotSession, kind) if leaseErr == nil { if earlyErr := wc.waitForEarlyReadFailure(ctx, newConnectionReadFailureGrace); earlyErr != nil { wc.session.RemovePendingRequest(pr.RequestID) leaseErr = earlyErr } } - m.releaseAccountConnectionCapacity(account.ID()) if leaseErr != nil { m.DiscardConnection(wc) lock.Unlock() @@ -1113,10 +950,13 @@ func (m *Manager) AcquireReusableConnection( // atomic from an acquire caller's perspective. On failure it rolls the pending // request back; the caller discards the unusable connection while holding its // pool-key acquisition lock. -func (m *Manager) addPendingAndBeginReadLease(wc *WsConnection, sessionKey string) (*PendingRequest, error) { +func (m *Manager) addPendingAndBeginReadLease(wc *WsConnection, sessionKey string, kind connectionCapacityKind) (*PendingRequest, error) { if wc == nil || wc.session == nil { return nil, fmt.Errorf("begin websocket read lease: connection has no session") } + if err := m.ensureConnectionBudget(wc, kind); err != nil { + return nil, err + } pr := wc.session.AddPendingRequest(sessionKey) if err := wc.BeginReadLease(pr.RequestID); err != nil { wc.session.RemovePendingRequest(pr.RequestID) @@ -1206,77 +1046,15 @@ func logCompressionNegotiation(resp *http.Response, accountID int64) { } } -// createConnection 创建新 WebSocket 连接 -func (m *Manager) createConnection( - ctx context.Context, - account *auth.Account, - wsURL string, - sessionKey string, - headers http.Header, - proxyOverride string, -) (*WsConnection, error) { - // 浅拷贝共享 dialer，继承全部调优字段（NetDialContext/KeepAlive、读写缓冲、压缩等）， - // 仅按需覆盖 Proxy；避免逐字段重建时漏抄字段（曾导致 NetDialContext/KeepAlive 失效）。 - dialerCopy := *m.dialer - dialer := &dialerCopy - - // 拨号代理与 ExecuteRequestViaWebsocket 同一套判定:Resin 承担出站时不配代理 - // (传入的 wsURL 已是 Resin 反代地址);poolKey 仍按第 2 层代理分池,保持既有键不变。 - proxyURL := effectiveProxyURL(account, proxyOverride) - dialProxy := proxy.CodexDialProxyURL(account, proxyURL) - if dialProxy != "" { - if err := configureWebsocketDialerProxy(dialer, dialProxy); err != nil { - return nil, err - } - } - - // 创建会话（先关闭旧 session 避免泄漏） - poolKey := m.poolKey(account.ID(), wsURL, sessionKey, proxyURL) - if oldSessionVal, ok := m.sessions.Load(poolKey); ok { - oldSession := oldSessionVal.(*Session) - oldSession.Close() - } - session := NewSession(account.ID(), m) - if trimmed := strings.TrimSpace(sessionKey); trimmed != "" { - session.ID = trimmed - } - m.sessions.Store(poolKey, session) - - // 拨号连接 - conn, resp, err := dialer.DialContext(ctx, wsURL, headers) - if err != nil { - m.sessions.Delete(poolKey) - session.Close() - // bad handshake 时 resp 常非空：附带上游 HTTP 状态/ body，便于测试连接定位。 - return nil, formatDialHandshakeError(err, resp) - } - - logCompressionNegotiation(resp, account.ID()) - - // 创建连接包装 - wc := NewWsConnection(conn, session, wsURL) - wc.account = account - wc.PoolKey = poolKey - wc.upstreamUserAgent = strings.TrimSpace(headers.Get("User-Agent")) - wc.upstreamUserAgentKnown = true - wc.upstreamClientIdentity = websocketClientIdentity(headers) - wc.httpResp = resp - wc.onDisconnected = m.getOnDisconnected() - wc.onReadFailure = m.DiscardConnection - session.SetConnected(true) - - // 控制帧处理器必须在唯一永久 reader 启动前安装。 - wc.installControlHandlers() - wc.StartReadPump() - - return wc, nil -} - // ReleaseConnection 释放连接（归还池） func (m *Manager) ReleaseConnection(wc *WsConnection) { if wc == nil { return } + if wc.session != nil { + m.retainReleasedConnection(wc) + return + } wc.Touch() } @@ -1285,6 +1063,7 @@ func (m *Manager) RemoveConnection(accountID int64, wsURL string, sessionKey str key := m.poolKey(accountID, wsURL, sessionKey, proxyURL) if v, ok := m.connections.LoadAndDelete(key); ok { wc := v.(*WsConnection) + wc.recordCloseCause(connectionCloseCause{source: "local", reason: closePoolRemoved, bound: m.hasLiveResponseBinding(wc)}) m.removeResponseConnBindings(wc) wc.Close() if wc.session != nil { @@ -1303,6 +1082,7 @@ func (m *Manager) DiscardConnection(wc *WsConnection) { if wc == nil { return } + wc.recordCloseCause(connectionCloseCause{source: "local", reason: closeUnusable, bound: m.hasLiveResponseBinding(wc)}) if wc.PoolKey != "" { m.connections.CompareAndDelete(wc.PoolKey, wc) if wc.session != nil { @@ -1360,6 +1140,11 @@ func (m *Manager) evictResponseConnBindingsLocked(now time.Time) { // BindResponseConn 记录 response_id 由哪条连接产出（续链亲和）。 func (m *Manager) BindResponseConn(responseID string, wc *WsConnection, sessionKey string, accountID int64, apiKey string) { + m.bindResponseConn(responseID, responseConnBinding{conn: wc, sessionKey: sessionKey, accountID: accountID, apiKey: apiKey}) +} + +func (m *Manager) bindResponseConn(responseID string, binding responseConnBinding) { + wc := binding.conn responseID = strings.TrimSpace(responseID) if m == nil || responseID == "" || wc == nil { return @@ -1375,21 +1160,25 @@ func (m *Manager) BindResponseConn(responseID string, wc *WsConnection, sessionK if !m.responseConnBindingLiveLocked(wc) { return } + if wc.session != nil && connectionBudgetKind(wc) == capacityChat { + wc.session.markUserContext() + } if m.respConnBindings == nil { m.respConnBindings = make(map[string]responseConnBinding, 64) } + // store:false 只承诺连接上的最新响应可续链，旧绑定必须失效。 + for id, binding := range m.respConnBindings { + if binding.conn == wc { + delete(m.respConnBindings, id) + } + } // 有界保护：先清一轮过期项，仍超限则拒绝新增（旧绑定比新绑定更可能被续链）。 if len(m.respConnBindings) >= responseConnBindingMaxEntries { m.evictResponseConnBindingsLocked(now) } if len(m.respConnBindings) < responseConnBindingMaxEntries { - m.respConnBindings[responseID] = responseConnBinding{ - conn: wc, - sessionKey: sessionKey, - accountID: accountID, - apiKey: apiKey, - expiresAt: now.Add(responseConnBindingTTL), - } + binding.expiresAt = now.Add(responseConnBindingTTL) + m.respConnBindings[responseID] = binding } } @@ -1420,10 +1209,14 @@ func (m *Manager) lookupResponseConn(responseID string, accountID int64, apiKey // AcquirePreferredConnection 尝试独占 response_id 绑定的原连接（续链亲和）。 // 成功返回 (连接, pendingRequest, 池内 sessionKey)；绑定失效或连接忙时返回 nil， -// 调用方回退到常规 acquire 路径。忙时不等待：续链上下文虽在原连接，但排队会 -// 阻塞在前一个长响应后面，且该场景（同会话并发续链）极少，退化为缓存 miss 更稳。 +// 此单次尝试保留兼容入口；生产续链通过 acquireContinuation 有界等待。 func (m *Manager) AcquirePreferredConnection(responseID string, accountID int64, apiKey string) (*WsConnection, *PendingRequest, string) { - wc, sessionKey := m.lookupResponseConn(responseID, accountID, apiKey) + return m.acquirePreferredConnection(websocketContinuation{responseID: responseID, accountID: accountID, apiKey: apiKey}) +} + +func (m *Manager) acquirePreferredConnection(input websocketContinuation) (*WsConnection, *PendingRequest, string) { + accountID := input.accountID + wc, sessionKey := m.lookupContinuationConn(input) if wc == nil { return nil, nil, "" } @@ -1451,17 +1244,18 @@ func (m *Manager) AcquirePreferredConnection(responseID string, accountID int64, // probe 期间连接被其它 pool key 的容量裁剪安全回收。 accountLock.Lock() defer accountLock.Unlock() - if v, exists := m.connections.Load(wc.PoolKey); !exists || v != wc || !canReuseConnection(wc) { + current, _ := m.lookupContinuationConn(input) + if current != wc || !canReuseConnection(wc) { return nil, nil, "" } - pr, err := m.addPendingAndBeginReadLease(wc, sessionKey) + pr, err := m.addPendingAndBeginReadLease(wc, sessionKey, capacityChat) if err != nil { - m.DiscardConnection(wc) + m.discardAcquisitionFailure(wc, err) return nil, nil, "" } wc.Touch() if wc.account != nil { - m.trimIdleAccountConnections(accountID, accountConnectionLimit(wc.account), wc) + m.trimIdleAccountConnections(accountID, connectionLimit(connectionBudgetKind(wc)), wc) } return wc, pr, sessionKey } diff --git a/proxy/wsrelay/manager_test.go b/proxy/wsrelay/manager_test.go index 37495538c..85b98111c 100644 --- a/proxy/wsrelay/manager_test.go +++ b/proxy/wsrelay/manager_test.go @@ -14,6 +14,7 @@ import ( "time" "github.com/codex2api/auth" + "github.com/codex2api/proxy" "github.com/gorilla/websocket" ) @@ -272,6 +273,7 @@ func TestCanReuseConnection(t *testing.T) { session.SetConnected(true) conn := &WsConnection{session: session} conn.SetState(StateConnected) + conn.Touch() conn.lastUsed.Store(time.Now().Add(-IdleTimeout - time.Second).UnixNano()) if canReuseConnection(conn) { @@ -314,7 +316,8 @@ func TestAcquireConnectionWaitsWhileSessionHasPendingRequest(t *testing.T) { } } -func TestAcquireConnectionCapsIdleConnectionsAtAccountConcurrency(t *testing.T) { +func TestAcquireConnectionCapsIdleConnectionsAtConfiguredSlots(t *testing.T) { + setBusyRuntimeSettings(t, func(s *proxy.RuntimeSettings) { s.CodexWSDownstreamKeepaliveSlots = 2 }) upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { conn, err := upgrader.Upgrade(w, r, nil) @@ -332,7 +335,7 @@ func TestAcquireConnectionCapsIdleConnectionsAtAccountConcurrency(t *testing.T) manager := NewManager() t.Cleanup(manager.Stop) - account := &auth.Account{DBID: 42, DynamicConcurrencyLimit: 2} + account := &auth.Account{DBID: 42, DynamicConcurrencyLimit: 1} wsURL := "ws" + strings.TrimPrefix(server.URL, "http") connections := make([]*WsConnection, 0, 3) @@ -355,7 +358,7 @@ func TestAcquireConnectionCapsIdleConnectionsAtAccountConcurrency(t *testing.T) } if got := manager.ConnectionCount(); got != 2 { - t.Fatalf("ConnectionCount = %d, want account concurrency cap 2", got) + t.Fatalf("ConnectionCount = %d, want configured connection cap 2", got) } if connections[0].IsConnected() { t.Fatal("oldest idle connection should be evicted when the account cap is reached") @@ -423,7 +426,8 @@ func TestAcquireConnectionKeepsActualHandshakeUserAgentWhenReused(t *testing.T) } } -func TestAcquireConnectionTrimsIdleConnectionsAfterDynamicLimitDecrease(t *testing.T) { +func TestAcquireConnectionTrimsIdleConnectionsAfterConfiguredSlotDecrease(t *testing.T) { + setBusyRuntimeSettings(t, func(s *proxy.RuntimeSettings) { s.CodexWSDownstreamKeepaliveSlots = 3 }) manager := NewManager() t.Cleanup(manager.Stop) manager.probeFunc = func(*WsConnection) bool { return true } @@ -439,9 +443,10 @@ func TestAcquireConnectionTrimsIdleConnectionsAfterDynamicLimitDecrease(t *testi t.Fatalf("ConnectionCount before limit decrease = %d, want 3", got) } - account.Mu().Lock() - account.DynamicConcurrencyLimit = 1 - account.Mu().Unlock() + proxy.UpdateRuntimeSettings(func(s proxy.RuntimeSettings) proxy.RuntimeSettings { + s.CodexWSDownstreamKeepaliveSlots = 1 + return s + }) protected := connections[1] got, pending, err := manager.AcquireConnection( context.Background(), account, wsURL, "session-1", http.Header{}, "", @@ -452,21 +457,25 @@ func TestAcquireConnectionTrimsIdleConnectionsAfterDynamicLimitDecrease(t *testi if got != protected { t.Fatal("existing session connection should be reused after the limit decrease") } - if count := manager.ConnectionCount(); count != 1 { - t.Fatalf("ConnectionCount after limit decrease = %d, want 1", count) + if count := manager.ConnectionCount(); count != 2 { + t.Fatalf("ConnectionCount after limit decrease = %d, want active 1 + idle 1", count) } if !protected.IsConnected() { t.Fatal("the connection selected for reuse must remain connected") } - for i, wc := range connections { - if wc != protected && wc.IsConnected() { - t.Fatalf("idle connection %d remained connected after the limit decreased", i) - } + if connections[0].IsConnected() || !connections[2].IsConnected() { + t.Fatal("limit decrease should keep the newest idle connection beside the active one") } got.session.RemovePendingRequest(pending.RequestID) + manager.ReleaseConnection(got) + if manager.ConnectionCount() != 1 || !protected.IsConnected() || connections[2].IsConnected() { + t.Fatal("completed request should replace the older idle connection") + } } -func TestAcquireConnectionCountsPendingDialTowardAccountCap(t *testing.T) { +func TestAcquireConnectionCountsPendingDialTowardBlankBudget(t *testing.T) { + setBusyRuntimeSettings(t, func(s *proxy.RuntimeSettings) { s.CodexWSStatelessSlots = 1 }) + blankCtx := withConnectionCapacityKind(context.Background(), nil, []byte(`{"generate":false,"input":[]}`)) upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} started := make(chan struct{}) release := make(chan struct{}) @@ -501,7 +510,7 @@ func TestAcquireConnectionCountsPendingDialTowardAccountCap(t *testing.T) { firstResult := make(chan acquireResult, 1) go func() { wc, pending, err := manager.AcquireConnection( - context.Background(), account, wsURL, "session-first", http.Header{}, "", + blankCtx, account, wsURL, "session-first", http.Header{}, "", ) firstResult <- acquireResult{wc: wc, pending: pending, err: err} }() @@ -512,13 +521,13 @@ func TestAcquireConnectionCountsPendingDialTowardAccountCap(t *testing.T) { t.Fatal("first websocket dial did not reach the server") } - secondCtx, cancel := context.WithTimeout(context.Background(), 40*time.Millisecond) + secondCtx, cancel := context.WithTimeout(blankCtx, 40*time.Millisecond) defer cancel() _, _, secondErr := manager.AcquireConnection( secondCtx, account, wsURL, "session-second", http.Header{}, "", ) if secondErr == nil { - t.Fatal("second dial should wait for account connection capacity") + t.Fatal("second dial should be rejected at account connection capacity") } if got := attempts.Load(); got != 1 { t.Fatalf("websocket dial attempts = %d, want 1 while first dial is pending", got) diff --git a/proxy/wsrelay/oneshot_conn_test.go b/proxy/wsrelay/oneshot_conn_test.go index 179d2bc21..ac4f8797e 100644 --- a/proxy/wsrelay/oneshot_conn_test.go +++ b/proxy/wsrelay/oneshot_conn_test.go @@ -13,6 +13,7 @@ func newPooledConn(t *testing.T, manager *Manager, accountID int64, sessionID, p session := NewSession(accountID, manager) session.ID = sessionID session.SetConnected(true) + session.markUserContext() conn := &WsConnection{session: session, URL: "wss://example.test/responses", PoolKey: poolKey} conn.SetState(StateConnected) conn.lastUsed.Store(lastUsed.UnixNano()) @@ -41,7 +42,7 @@ func TestSortIdleForEvictionPrefersOneShotConns(t *testing.T) { } } -func TestTrimIdleAccountConnectionsEvictsOneShotFirst(t *testing.T) { +func TestTrimIdleChatConnectionsEvictsOldestFirst(t *testing.T) { manager := NewManager() t.Cleanup(manager.Stop) @@ -56,18 +57,18 @@ func TestTrimIdleAccountConnectionsEvictsOneShotFirst(t *testing.T) { manager.trimIdleAccountConnections(1, 2, nil) - if _, ok := manager.connections.Load(oneShotKey); ok { - t.Fatal("expected one-shot connection to be evicted first under capacity pressure") + if _, ok := manager.connections.Load(oneShotKey); !ok { + t.Fatal("newest one-shot chat should remain retained") } - if _, ok := manager.connections.Load(slotOldKey); !ok { - t.Fatal("expected LRU slot connection to survive while one-shot conn exists") + if _, ok := manager.connections.Load(slotOldKey); ok { + t.Fatal("oldest chat should be evicted regardless of its pool key") } if _, ok := manager.connections.Load(slotMidKey); !ok { t.Fatal("expected newer slot connection to survive") } } -func TestEnsureAccountConnectionCapacityEvictsOneShotFirst(t *testing.T) { +func TestChatInferenceReservationDoesNotConsumeIdleRetention(t *testing.T) { manager := NewManager() t.Cleanup(manager.Stop) @@ -78,12 +79,12 @@ func TestEnsureAccountConnectionCapacityEvictsOneShotFirst(t *testing.T) { newPooledConn(t, manager, 1, "det-key#0", slotKey, now.Add(-2*time.Minute)) newPooledConn(t, manager, 1, "stateless-11111111-aaaa-bbbb-cccc-000000000003", oneShotKey, now) - // count=2, pending=1, limit=3：需腾出 1 个槽位，应先逐出一次性连接。 - if !manager.ensureAccountConnectionCapacity(1, 3, "", 1) { - t.Fatal("expected capacity reservation to succeed after evicting one-shot conn") + // 两条空闲连接未超保留上限，新增推理不要求它们让位。 + if !manager.ensureAccountConnectionCapacity(connectionCapacityRequest{accountID: 1, limit: 3, extraPending: 1}) { + t.Fatal("inference should start without evicting retained idle connections") } - if _, ok := manager.connections.Load(oneShotKey); ok { - t.Fatal("expected one-shot connection to be evicted first") + if _, ok := manager.connections.Load(oneShotKey); !ok { + t.Fatal("inference admission should not consume idle retention") } if _, ok := manager.connections.Load(slotKey); !ok { t.Fatal("expected reusable slot connection to survive") diff --git a/proxy/wsrelay/read_pump.go b/proxy/wsrelay/read_pump.go index 02b48968d..489662148 100644 --- a/proxy/wsrelay/read_pump.go +++ b/proxy/wsrelay/read_pump.go @@ -139,10 +139,7 @@ func (wc *WsConnection) installControlHandlers() { ) }) wc.conn.SetPongHandler(func(appData string) error { - if wc.session != nil { - wc.session.HandlePong() - } - wc.Touch() + // Pong 只证明传输层存活，不能延长用户聊天的业务空闲窗口。 wc.touchInbound() wc.notifyProbePong(appData) return nil @@ -203,11 +200,10 @@ func (wc *WsConnection) runReadPump() { } if captured.idleSniff { // 空闲期帧从不投递给任何后续租约（帧在无租约时刻开始，归属已定）， - // 只在"丢弃续命"与"销毁连接"之间裁决。 + // 只在"丢弃元数据"与"销毁连接"之间裁决，不延长业务空闲窗口。 eventType := strings.TrimSpace(gjson.GetBytes(payload, "type").String()) if isIdleDroppableMetadataFrame(eventType) { wc.endIdleSniff() - wc.Touch() wc.touchInbound() continue } @@ -363,7 +359,7 @@ func (wc *WsConnection) enqueueBusinessFrameForCapturedLease(messageType int, pa func isReadLeaseTerminal(payload []byte) bool { switch gjson.GetBytes(payload, "type").String() { - case "response.completed", "response.failed", "response.done", "error": + case "response.completed", "response.incomplete", "response.failed", "response.done", "error": return true default: return false @@ -448,6 +444,7 @@ func (wc *WsConnection) finalizeReadPumpFailure(state *wsReadState) { state.mu.Lock() readerErr := state.readerErr state.mu.Unlock() + wc.recordReadClose(readerErr) idle := time.Duration(0) if ts := wc.lastUsed.Load(); ts > 0 { idle = time.Since(time.Unix(0, ts)).Round(time.Millisecond) diff --git a/proxy/wsrelay/session.go b/proxy/wsrelay/session.go index 6036d982d..e25e0b8ea 100644 --- a/proxy/wsrelay/session.go +++ b/proxy/wsrelay/session.go @@ -36,8 +36,8 @@ const ( // 写超时：30 秒 WriteTimeout = 30 * time.Second - // 空闲超时：5 分钟无活动则断开 - IdleTimeout = 5 * time.Minute + // 业务空闲超时：30 分钟未处理聊天请求或响应则断开，心跳不续期。 + IdleTimeout = 30 * time.Minute // 连接最大寿命：上游 chatgpt backend 对每条 Responses WS 连接强制 60 分钟 // 寿命上限，超限后该连接上的 response.create 一律返回 @@ -132,6 +132,9 @@ type Session struct { // 最后活跃时间 LastActiveAt time.Time + // 曾用于用户聊天；续链绑定过期后仍保留，不能降回空白预热槽位。 + userContext bool + // 连接状态 Connected bool @@ -163,20 +166,6 @@ func NewSession(accountID int64, manager *Manager) *Session { } } -// Touch 更新最后活跃时间 -func (s *Session) Touch() { - s.mu.Lock() - s.LastActiveAt = time.Now() - s.mu.Unlock() -} - -// IsExpired 检查会话是否过期（空闲超时） -func (s *Session) IsExpired() bool { - s.mu.RLock() - defer s.mu.RUnlock() - return time.Since(s.LastActiveAt) > connectionIdleTimeout() -} - // SetConnected 设置连接状态 func (s *Session) SetConnected(connected bool) { s.mu.Lock() @@ -343,11 +332,6 @@ func (s *Session) StopHeartbeat() { s.mu.Unlock() } -// HandlePong 处理 Pong 响应并刷新会话活动时间。 -func (s *Session) HandlePong() { - s.Touch() -} - // Close 关闭会话 func (s *Session) Close() { s.StopHeartbeat() diff --git a/proxy/wsrelay/session_activity.go b/proxy/wsrelay/session_activity.go new file mode 100644 index 000000000..0ce90ca84 --- /dev/null +++ b/proxy/wsrelay/session_activity.go @@ -0,0 +1,37 @@ +package wsrelay + +import "time" + +// Touch 更新最后业务活跃时间，并永久标记该会话曾承载用户聊天。 +func (s *Session) Touch() { + s.touchActivity(true) +} + +func (s *Session) touchActivity(userContext bool) { + s.mu.Lock() + s.LastActiveAt = time.Now() + s.userContext = s.userContext || userContext + s.mu.Unlock() +} + +func (s *Session) hasUserContext() bool { + s.mu.RLock() + defer s.mu.RUnlock() + return s.userContext +} + +func (s *Session) markUserContext() { + s.mu.Lock() + s.userContext = true + s.mu.Unlock() +} + +// IsExpired 检查业务空闲超时；空白会话在普通模式下继续保持连接。 +func (s *Session) IsExpired() bool { + s.mu.RLock() + defer s.mu.RUnlock() + if !s.userContext && !weakNetworkModeEnabled() { + return false + } + return time.Since(s.LastActiveAt) >= connectionIdleTimeout() +} diff --git a/proxy/wsrelay/session_test.go b/proxy/wsrelay/session_test.go index de7f1057d..bf69f075f 100644 --- a/proxy/wsrelay/session_test.go +++ b/proxy/wsrelay/session_test.go @@ -103,13 +103,14 @@ func TestSessionExpiration(t *testing.T) { } // 手动设置 LastActiveAt 为很久以前 + session.Touch() session.mu.Lock() - session.LastActiveAt = time.Now().Add(-10 * time.Minute) + session.LastActiveAt = time.Now().Add(-IdleTimeout - time.Minute) session.mu.Unlock() - // 现在应该过期（IdleTimeout = 5分钟） + // 超过当前业务空闲窗口后应过期。 if !session.IsExpired() { - t.Error("session should be expired after 10 minutes of inactivity") + t.Error("session should be expired after IdleTimeout of inactivity") } } diff --git a/proxy/wsrelay/subagent_transport_lane_test.go b/proxy/wsrelay/subagent_transport_lane_test.go index 4e165a56e..2b7352d3c 100644 --- a/proxy/wsrelay/subagent_transport_lane_test.go +++ b/proxy/wsrelay/subagent_transport_lane_test.go @@ -185,7 +185,8 @@ func TestSubagentTransportLanesAvoidSharedSessionBusyWait(t *testing.T) { } }) - t.Run("account capacity remains authoritative", func(t *testing.T) { + t.Run("inference lanes do not consume idle retention", func(t *testing.T) { + setBusyRuntimeSettings(t, func(s *proxy.RuntimeSettings) { s.CodexWSDownstreamKeepaliveSlots = 1 }) manager := NewManager() t.Cleanup(manager.Stop) limitedAccount := &auth.Account{DBID: 43, DynamicConcurrencyLimit: 1} @@ -209,8 +210,16 @@ func TestSubagentTransportLanesAvoidSharedSessionBusyWait(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) defer cancel() - if _, _, err := manager.AcquireConnection(ctx, limitedAccount, wsURL, childLane, childHeaders, ""); err == nil { - t.Fatal("child bypassed the account connection capacity limit") + child, childPending, err := manager.AcquireConnection(ctx, limitedAccount, wsURL, childLane, childHeaders, "") + if err != nil { + t.Fatalf("idle retention rejected active child lane: %v", err) + } + t.Cleanup(func() { + child.session.RemovePendingRequest(childPending.RequestID) + manager.DiscardConnection(child) + }) + if child == parent || !parent.IsConnected() || !child.IsConnected() { + t.Fatal("idle retention should preserve both active inference lanes") } }) } From 0d47b87857c33c32ec3ddc37f21f3ecaeda95120 Mon Sep 17 00:00:00 2001 From: bxb1337 Date: Fri, 9 Oct 2026 19:46:26 +0800 Subject: [PATCH 4/4] =?UTF-8?q?docs(wsrelay):=20=E8=AF=B4=E6=98=8E?= =?UTF-8?q?=E8=BF=9E=E6=8E=A5=E4=BF=9D=E7=95=99=E5=9B=9E=E6=94=B6=E4=B8=8E?= =?UTF-8?q?=E7=BB=AD=E9=93=BE=E6=81=A2=E5=A4=8D=E6=9C=BA=E5=88=B6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- docs/API.md | 4 ++-- docs/ARCHITECTURE.md | 4 ++-- docs/CONFIGURATION.md | 24 +++++++++++++++++++----- 3 files changed, 23 insertions(+), 9 deletions(-) diff --git a/docs/API.md b/docs/API.md index 2ce3748da..e1d3f326f 100644 --- a/docs/API.md +++ b/docs/API.md @@ -222,8 +222,8 @@ Messages 的 `tool_use.input` 必须使用对象，因此自由文本工具输 | include | array | 否 | 包含的额外字段 | | previous_response_id | string | 否 | 上一响应 ID，用于上下文连续 | -`previous_response_id` 的上下文先查当前进程的有界 L1。已认证请求按 API Key ID 隔离；未配置任何 API Key、显式启用 `CODEX_ALLOW_ANONYMOUS=true` 后放行的请求共用 `anon` 命名空间。Redis 模式在 L1 未命中时可从共享后端重建；后端值未超过重建上限但超过 L1 准入预算时仍可服务本次请求，只是不提升到 L1。Memory 模式没有共享 response context 后备，依赖上下文被判定为超限、已淘汰或缺失时可能返回 HTTP `409 response_context_unavailable`。共享后端暂时不可用且请求依赖该上下文时可能返回 HTTP `503 service_unavailable`。如果账号池存在可用的 relay-style 后备，网关可保留原始 `previous_response_id` 继续转发，而不是立即返回上述错误。客户端原生 Responses WebSocket 入口在上游连接正常时保留 `previous_response_id` 并只发送当轮增量，轮次开始时保留仍有效的 L1 祖先引用，快照合并与序列化在响应成功提交后才执行并写入本地缓存，共享后端（Redis）写入在后台完成，不占首字路径。缓存会收集 `response.output_item.done`，因此最终 `response.output` 为空时也能保留消息、`phase`、工具调用/结果及 Lite `additional_tools` 声明。`response.completed` 和 `response.incomplete` 成功提交后均可写入，仍受写入策略、TTL 和容量限制约束；按需写入模式下，`store` 未显式为 `false` 的 WebSocket 会话从根轮起即有写入资格；显式 `store:false` 的会话（如 Codex CLI 全量上下文）不写缓存，其历史无法事后恢复。续链失效、换号或转为 HTTP 时，只有能够恢复所需上下文才继续发送；缺失、超限或不可移植的加密状态返回 `response_context_unavailable` 错误帧，共享后端故障返回 `service_unavailable`，随后关闭连接（409 使用 1008，后端暂时不可用使用 1011）。客户端应重发完整上下文并开始新的响应链。 -祖先在轮次开始前已过期或被淘汰时，网关不会把增量伪装成完整快照；该链后续健康轮也不能自行补全历史。缓存 TTL 为 10 分钟，Memory 模式重启会丢失缓存。设置环境变量 `CODEX_WS_CONTINUATION_FAIL_OPEN=true` 可退回旧行为：上下文不可恢复时剥离 `previous_response_id` 后按原样转发，上游将看不到历史；这次有损降级产生的响应不写回放缓存，关闭逃生阀后也不会信任残缺快照。 +`previous_response_id` 的上下文先查当前进程的有界 L1。本地与 Redis 共享上下文的默认绝对 TTL 为 45 分钟，覆盖聊天上游连接 30 分钟空闲回收后的恢复窗口，仍受缓存容量和账号兼容性限制。已认证请求按 API Key ID 隔离；未配置任何 API Key、显式启用 `CODEX_ALLOW_ANONYMOUS=true` 后放行的请求共用 `anon` 命名空间。Redis 模式在 L1 未命中时可从共享后端重建；后端值未超过重建上限但超过 L1 准入预算时仍可服务本次请求，只是不提升到 L1。Memory 模式没有共享 response context 后备，依赖上下文被判定为超限、已淘汰或缺失时可能返回 HTTP `409 response_context_unavailable`。共享后端暂时不可用且请求依赖该上下文时可能返回 HTTP `503 service_unavailable`。如果账号池存在可用的 relay-style 后备，网关可保留原始 `previous_response_id` 继续转发，而不是立即返回上述错误。客户端原生 Responses WebSocket 入口在上游连接正常时保留 `previous_response_id` 并只发送当轮增量，轮次开始时保留仍有效的 L1 祖先引用，快照合并与序列化在响应成功提交后才执行并写入本地缓存，共享后端（Redis）写入在后台完成，不占首字路径。缓存会收集 `response.output_item.done`，因此最终 `response.output` 为空时也能保留消息、`phase`、工具调用/结果及 Lite `additional_tools` 声明。`response.completed` 和 `response.incomplete` 成功提交后均可写入，仍受写入策略、TTL 和容量限制约束；按需写入模式下，`store` 未显式为 `false` 的 WebSocket 会话从根轮起即有写入资格；显式 `store:false` 的会话（如 Codex CLI 全量上下文）不写缓存，其历史无法事后恢复。续链失效、换号或转为 HTTP 时，只有能够恢复所需上下文才继续发送；缺失、超限或不可移植的加密状态返回 `response_context_unavailable` 错误帧，共享后端故障返回 `service_unavailable`，随后关闭连接（409 使用 1008，后端暂时不可用使用 1011）。客户端应重发完整上下文并开始新的响应链。 +祖先在轮次开始前已过期或被淘汰时，网关不会把增量伪装成完整快照；该链后续健康轮也不能自行补全历史。缓存 TTL 为 45 分钟，Memory 模式重启会丢失缓存。设置环境变量 `CODEX_WS_CONTINUATION_FAIL_OPEN=true` 可退回旧行为：上下文不可恢复时剥离 `previous_response_id` 后按原样转发，上游将看不到历史；这次有损降级产生的响应不写回放缓存，关闭逃生阀后也不会信任残缺快照。 对于原生 WebSocket 的结构化输出，`gpt-6-astra` 和 `gpt-5.6-luna` 在显式启用 Responses Lite 时保留 JSON Schema 的 `minLength` / `maxLength`。Lite 信号可来自 `client_metadata.ws_request_header_x_openai_internal_codex_responses_lite=true`、`X-OpenAI-Internal-Codex-Responses-Lite: true` 请求头，或由 Payload Rules 注入该元数据标记。该放行仅用于结构化输出；已有工具参数清洗继续使用保守规则。长度约束在入口准备阶段暂时保留，最终出站前才根据最终模型、规则改写后的 Lite 信号、账号 Lite 能力和实际传输统一处理。其他模型、最终未启用 Lite、HTTP 和 Compact 请求沿用原有清洗策略，HTTP 降级请求不会携带不适用的约束。 diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index 352d66d2b..7747dc433 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -528,14 +528,14 @@ L1 受 4,096 条、16 MiB 逻辑 JSON 字节和单条 64 KiB 预算限制；大 #### Responses 上下文缓存 -`previous_response_id` 连续请求使用一层每进程独立的有界 L1。完成响应时先按完整 call/output 组保留尾部最多 200 个 raw item，再写入 L1；默认总量 64 MiB、单条准入 8 MiB、最多 2,000 条、绝对 TTL 10 分钟。LRU 同时受条数和逻辑字节预算约束，命中不会延长绝对 TTL。 +`previous_response_id` 连续请求使用一层每进程独立的有界 L1。完成响应时先按完整 call/output 组保留尾部最多 200 个 raw item，再写入 L1；默认总量 64 MiB、单条准入 8 MiB、最多 2,000 条、绝对 TTL 45 分钟。LRU 同时受条数和逻辑字节预算约束，命中不会延长绝对 TTL。 ```text 完成响应 │ ├─ pair-safe 尾部裁剪（最多 200 items） │ - ├─ 本进程 L1（64 MiB total / 8 MiB entry / 2,000 entries / 10m） + ├─ 本进程 L1（64 MiB total / 8 MiB entry / 2,000 entries / 45m） │ └─ Redis 共享 response context（仅 Redis 模式） diff --git a/docs/CONFIGURATION.md b/docs/CONFIGURATION.md index d33ddb03f..7f07ee404 100644 --- a/docs/CONFIGURATION.md +++ b/docs/CONFIGURATION.md @@ -233,13 +233,17 @@ Responses 连续请求会按 `previous_response_id` 重建上下文。每个 Cod | `response_cache_reconstruct_max_bytes` | 67,108,864 bytes（64 MiB） | 整数 8-512 MiB | 从共享后端读取并重建一条上下文时允许的逻辑 payload 上限 | | `response_cache_config_generation` | 1 | 只读 | 配置发生实际变化时递增；客户端不能写入 | -设置页会把三个预算作为一个原子更新发送。管理 API 也支持只提交部分预算，服务端会在数据库事务中与当前值合并并校验，提交成功后才应用到本实例。固定边界不随这三个设置变化：最多 2,000 条、10 分钟绝对 TTL、每条最多 200 个 raw item。降低预算会立即收缩本地 L1，并可能淘汰已有上下文。 +设置页会把三个预算作为一个原子更新发送。管理 API 也支持只提交部分预算，服务端会在数据库事务中与当前值合并并校验，提交成功后才应用到本实例。固定边界不随这三个设置变化：最多 2,000 条、45 分钟绝对 TTL、每条最多 200 个 raw item。降低预算会立即收缩本地 L1，并可能淘汰已有上下文。 Redis 模式会把 response context 保存到共享后端。后端值在重建上限内但超过 L1 单条或总量准入预算时，仍可服务当前请求，但不会提升到本地 L1。Memory 模式只保留本进程 L1，不存在第二份共享 response context；已知超限/淘汰，或依赖的必需上下文缺失/过期时，可能导致 HTTP `409 response_context_unavailable`。Redis 值损坏或超过重建上限且无法走 relay 后备时也可能返回 409。共享后端暂时不可用时，依赖该上下文且无法走 relay 后备的请求可能返回 HTTP `503`。 只有预算实际变化时才会分配并递增 generation；同值更新或空更新不会递增。当前实例在数据库提交后立即应用，其他实例每 5 秒轮询一次，只应用更新的 generation；单次读取最多等待 3 秒。同步失败时保留最后一次有效配置，并在运维页显示错误，后续轮询成功后自动恢复。 -这些预算覆盖 HTTP Responses/Compact 和原生 Responses WebSocket 的本地回放上下文。健康的原生 WS 续链仍保留 `previous_response_id` 交给上游；需要降级时可使用完整本地快照。原生 WS 显式 `store:false` 的请求不写回放缓存。 +这些预算覆盖 HTTP Responses/Compact 和原生 Responses WebSocket 的回放上下文。健康的原生 WS 续链仍保留 `previous_response_id` 交给原上游连接。显式 `store:false` 的原生 WS 请求会在响应成功过滤并提交后保存完整快照：Memory 使用本进程 L1，Redis 同时保存共享快照，支持下游重连到其他实例时恢复；此快照不受按需写入策略影响。快照按数据库作用域、API Key、权限范围、稳定会话、执行通道和 `stream_id` 隔离，沿用上述预算与绝对 TTL；读取或提升到 L1 不延长原过期时间。超过 L1 预算但仍在重建上限内时可从 Redis 服务，超过条目数或重建上限则拒绝整条快照，不截尾丢失祖先。 + +Redis 原生 WS 快照将完整历史、账号 ID、凭证代际、兼容域和过期时间原子保存为同一记录。恢复加密 reasoning/compaction 时优先选择原账号，再核对凭证与兼容域；账号不兼容时不丢弃必要推理来强行续链。终态发送前登记共享提交屏障，其他实例等待快照提交后再读取；共享快照暂时不可用且 L1 无法恢复时返回可恢复的 `503`，完整 L1 命中不额外查询 Redis。上游实体连接与连接池仍属于原实例，Redis 分享上下文与协调状态，不迁移实体 WS。 + +原连接丢失或上游明确拒绝旧 ID 时，网关在尚未交付内容或工具调用的条件下，最多恢复一次：移除 `previous_response_id` 并发送完整快照。加密推理/压缩项仅在原账号、凭证代次和兼容域均匹配时回放。祖先缺失、超限或兼容性无法确认时，返回 `status:409`、`code:previous_response_not_found`，供客户端重发完整输入；健康下游 WS 保持可用。Memory 模式下进程重启会丢失本地快照；Redis 模式可在原绝对 TTL 内跨实例或重启恢复共享快照。 共享后端写入的异步与同步路径统一限制为最多 16 个在途写、64 MiB 在途逻辑正文、64 个等待者；在复制/编码之前获取额度，最多等待 5 秒。超过 64 MiB 的既有合法单条上下文可独占写入器，因此其逻辑上限为普通预算与最大在途单条的较大值，不会静默丢弃大快照。普通写入 I/O deadline 为 2 秒，关停同步写为 500 毫秒。饱和时响应收尾及同 WS 后续轮次可能等待写入额度；写失败后 L1 仍可服务，必须依赖该快照但 L1/共享后端均缺失时返回 503。 @@ -262,6 +266,8 @@ Redis 模式会把 response context 保存到共享后端。后端值在重建 | `CodexForceWebsocket` | bool | false | - | 强制 Codex 上游走 WebSocket 长连接（复用连接池），更接近官方 CLI 体验；关闭时走原有 HTTP 请求 | | `CodexWSKeepaliveEnabled` | bool | false | - | 启用上游 WS 空闲连接保活（后台仅发 Ping，不发起新请求、不消耗账号额度） | | `CodexWSKeepaliveIntervalSec` | int | 60 | 10-600 | WS 保活 Ping 间隔（秒），仅在 `CodexWSKeepaliveEnabled` 开启时生效 | +| `CodexWSDownstreamKeepaliveSlots` | int | 8 | 0-32 | 每账号在推理输出结束后保留的空闲聊天上游连接数，超额按 LRU 回收 | +| `CodexWSStatelessSlots` | int | 8 | 0-32 | 每账号空白上游连接预算及无状态池扫描窗口 | | `CodexWSHideUpstreamErrors` | bool | true | - | WS 上游最终失败时向客户端隐藏原始错误，返回统一友好提示；原始错误仍记录在后台日志/用量记录 | | `CodexWSSilentRetryEnabled` | bool | true | - | WS 首包前遇到限流、额度耗尽、5xx、读取错误或超时时，静默换账号并重建上游 WS | | `CodexWSSilentMaxRetries` | int | 2 | 0-10 | WS 首包前静默重试上限；`0` 禁用该预算 | @@ -278,7 +284,7 @@ Redis 模式会把 response context 保存到共享后端。后端值在重建 索引选号的过滤器与准入回调在调度锁外执行，返回后重新检查候选代次、账号状态和并发；`Disabled` / `DispatchPaused` 同样阻止最终占位。已有会话绑定、容量借号保护和有状态续链的账号约束保持生效。 -账号满载时，等待队列同时受全局与单个 API Key 上限约束，所有使用该账号池等待路径的协议和上游共用预算。HTTP 队列溢出返回 `503` 和 `Retry-After: 1`；已提交的 SSE 输出对应协议的失败事件，WebSocket 返回错误帧并以 `1013` 关闭，文案提示 1 秒后重试。该本地过载不会进入持续重试的上游换号循环，也不会被误报为账号额度耗尽。已有等待者不因调低上限而被取消；环境变量不是跨实例配额，也不限制已经在上游执行的请求。 +账号满载时，等待队列同时受全局与单个 API Key 上限约束，所有使用该账号池等待路径的协议和上游共用预算。HTTP 队列溢出返回 `503` 和 `Retry-After: 1`；已提交的 SSE 输出对应协议的失败事件，Responses WebSocket 返回带 `status:503` 的错误帧并保持连接可用，文案提示 1 秒后重试。该本地过载不会进入持续重试的上游换号循环，也不会被误报为账号额度耗尽。已有等待者不因调低上限而被取消；环境变量不是跨实例配额，也不限制已经在上游执行的请求。 队列内按 API Key 轮转，同一 Key 按可尝试请求的入队顺序唤醒。普通单槽释放只唤醒一个等待者；已有续链绑定和排除账号用于跳过不匹配的通知，剩余模型、分组和 scope 过滤仍在锁外执行，失败后把机会交给后续等待者。一轮通知最多尝试当前等待集合一次，同时最多有 8 个通知驱动的选号；密集释放合并为后续容量检查。公平性针对已排队的请求，不承诺绕过快路径新请求的全局先来先服务，也不保证不同过滤条件获得相同吞吐。SSE/WS 心跳不重新入队。每个有等待者的账号池只使用一个每秒恢复检查定时器，空队列自动停止；账号冷却自身的到期恢复通知仍然生效。 @@ -368,8 +374,16 @@ Claude / Grok / Antigravity 等中继型账号不经 Resin，始终按第 2、3 ### WebSocket 连接池与 1009 降级 -- 每个账号的上游物理 WebSocket 连接数受其当前 `DynamicConcurrencyLimit` 限制。 -- 新建或复用连接时如果超过新上限，只淘汰最老的空闲连接；当前请求使用的连接和其他活跃连接不会被中断。 +- 管理后台 **设置 → WebSocket（Codex 上游）→ 聊天续链连接保留数**（`codex_ws_downstream_keepalive_slots`）限制每账号在 Codex 输出完毕、等待用户继续输入时保留的空闲聊天上游连接数，默认 8、范围 0～32，所有 Key 共用，保存后热更新。设为 0 时聊天输出收尾即关闭上游连接，后续通过仍有效的上下文恢复。拨号及推理中的连接不占空闲保留额度；普通会话、无状态客户端、overflow、一次性连接与后台执行道的聊天输出收尾后，均按相同规则保留。 +- 原 `codex_ws_stateless_slots`（空白预热连接保留数）独立限制每账号未承载用户聊天上下文的空白连接，并控制无状态连接池扫描窗口，默认 8、范围 0～32。没有聊天输入、指令、工具定义、模板或续链 ID 的明确预热请求使用空白预算。首次在空白连接上开始聊天前，连接退出空白预算，进入推理状态，输出结束后才进入聊天空闲保留池；两个保留池不互相占位。设为 0 时不扫描通用池、不保留空白连接，预热请求仍可临时建连并在收尾后关闭。两项之和是正常空闲保留量，不是全部实体连接的硬上限；推理连接另外由原账号/Key 请求并发约束。 +- 连接预算独立于 `DynamicConcurrencyLimit`。API Key 和账号的请求并发继续沿用原配置、队列与计数方式。例如 4 个 Key 各限并发 3、账号限并发 3，仍最多同时处理 3 个请求；若希望保留 12 个独立聊天会话的上游连接，可将“聊天续链连接保留数”设为 12 或 16，不会自动修改现有配置。 +- 空白拨号仍按原预算原子预留；聊天拨号、在途以及摘池后的收尾只登记实体 socket 生命周期，直到实际关闭后幂等解除登记。聊天推理可以在空闲保留池已满时正常开始，不会因这个设置返回 503；原来的 Key/账号请求并发和 busy 等待规则继续生效。 +- 聊天输出收尾后更新最后业务使用时间，空闲聊天连接超过上限时按 LRU 关闭最久未使用的连接，包括仍有有效续链绑定的连接。回收在账号锁下复验连接身份和 pending 请求，跳过在途请求、拨号及空白连接；旧连接的晚到回调不会影响新连接。被回收会话再次输入时，可从仍有效且兼容的完整上下文快照重建上游连接。空白预算不足时仍可能返回 `503 upstream_ws_connection_capacity_exceeded`。 +- 下调聊天保留数后，取连、输出收尾或每 30 秒后台巡检会按新上限淘汰最旧空闲连接；推理不受影响。普通模式下，曾用于用户聊天的上游 WS 在最后一次业务使用后闲置 30 分钟也会关闭，Ping/Pong、下游连接保持打开和后台元数据不续期。空白连接不受聊天闲置限制，弱网模式及到龄轮转沿用原规则。 +- 空闲回收会实际关闭旧上游 socket，清除对应会话和续链绑定。后续请求重新握手创建新的上游 WS，并在需要时恢复完整上下文；旧连接的晚到清理不会移除同池键下的新连接。 +- `CODEX_WS_STATELESS_ONESHOT=1` 禁止通用槽位复用；响应正常收尾后，有有效续链绑定的连接按聊天空闲上限保留，无绑定的一次性连接直接销毁。异常或未完整消费的响应继续销毁对应坏连接。 +- 连接关闭日志区分本地回收、空闲超时、到龄轮转、上游关闭帧与网络错误，并记录关闭码、续链绑定状态和剩余连接数。Redis 继续用于完整续链快照及恢复元数据，实体连接预算在单实例内协调。本地与 Redis 上下文默认保存 45 分钟，覆盖聊天连接 30 分钟空闲回收后的恢复窗口；缓存容量淘汰、账号不兼容或已过期仍可能使恢复失败。 +- 带 `previous_response_id` 的请求只取回原连接，并复验账号、模型、出口和客户端身份；忙时可取消地等待，遵守现有 busy 等待上限。等待耗尽返回 `status:503`，API Key 并发拒绝返回 `status:429`，均保留下游 WS 供下一轮使用。 - 上游在尚未向下游输出内容时返回 close 1009，或本地读取触发等价的 read-limit 错误，网关会保留同一账号租约和已解析代理，最多降级一次 HTTP。 - 1009 属于传输限制，不降低账号健康度，也不触发鉴权探针；一旦已向下游输出内容，就不会再发起 HTTP 降级，避免重复请求和重复计费。