From 4aa2ae01413fbfbcaea07cc54c99edd29852bd13 Mon Sep 17 00:00:00 2001
From: dundunge <105058271+dundunge@users.noreply.github.com>
Date: Sat, 22 Aug 2026 00:50:54 +0800
Subject: [PATCH] fix(proxy): keep long responses streams alive

---
 main.go                            | 12 +++--
 main_test.go                       | 35 ++++++++++++
 proxy/downstream_keepalive.go      | 57 ++++++++++++++++++++
 proxy/downstream_keepalive_test.go | 85 ++++++++++++++++++++++++++++++
 proxy/handler.go                   | 44 +++++++++++++++-
 proxy/handler_test.go              | 53 +++++++++++++++++++
 6 files changed, 282 insertions(+), 4 deletions(-)
 create mode 100644 proxy/downstream_keepalive.go
 create mode 100644 proxy/downstream_keepalive_test.go

diff --git a/main.go b/main.go
index 37790db33..f619534fb 100644
--- a/main.go
+++ b/main.go
@@ -602,7 +602,13 @@ func loggerMiddleware() gin.HandlerFunc {
 		start := time.Now()
 		c.Next()
 		latency := time.Since(start)
-		if shouldSkipAccessLog(c.Request.Method, c.Request.URL.Path, c.Writer.Status()) {
+		statusCode := c.Writer.Status()
+		if override, ok := c.Get(proxy.AccessLogStatusContextKey); ok {
+			if status, valid := override.(int); valid && status >= 100 && status <= 599 {
+				statusCode = status
+			}
+		}
+		if shouldSkipAccessLog(c.Request.Method, c.Request.URL.Path, statusCode) {
 			return
 		}
 
@@ -639,9 +645,9 @@ func loggerMiddleware() gin.HandlerFunc {
 		}
 
 		if emailStr != "" {
-			log.Printf("%s %s %d %v%s [%s] [%s]", c.Request.Method, c.Request.URL.Path, c.Writer.Status(), latency, tagStr, emailStr, proxyStr)
+			log.Printf("%s %s %d %v%s [%s] [%s]", c.Request.Method, c.Request.URL.Path, statusCode, latency, tagStr, emailStr, proxyStr)
 		} else {
-			log.Printf("%s %s %d %v%s", c.Request.Method, c.Request.URL.Path, c.Writer.Status(), latency, tagStr)
+			log.Printf("%s %s %d %v%s", c.Request.Method, c.Request.URL.Path, statusCode, latency, tagStr)
 		}
 	}
 }
diff --git a/main_test.go b/main_test.go
index 997731e03..94139a211 100644
--- a/main_test.go
+++ b/main_test.go
@@ -8,6 +8,7 @@ import (
 	"strings"
 	"testing"
 
+	"github.com/codex2api/proxy"
 	"github.com/gin-gonic/gin"
 )
 
@@ -171,3 +172,37 @@ func TestLoggerMiddlewareRedactsSensitiveContext(t *testing.T) {
 		}
 	}
 }
+
+func TestLoggerMiddlewareUsesStreamingOutcomeOverride(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+
+	var logs bytes.Buffer
+	previousOutput := log.Writer()
+	previousFlags := log.Flags()
+	log.SetOutput(&logs)
+	log.SetFlags(0)
+	t.Cleanup(func() {
+		log.SetOutput(previousOutput)
+		log.SetFlags(previousFlags)
+	})
+
+	r := gin.New()
+	r.Use(loggerMiddleware())
+	r.POST("/v1/responses", func(c *gin.Context) {
+		// 模拟 SSE 已提交 200 后才发现下游断开；真实 HTTP 状态无法回写，
+		// 但访问日志必须记录最终内部结果 499。
+		c.Status(http.StatusOK)
+		c.Set(proxy.AccessLogStatusContextKey, 499)
+	})
+
+	w := httptest.NewRecorder()
+	req := httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
+	r.ServeHTTP(w, req)
+
+	if w.Code != http.StatusOK {
+		t.Fatalf("wire status = %d, want %d", w.Code, http.StatusOK)
+	}
+	if got := logs.String(); !strings.Contains(got, "POST /v1/responses 499") {
+		t.Fatalf("access log must use stream outcome override: %s", got)
+	}
+}
diff --git a/proxy/downstream_keepalive.go b/proxy/downstream_keepalive.go
new file mode 100644
index 000000000..d181a456b
--- /dev/null
+++ b/proxy/downstream_keepalive.go
@@ -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
+		})
+	}
+}
diff --git a/proxy/downstream_keepalive_test.go b/proxy/downstream_keepalive_test.go
new file mode 100644
index 000000000..f1fe59572
--- /dev/null
+++ b/proxy/downstream_keepalive_test.go
@@ -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)
+	}
+}
diff --git a/proxy/handler.go b/proxy/handler.go
index 701f41708..ec9b57be6 100644
--- a/proxy/handler.go
+++ b/proxy/handler.go
@@ -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
+				})
+			}
 			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)
