From 987bd0a03cd5ff188f9b4692f821c54daf89b048 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=85=AC=E6=98=8E?= <83812544+Ed1s0nZ@users.noreply.github.com> Date: Sat, 11 Jul 2026 12:10:57 +0800 Subject: [PATCH] Add files via upload --- internal/handler/agent.go | 5 +-- internal/handler/eino_single_agent.go | 5 +-- internal/handler/multi_agent.go | 5 +-- internal/handler/sse_keepalive.go | 62 +++++++++++++++++++------- internal/handler/sse_keepalive_test.go | 61 +++++++++++++++++++++++++ 5 files changed, 113 insertions(+), 25 deletions(-) create mode 100644 internal/handler/sse_keepalive_test.go diff --git a/internal/handler/agent.go b/internal/handler/agent.go index 640bfdd9..d4afd4ed 100644 --- a/internal/handler/agent.go +++ b/internal/handler/agent.go @@ -1611,9 +1611,8 @@ func (h *AgentHandler) SubscribeAgentTaskEvents(c *gin.Context) { flusher, _ := c.Writer.(http.Flusher) ctx := c.Request.Context() var writeMu sync.Mutex - stopKeepalive := make(chan struct{}) - go sseKeepalive(c, stopKeepalive, &writeMu) - defer close(stopKeepalive) + stopKeepalive := runSSEKeepalive(c, &writeMu) + defer stopKeepalive() for { select { diff --git a/internal/handler/eino_single_agent.go b/internal/handler/eino_single_agent.go index d9faf2be..074e6cbd 100644 --- a/internal/handler/eino_single_agent.go +++ b/internal/handler/eino_single_agent.go @@ -139,9 +139,8 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) { "conversationId": conversationID, }) - stopKeepalive := make(chan struct{}) - go sseKeepalive(c, stopKeepalive, &sseWriteMu) - defer close(stopKeepalive) + stopKeepalive := runSSEKeepalive(c, &sseWriteMu) + defer stopKeepalive() if h.config == nil { taskStatus = "failed" diff --git a/internal/handler/multi_agent.go b/internal/handler/multi_agent.go index 0760bcac..03e6b4a8 100644 --- a/internal/handler/multi_agent.go +++ b/internal/handler/multi_agent.go @@ -156,9 +156,8 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) { "conversationId": conversationID, }) - stopKeepalive := make(chan struct{}) - go sseKeepalive(c, stopKeepalive, &sseWriteMu) - defer close(stopKeepalive) + stopKeepalive := runSSEKeepalive(c, &sseWriteMu) + defer stopKeepalive() var result *multiagent.RunResult var runErr error diff --git a/internal/handler/sse_keepalive.go b/internal/handler/sse_keepalive.go index ae750ecd..dac366d9 100644 --- a/internal/handler/sse_keepalive.go +++ b/internal/handler/sse_keepalive.go @@ -1,6 +1,7 @@ package handler import ( + "context" "fmt" "net/http" "sync" @@ -13,33 +14,51 @@ import ( // some proxies that treat connections as idle; 10s is a reasonable balance with traffic. const sseKeepaliveInterval = 10 * time.Second -// sseKeepalive sends periodic SSE traffic so proxies (e.g. nginx proxy_read_timeout), NATs, +// runSSEKeepalive starts periodic SSE heartbeats in a background goroutine. +// The returned stop function must be deferred (or called) before the handler returns so the +// goroutine exits before Gin finalizes the ResponseWriter (avoids "Write called after Handler finished"). +// +// writeMu must be the same mutex used by the handler's event writes for this request: concurrent +// writes to http.ResponseWriter break chunked transfer encoding (browser: net::ERR_INVALID_CHUNKED_ENCODING). +func runSSEKeepalive(c *gin.Context, writeMu *sync.Mutex) func() { + if writeMu == nil { + return func() {} + } + stop := make(chan struct{}) + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + sseKeepaliveLoop(c, stop, writeMu) + }() + var once sync.Once + return func() { + once.Do(func() { + close(stop) + wg.Wait() + }) + } +} + +// sseKeepaliveLoop sends periodic SSE traffic so proxies (e.g. nginx proxy_read_timeout), NATs, // and load balancers do not close long-running streams. Some intermediaries ignore comment-only // lines, so we send both a comment and a minimal data frame (type heartbeat) per tick. -// -// writeMu must be the same mutex used by sendEvent for this request: concurrent writes to -// http.ResponseWriter break chunked transfer encoding (browser: net::ERR_INVALID_CHUNKED_ENCODING). -func sseKeepalive(c *gin.Context, stop <-chan struct{}, writeMu *sync.Mutex) { - if writeMu == nil { - return - } +func sseKeepaliveLoop(c *gin.Context, stop <-chan struct{}, writeMu *sync.Mutex) { ticker := time.NewTicker(sseKeepaliveInterval) defer ticker.Stop() + ctx := c.Request.Context() for { select { case <-stop: return - case <-c.Request.Context().Done(): + case <-ctx.Done(): return case <-ticker.C: - select { - case <-stop: - return - case <-c.Request.Context().Done(): - return - default: - } writeMu.Lock() + if sseShuttingDown(stop, ctx) { + writeMu.Unlock() + return + } if _, err := fmt.Fprintf(c.Writer, ": keepalive\n\n"); err != nil { writeMu.Unlock() return @@ -56,3 +75,14 @@ func sseKeepalive(c *gin.Context, stop <-chan struct{}, writeMu *sync.Mutex) { } } } + +func sseShuttingDown(stop <-chan struct{}, ctx context.Context) bool { + select { + case <-stop: + return true + case <-ctx.Done(): + return true + default: + return false + } +} diff --git a/internal/handler/sse_keepalive_test.go b/internal/handler/sse_keepalive_test.go new file mode 100644 index 00000000..a468708d --- /dev/null +++ b/internal/handler/sse_keepalive_test.go @@ -0,0 +1,61 @@ +package handler + +import ( + "context" + "net/http/httptest" + "sync" + "testing" + "time" + + "github.com/gin-gonic/gin" +) + +func TestRunSSEKeepaliveStopsBeforeHandlerReturns(t *testing.T) { + gin.SetMode(gin.TestMode) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest("GET", "/events", nil) + + var writeMu sync.Mutex + stop := runSSEKeepalive(c, &writeMu) + stop() + + // A second stop must be safe (channel already closed, goroutine already exited). + stop() +} + +func TestRunSSEKeepaliveExitsOnClientDisconnect(t *testing.T) { + gin.SetMode(gin.TestMode) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + ctx, cancel := context.WithCancel(context.Background()) + c.Request = httptest.NewRequest("GET", "/events", nil).WithContext(ctx) + + var writeMu sync.Mutex + stop := runSSEKeepalive(c, &writeMu) + cancel() + + done := make(chan struct{}) + go func() { + stop() + close(done) + }() + + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("keepalive stop did not complete after client disconnect") + } +} + +func TestRunSSEKeepaliveNilMutexIsNoop(t *testing.T) { + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest("GET", "/events", nil) + + stop := runSSEKeepalive(c, nil) + stop() +}