From 06af8115f7fda82c70075a675bb581a25c3ed4d7 Mon Sep 17 00:00:00 2001
From: jjaw <jjaw@jjaw.cn>
Date: Sun, 12 Jul 2026 04:53:41 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=20compact=20=E5=BF=83?=
 =?UTF-8?q?=E8=B7=B3=20writer=20=E7=94=9F=E5=91=BD=E5=91=A8=E6=9C=9F?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../handler/ops_capture_writer_nil_test.go    | 31 +++++++++++++++++++
 .../service/openai_compact_sse_keepalive.go   | 15 +++++++--
 2 files changed, 43 insertions(+), 3 deletions(-)

diff --git a/backend/internal/handler/ops_capture_writer_nil_test.go b/backend/internal/handler/ops_capture_writer_nil_test.go
index 4e96333f9d..88aa7c043f 100644
--- a/backend/internal/handler/ops_capture_writer_nil_test.go
+++ b/backend/internal/handler/ops_capture_writer_nil_test.go
@@ -1,9 +1,15 @@
 package handler
 
 import (
+	"net/http"
+	"net/http/httptest"
 	"testing"
+	"time"
 
+	"github.com/Wei-Shaw/sub2api/internal/service"
+	"github.com/gin-gonic/gin"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestOpsCaptureWriter_NilInnerWriter_NoPanic(t *testing.T) {
@@ -57,3 +63,28 @@ func TestOpsCaptureWriter_NilInnerWriter_NoPanic(t *testing.T) {
 		assert.Nil(t, p)
 	})
 }
+
+func TestOpsCaptureWriter_CompactKeepaliveRestoresOriginalWriter(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+	router := gin.New()
+	outerStatus := -1
+	router.Use(func(c *gin.Context) {
+		c.Next()
+		outerStatus = c.Writer.Status()
+	})
+	router.Use(OpsErrorLoggerMiddleware(nil))
+	router.GET("/compact", func(c *gin.Context) {
+		service.MarkOpenAICompactClientStream(c)
+		stop := service.StartOpenAICompactSSEKeepalive(c, time.Hour)
+		defer stop()
+		c.Status(http.StatusOK)
+	})
+
+	recorder := httptest.NewRecorder()
+	request := httptest.NewRequest(http.MethodGet, "/compact", nil)
+	require.NotPanics(t, func() {
+		router.ServeHTTP(recorder, request)
+	})
+	require.Equal(t, http.StatusOK, outerStatus)
+	require.Equal(t, http.StatusOK, recorder.Code)
+}
diff --git a/backend/internal/service/openai_compact_sse_keepalive.go b/backend/internal/service/openai_compact_sse_keepalive.go
index 70ef3fc01a..776fde6962 100644
--- a/backend/internal/service/openai_compact_sse_keepalive.go
+++ b/backend/internal/service/openai_compact_sse_keepalive.go
@@ -43,12 +43,14 @@ func StartOpenAICompactSSEKeepalive(c *gin.Context, interval time.Duration) func
 	if c == nil || c.Writer == nil || interval <= 0 || !openAICompactClientWantsStream(c) {
 		return func() {}
 	}
+	originalWriter := c.Writer
 	k := &openAICompactSSEKeepalive{
-		writer: c.Writer,
+		writer: originalWriter,
 		stop:   make(chan struct{}),
 	}
 	c.Set(openAICompactSSEKeepaliveKey, k)
-	c.Writer = &openAICompactKeepaliveWriter{ResponseWriter: c.Writer, k: k}
+	wrappedWriter := &openAICompactKeepaliveWriter{ResponseWriter: originalWriter, k: k}
+	c.Writer = wrappedWriter
 
 	var reqDone <-chan struct{}
 	if c.Request != nil {
@@ -71,7 +73,14 @@ func StartOpenAICompactSSEKeepalive(c *gin.Context, interval time.Duration) func
 			timer.Reset(interval)
 		}
 	}()
-	return k.Stop
+	return func() {
+		k.Stop()
+		// Do not leave a pooled middleware writer reachable through the compact
+		// wrapper after the request finishes.
+		if current, ok := c.Writer.(*openAICompactKeepaliveWriter); ok && current == wrappedWriter {
+			c.Writer = originalWriter
+		}
+	}
 }
 
 // beat 在锁内提交（首次）响应头并写出一条 SSE 注释行；返回 false 表示心跳已