+		}
 		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 {
diff --git a/proxy/handler_test.go b/proxy/handler_test.go
index e63696eb8..647ddc6eb 100644
--- a/proxy/handler_test.go
+++ b/proxy/handler_test.go
@@ -1570,6 +1570,59 @@ func TestResponsesHTTPIngressFallsBackToHTTPWhenForcedWebsocketMessageTooBig(t *
 	}
 }
 
+func TestResponsesHTTPIngressKeepsDownstreamAliveDuringUpstreamSilence(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+
+	previousExec := WebsocketExecuteFunc
+	previousSettings := CurrentRuntimeSettings()
+	previousInterval := downstreamSSEKeepaliveInterval
+	t.Cleanup(func() {
+		WebsocketExecuteFunc = previousExec
+		ApplyRuntimeSettings(previousSettings)
+		downstreamSSEKeepaliveInterval = previousInterval
+	})
+
+	nextSettings := previousSettings
+	nextSettings.CodexForceWebsocket = true
+	nextSettings.CodexContinueThinking = false
+	ApplyRuntimeSettings(nextSettings)
+	downstreamSSEKeepaliveInterval = 5 * time.Millisecond
+
+	WebsocketExecuteFunc = func(ctx context.Context, account *auth.Account, requestBody []byte, sessionID string, proxyOverride string, apiKey string, deviceCfg *DeviceProfileConfig, headers http.Header, poolRouteKey string) (*http.Response, error) {
+		pr, pw := io.Pipe()
+		go func() {
+			_, _ = pw.Write([]byte(`data: {"type":"response.output_text.delta","delta":"started"}` + "\n\n"))
+			time.Sleep(30 * time.Millisecond)
+			_, _ = pw.Write([]byte(`data: {"type":"response.completed","response":{"usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}}` + "\n\n"))
+			_ = pw.Close()
+		}()
+		return &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: pr}, nil
+	}
+
+	store := auth.NewStore(nil, nil, &database.SystemSettings{MaxConcurrency: 1, TestConcurrency: 1, TestModel: "gpt-5.6-sol"})
+	store.AddAccount(&auth.Account{DBID: 1, AccessToken: "at-1", PlanType: "pro", AccountID: "acct-1"})
+	handler := NewHandler(store, nil, &config.Config{AllowAnonymousV1: true}, nil)
+
+	body := []byte(`{"model":"gpt-5.6-sol","input":"hello","stream":true}`)
+	req := httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body))
+	req.Header.Set("Content-Type", "application/json")
+	recorder := httptest.NewRecorder()
+	ctx, _ := gin.CreateTestContext(recorder)
+	ctx.Request = req
+
+	handler.Responses(ctx)
+
+	if recorder.Code != http.StatusOK {
+		t.Fatalf("status = %d, want 200; body=%s", recorder.Code, recorder.Body.String())
+	}
+	got := recorder.Body.String()
+	for _, want := range []string{`"delta":"started"`, downstreamSSEKeepaliveComment, `"type":"response.completed"`} {
+		if !strings.Contains(got, want) {
+			t.Fatalf("stream missing %q; body=%q", want, got)
+		}
+	}
+}
+
 func TestResponsesSkipsWebsocketWhenBodyReachesLearnedTooBigThreshold(t *testing.T) {
 	gin.SetMode(gin.TestMode)