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 表示心跳已