Add files via upload

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{})
+	stopKeepalive := runSSEKeepalive(c, &writeMu)
-	go sseKeepalive(c, stopKeepalive, &writeMu)
+	defer stopKeepalive()
-	defer close(stopKeepalive)
 
 	for {
 		select {

internal/handler/eino_single_agent.go

@@ -139,9 +139,8 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
 		"conversationId": conversationID,
 	})
 
-	stopKeepalive := make(chan struct{})
+	stopKeepalive := runSSEKeepalive(c, &sseWriteMu)
-	go sseKeepalive(c, stopKeepalive, &sseWriteMu)
+	defer stopKeepalive()
-	defer close(stopKeepalive)
 
 	if h.config == nil {
 		taskStatus = "failed"

internal/handler/multi_agent.go

@@ -156,9 +156,8 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
 		"conversationId": conversationID,
 	})
 
-	stopKeepalive := make(chan struct{})
+	stopKeepalive := runSSEKeepalive(c, &sseWriteMu)
-	go sseKeepalive(c, stopKeepalive, &sseWriteMu)
+	defer stopKeepalive()
-	defer close(stopKeepalive)
 
 	var result *multiagent.RunResult
 	var runErr error

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.
-//
+func sseKeepaliveLoop(c *gin.Context, stop <-chan struct{}, writeMu *sync.Mutex) {
-// 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
-	}
 	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
+	}
+}

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()
+}