-
Notifications
You must be signed in to change notification settings - Fork 345
fix(proxy): keep long Responses streams alive #561
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,57 @@ | ||
| package proxy | ||
|
|
||
| import ( | ||
| "context" | ||
| "sync" | ||
| "time" | ||
| ) | ||
|
|
||
| const ( | ||
| // 普通 Responses SSE 在上游长时间只思考、不产出可转发事件时,也要持续 | ||
| // 刷新下游链路的 idle timer。10 秒低于常见的 30/60 秒反代超时,同时 | ||
| // 每分钟仅增加几十字节;SSE 注释不会被 Codex 客户端当作模型输出。 | ||
| defaultDownstreamSSEKeepaliveInterval = 10 * time.Second | ||
| downstreamSSEKeepaliveComment = ": keepalive\n\n" | ||
| ) | ||
|
|
||
| // 变量形式只为处理器级测试缩短等待;生产运行保持默认 10 秒。 | ||
| var downstreamSSEKeepaliveInterval = defaultDownstreamSSEKeepaliveInterval | ||
|
|
||
| // startDownstreamSSEKeepalive 周期执行 writeKeepalive,直到请求取消、写失败 | ||
| // 或调用 stop。stop 会等待 goroutine 完整退出,保证流收尾后不再并发写入。 | ||
| func startDownstreamSSEKeepalive(ctx context.Context, interval time.Duration, writeKeepalive func() bool) func() { | ||
| if interval <= 0 || writeKeepalive == nil { | ||
| return func() {} | ||
| } | ||
| if ctx == nil { | ||
| ctx = context.Background() | ||
| } | ||
|
|
||
| stopCh := make(chan struct{}) | ||
| done := make(chan struct{}) | ||
| var stopOnce sync.Once | ||
| go func() { | ||
| defer close(done) | ||
| ticker := time.NewTicker(interval) | ||
| defer ticker.Stop() | ||
| for { | ||
| select { | ||
| case <-ticker.C: | ||
| if !writeKeepalive() { | ||
| return | ||
| } | ||
| case <-ctx.Done(): | ||
| return | ||
| case <-stopCh: | ||
| return | ||
| } | ||
| } | ||
| }() | ||
|
|
||
| return func() { | ||
| stopOnce.Do(func() { | ||
| close(stopCh) | ||
| <-done | ||
| }) | ||
| } | ||
| } |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,85 @@ | ||
| package proxy | ||
|
|
||
| import ( | ||
| "context" | ||
| "sync/atomic" | ||
| "testing" | ||
| "time" | ||
| ) | ||
|
|
||
| func TestDownstreamSSEKeepaliveStopsAndJoins(t *testing.T) { | ||
| var writes atomic.Int32 | ||
| firstWrite := make(chan struct{}, 1) | ||
| stop := startDownstreamSSEKeepalive(context.Background(), time.Millisecond, func() bool { | ||
| writes.Add(1) | ||
| select { | ||
| case firstWrite <- struct{}{}: | ||
| default: | ||
| } | ||
| return true | ||
| }) | ||
|
|
||
| select { | ||
| case <-firstWrite: | ||
| case <-time.After(100 * time.Millisecond): | ||
| t.Fatal("keepalive did not fire") | ||
| } | ||
| stop() | ||
| stoppedAt := writes.Load() | ||
| time.Sleep(10 * time.Millisecond) | ||
| if got := writes.Load(); got != stoppedAt { | ||
| t.Fatalf("keepalive wrote after stop returned: %d -> %d", stoppedAt, got) | ||
| } | ||
| } | ||
|
|
||
| func TestDownstreamSSEKeepaliveStopsOnContextCancel(t *testing.T) { | ||
| ctx, cancel := context.WithCancel(context.Background()) | ||
| var writes atomic.Int32 | ||
| firstWrite := make(chan struct{}, 1) | ||
| stop := startDownstreamSSEKeepalive(ctx, time.Millisecond, func() bool { | ||
| writes.Add(1) | ||
| select { | ||
| case firstWrite <- struct{}{}: | ||
| default: | ||
| } | ||
| return true | ||
| }) | ||
| defer stop() | ||
|
|
||
| select { | ||
| case <-firstWrite: | ||
| case <-time.After(100 * time.Millisecond): | ||
| t.Fatal("keepalive did not fire") | ||
| } | ||
| cancel() | ||
| stop() | ||
| stoppedAt := writes.Load() | ||
| time.Sleep(10 * time.Millisecond) | ||
| if got := writes.Load(); got != stoppedAt { | ||
| t.Fatalf("keepalive wrote after context cancellation: %d -> %d", stoppedAt, got) | ||
| } | ||
| } | ||
|
|
||
| func TestDownstreamSSEKeepaliveStopsWhenWriterFails(t *testing.T) { | ||
| var writes atomic.Int32 | ||
| firstWrite := make(chan struct{}, 1) | ||
| stop := startDownstreamSSEKeepalive(context.Background(), time.Millisecond, func() bool { | ||
| writes.Add(1) | ||
| select { | ||
| case firstWrite <- struct{}{}: | ||
| default: | ||
| } | ||
| return false | ||
| }) | ||
| defer stop() | ||
|
|
||
| select { | ||
| case <-firstWrite: | ||
| case <-time.After(100 * time.Millisecond): | ||
| t.Fatal("keepalive did not fire") | ||
| } | ||
| stop() | ||
| if got := writes.Load(); got != 1 { | ||
| t.Fatalf("writer failure must stop keepalive after one write, got %d", got) | ||
| } | ||
| } |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -2541,6 +2541,9 @@ func (h *Handler) effectiveMaxRateLimitRetries(account *auth.Account, fallback i | |
| const ( | ||
| logStatusClientClosed = 499 | ||
| logStatusUpstreamStreamBreak = 598 | ||
| // AccessLogStatusContextKey 允许流处理器在 HTTP 200 header 已提交后, | ||
| // 把最终的内部结果(如客户端断开的 499)提供给访问日志中间件。 | ||
| AccessLogStatusContextKey = "x-access-log-status" | ||
| ) | ||
|
|
||
| // upstreamStreamBreakMessage 是断流反馈给下游的稳定可读消息;机器识别用 | ||
|
|
@@ -3845,6 +3848,7 @@ func (h *Handler) Responses(c *gin.Context) { | |
| // 并发方,锁零竞争。 | ||
| var downstreamMu sync.Mutex | ||
| var pendingFirstTokenEvents bytes.Buffer | ||
| contEnabled, contMaxRounds := codexContinueThinkingSettings() | ||
| // 前置元数据事件立即透传(旧版兼容,issue #425):每个 attempt 取一次快照, | ||
| // 热更新对新请求生效,流转发中途不切换缓冲策略。 | ||
| preflightPassthrough := CurrentRuntimeSettings().CodexPreflightSSEPassthrough | ||
|
|
@@ -3853,6 +3857,12 @@ func (h *Handler) Responses(c *gin.Context) { | |
| h.recordCompactionProvenanceFromPayload(context.Background(), account, data) | ||
| downstreamMu.Lock() | ||
| defer downstreamMu.Unlock() | ||
| // 上游 context 为了提取 usage 会在客户端断开后再排空最多 5 秒; | ||
| // 但下游 context 一旦取消,绝不能再尝试写 SSE,否则下一帧必然 | ||
| // 变成 broken pipe。继续解析帧只用于拿 response.completed/usage。 | ||
| if c.Request.Context().Err() != nil { | ||
| clientGone = true | ||
| } | ||
| parsed := gjson.ParseBytes(data) | ||
| eventType := parsed.Get("type").String() | ||
|
|
||
|
|
@@ -3951,7 +3961,35 @@ func (h *Handler) Responses(c *gin.Context) { | |
| // 思考截断自动续想(默认关闭):开启时用折叠状态机包裹 forward, | ||
| // 命中 518n-2 截断指纹则用同一账号续发上游并折叠成单响应; | ||
| // 关闭时保持原有逐事件透传路径,字节级零变化。 | ||
| contEnabled, contMaxRounds := codexContinueThinkingSettings() | ||
| // 默认(未启用自动续想)路径也可能在 xhigh/max 的长推理阶段数十秒 | ||
| // 没有可转发帧。定期写标准 SSE 注释,避免本机反代/Tailscale | ||
| // 把健康长流误判为空闲连接。自动续想路径已有自己的隐藏轮保活, | ||
| // 不重复启动第二个 ticker。 | ||
| stopDownstreamKeepalive := func() {} | ||
| if !contEnabled { | ||
| stopDownstreamKeepalive = startDownstreamSSEKeepalive(c.Request.Context(), downstreamSSEKeepaliveInterval, func() bool { | ||
| downstreamMu.Lock() | ||
| defer downstreamMu.Unlock() | ||
| if c.Request.Context().Err() != nil { | ||
| clientGone = true | ||
| return false | ||
| } | ||
| if clientGone { | ||
| return false | ||
| } | ||
| // 首个真实字节前不能写注释,否则会提前提交 HTTP 200, | ||
| // 破坏首包前 response.failed 的真实状态码与换号重试语义。 | ||
| if !wroteAnyBody { | ||
| return true | ||
| } | ||
| if err := streamWriter.WriteSSEComment(downstreamSSEKeepaliveComment); err != nil { | ||
| writeErr = err | ||
| clientGone = true | ||
| return false | ||
| } | ||
| return true | ||
| }) | ||
| } | ||
|
Comment on lines
+3968
to
+3992
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift Add downstream keepalives to relay Responses streams. Lines 3968-3992 install the ticker only in the non-relay Responses path. The Apply the same post-first-byte, serialized keepalive behavior to both relay streaming paths. 🤖 Prompt for AI Agents |
||
| if contEnabled { | ||
| fold := &continueFold{ | ||
| baseBody: upstreamBody, | ||
|
|
@@ -4032,6 +4070,7 @@ func (h *Handler) Responses(c *gin.Context) { | |
| } else { | ||
| readErr = ReadSSEStream(resp.Body, forward) | ||
| } | ||
| stopDownstreamKeepalive() | ||
| // 仅在真的写过 body 时才做收尾 flush:flusher.Flush 会先提交 HTTP 200 header, | ||
| // 零写入时提前 flush 会让循环外的 c.JSON(4xx) 失效(status 已定型为 200)。 | ||
| if writeErr == nil && wroteAnyBody { | ||
|
|
@@ -4173,6 +4212,9 @@ func (h *Handler) Responses(c *gin.Context) { | |
|
|
||
| h.store.BindSessionAffinity(affinityKey, account, proxyURL) | ||
| logStatusCode := outcome.logStatusCode | ||
| if logStatusCode != http.StatusOK { | ||
| c.Set(AccessLogStatusContextKey, logStatusCode) | ||
| } | ||
|
Comment on lines
+4215
to
+4217
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift Set the final access-log status before relay stream returns. Lines 4215-4217 run only after the non-relay Responses path completes. The relay branches return earlier after Set the context status in each relay stream finalization path before it returns. 🤖 Prompt for AI Agents |
||
| if outcome.logStatusCode != http.StatusOK { | ||
| log.Printf("流异常结束 (account %d, /v1/responses, status %d): %s,已转发约 %d 字符", account.ID(), outcome.logStatusCode, outcome.failureMessage, deltaCharCount) | ||
| if deltaCharCount > 0 { | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Cover the access-log suppression branch.
The test uses
POST /v1/responses, whichshouldSkipAccessLogdoes not suppress. It verifies the logged499, but it does not verify that the override is applied before suppression. Add a case forGET /api/admin/healthwith wire status200and override499, then assert that the access log is emitted.🤖 Prompt for AI Agents