fix(server): extend SSE write deadline for chat and run/stream

Wire heartbeat + SetWriteDeadline into handleChatStream and handleRunTaskStream; reject malformed run/stream JSON with 400; move CreateAndStart/Subscribe before SSE Flush so failures stay JSON errors. Co-authored-by: Cursor <cursoragent@cursor.com>
2026-09-17 16:42:56 +00:00
parent 64efaa32c8
commit 54244ed97a
2 changed files with 248 additions and 12 deletions
@@ -204,6 +204,7 @@ func (s *Server) handleChatStream(w http.ResponseWriter, r *http.Request) {
 	s.writeCORS(w, r)
 	w.Header().Set("X-Accel-Buffering", "no") // 禁用 Nginx 等反代缓冲
 	flusher.Flush()
+	extendSSEWriteDeadline(w)

 	ctx := stream.WithRequestID(r.Context(), requestID)

@@ -228,14 +229,37 @@ func (s *Server) handleChatStream(w http.ResponseWriter, r *http.Request) {
 		}
 	}

-	// 3. 构建线程安全的 SSE 监听分发闭包
+	// 3. 构建线程安全的 SSE 监听分发闭包；心跳与业务写共用 writeMu，避免与 WriteTimeout 竞态
 	var writeMu sync.Mutex
 	listener := func(ev stream.StreamEvent) {
 		writeMu.Lock()
 		defer writeMu.Unlock()
+		extendSSEWriteDeadline(w)
 		writeSSE(w, flusher, ev)
 	}

+	heartbeatDone := make(chan struct{})
+	go func() {
+		ticker := time.NewTicker(sseHeartbeatInterval)
+		defer ticker.Stop()
+		for {
+			select {
+			case <-heartbeatDone:
+				return
+			case <-r.Context().Done():
+				return
+			case <-ticker.C:
+				writeMu.Lock()
+				err := writeSSEHeartbeat(w, flusher)
+				writeMu.Unlock()
+				if err != nil {
+					return
+				}
+			}
+		}
+	}()
+	defer close(heartbeatDone)
+
 	clientTurnID := req.ClientTurnID
 	if clientTurnID == "" {
 		clientTurnID = r.Header.Get("Idempotency-Key")
@@ -312,7 +336,10 @@ func (s *Server) handleRunTaskStream(w http.ResponseWriter, r *http.Request) {
 	var req taskStreamRequest
 	includeDetails := false
 	if r.Method == http.MethodPost {
-		_ = json.NewDecoder(r.Body).Decode(&req)
+		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+			writeSafeError(w, r, http.StatusBadRequest, "invalid request body: "+err.Error())
+			return
+		}
 	} else {
 		q := r.URL.Query()
 		req.TaskKey = q.Get("task_key")
@@ -349,22 +376,15 @@ func (s *Server) handleRunTaskStream(w http.ResponseWriter, r *http.Request) {
 		}
 	}

-	// 请求追踪 ID：与 chat stream 对齐，注入 ctx 并回写响应头
-	requestID := fmt.Sprintf("req_%d", time.Now().UnixNano())
-	w.Header().Set("X-Request-Id", requestID)
-	w.Header().Set("Content-Type", "text/event-stream")
-	w.Header().Set("Cache-Control", "no-cache")
-	w.Header().Set("Connection", "keep-alive")
-	s.writeCORS(w, r)
-	w.Header().Set("X-Accel-Buffering", "no") // 禁用 Nginx 等反代缓冲
-	flusher.Flush()
-
 	if req.TaskKey == "" && req.Stage == "" {
 		req.TaskKey = "intake"
 	}
 	req.UserID = user.ID
 	req.WorkDir = ""

+	// 请求追踪 ID：与 chat stream 对齐；CreateAndStart/Subscribe 必须在 Flush 之前，
+	// 失败时返回普通 JSON 错误，避免污染已声明的 text/event-stream 连接。
+	requestID := fmt.Sprintf("req_%d", time.Now().UnixNano())
 	runID := fmt.Sprintf("run_%d", time.Now().UnixNano())
 	record := &taskrun.RunRecord{
 		RunID:        runID,
@@ -394,9 +414,19 @@ func (s *Server) handleRunTaskStream(w http.ResponseWriter, r *http.Request) {
 	}
 	defer s.taskRuns.Unsubscribe(created.RunID, ch)

+	w.Header().Set("X-Request-Id", requestID)
+	w.Header().Set("Content-Type", "text/event-stream")
+	w.Header().Set("Cache-Control", "no-cache")
+	w.Header().Set("Connection", "keep-alive")
+	s.writeCORS(w, r)
+	w.Header().Set("X-Accel-Buffering", "no") // 禁用 Nginx 等反代缓冲
+	flusher.Flush()
+	extendSSEWriteDeadline(w)
+
 	writeStored := func(stored taskrun.StoredEvent) {
 		_ = (transport.Encoder{W: w}).WriteWithID(stored.Seq, stored.Event)
 		flusher.Flush()
+		extendSSEWriteDeadline(w)
 	}

 	for _, stored := range replay {
@@ -406,11 +436,18 @@ func (s *Server) handleRunTaskStream(w http.ResponseWriter, r *http.Request) {
 		}
 	}

+	heartbeat := time.NewTicker(sseHeartbeatInterval)
+	defer heartbeat.Stop()
+
 	for {
 		select {
 		case <-r.Context().Done():
 			// SSE 断线仅停止推送，不取消后台 worker
 			return
+		case <-heartbeat.C:
+			if err := writeSSEHeartbeat(w, flusher); err != nil {
+				return
+			}
 		case stored, ok := <-ch:
 			if !ok {
 				return
@@ -0,0 +1,199 @@
+package server
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"net"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+	"time"
+
+	"chat-agent-system/internal/bizdev"
+	"chat-agent-system/internal/config"
+	"chat-agent-system/internal/stream"
+	"chat-agent-system/internal/taskrun"
+)
+
+// failInsertStore：Available=true 但 InsertRun 恒失败，用于验证 Flush 前 CreateAndStart 失败路径。
+type failInsertStore struct{}
+
+func (failInsertStore) Available() bool { return true }
+func (failInsertStore) InsertRun(context.Context, *taskrun.RunRecord) error {
+	return errors.New("injected insert failure")
+}
+func (failInsertStore) UpdateRun(context.Context, *taskrun.RunRecord) error { return nil }
+func (failInsertStore) GetByIdempotencyKey(context.Context, string) (*taskrun.RunRecord, error) {
+	return nil, nil
+}
+func (failInsertStore) AppendEvent(context.Context, string, int64, stream.StreamEvent) error {
+	return nil
+}
+func (failInsertStore) GetRun(context.Context, string) (*taskrun.RunRecord, error) {
+	return nil, errors.New("not found")
+}
+func (failInsertStore) ListEvents(context.Context, string, int64) ([]taskrun.StoredEvent, error) {
+	return nil, nil
+}
+func (failInsertStore) AcquireLease(context.Context, string, string, time.Duration) error {
+	return nil
+}
+func (failInsertStore) HeartbeatLease(context.Context, string, string, time.Duration) error {
+	return nil
+}
+func (failInsertStore) ReleaseLease(context.Context, string, string) error { return nil }
+func (failInsertStore) ResetForResume(context.Context, string, int) error  { return nil }
+func (failInsertStore) MarkOrphanedRunning(context.Context) (int, error)   { return 0, nil }
+
+func TestRunTaskStream_InvalidJSON_Returns400(t *testing.T) {
+	srv := newTestServer(t)
+	token := loginTestAdmin(t, srv)
+
+	req := httptest.NewRequest(http.MethodPost, "/api/v1/tasks/run/stream", bytes.NewReader([]byte(`{not-json`)))
+	withBearerAuth(req, token)
+	req.Header.Set("Content-Type", "application/json")
+	w := httptest.NewRecorder()
+	srv.Routes().ServeHTTP(w, req)
+
+	if w.Code != http.StatusBadRequest {
+		t.Fatalf("want 400 for malformed body, got %d body=%s", w.Code, w.Body.String())
+	}
+	ct := w.Header().Get("Content-Type")
+	if strings.Contains(ct, "text/event-stream") {
+		t.Fatalf("malformed body must not open SSE; Content-Type=%q", ct)
+	}
+}
+
+func TestRunTaskStream_CreateAndStartFails_BeforeSSEFlush(t *testing.T) {
+	srv := newTestServer(t)
+	token := loginTestAdmin(t, srv)
+
+	mgr := taskrun.NewManager(nil, nil)
+	mgr.SetStore(failInsertStore{})
+	srv.taskRuns = mgr
+
+	body, _ := json.Marshal(map[string]any{
+		"task_key":     "intake",
+		"project_name": "wave0-fail-start",
+	})
+	req := httptest.NewRequest(http.MethodPost, "/api/v1/tasks/run/stream", bytes.NewReader(body))
+	withBearerAuth(req, token)
+	req.Header.Set("Content-Type", "application/json")
+	w := httptest.NewRecorder()
+	srv.Routes().ServeHTTP(w, req)
+
+	if w.Code < 400 {
+		t.Fatalf("want error status before SSE flush, got %d body=%s", w.Code, w.Body.String())
+	}
+	ct := w.Header().Get("Content-Type")
+	if strings.Contains(ct, "text/event-stream") {
+		t.Fatalf("CreateAndStart failure must not declare SSE; Content-Type=%q body=%s", ct, w.Body.String())
+	}
+	if !strings.Contains(w.Body.String(), "start task run failed") && !strings.Contains(w.Body.String(), "injected insert") {
+		t.Fatalf("expected start failure message, got %s", w.Body.String())
+	}
+}
+
+func TestSSEHeartbeat_HelpersSurviveShortWriteTimeout(t *testing.T) {
+	// 复用既有心跳机制：短 WriteTimeout 下流时长 > timeout 仍可读完（已有机制的回归钉）。
+	mux := http.NewServeMux()
+	mux.HandleFunc("/sse", func(w http.ResponseWriter, r *http.Request) {
+		flusher, ok := w.(http.Flusher)
+		if !ok {
+			http.Error(w, "no flush", 500)
+			return
+		}
+		w.Header().Set("Content-Type", "text/event-stream")
+		flusher.Flush()
+		extendSSEWriteDeadline(w)
+		tick := time.NewTicker(60 * time.Millisecond)
+		defer tick.Stop()
+		end := time.After(700 * time.Millisecond)
+		for {
+			select {
+			case <-end:
+				_, _ = fmt.Fprintf(w, "event: done\ndata: {\"type\":\"done\"}\n\n")
+				flusher.Flush()
+				return
+			case <-tick.C:
+				if err := writeSSEHeartbeat(w, flusher); err != nil {
+					return
+				}
+			case <-r.Context().Done():
+				return
+			}
+		}
+	})
+
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	httpSrv := &http.Server{
+		Handler:      mux,
+		WriteTimeout: 200 * time.Millisecond,
+	}
+	go func() { _ = httpSrv.Serve(ln) }()
+	t.Cleanup(func() { _ = httpSrv.Close() })
+
+	resp, err := http.Get("http://" + ln.Addr().String() + "/sse")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer resp.Body.Close()
+	raw, err := io.ReadAll(resp.Body)
+	if err != nil {
+		t.Fatalf("read body: %v", err)
+	}
+	if !strings.Contains(string(raw), "event: done") {
+		t.Fatalf("expected done event after heartbeats; body=%q", string(raw))
+	}
+}
+
+func TestRunTaskStream_ValidBody_OpensSSEWithHeartbeatPath(t *testing.T) {
+	cfg := &config.Config{
+		Env:           "test",
+		AgentBackend:  "stub",
+		DataDir:       t.TempDir(),
+		KbDbPort:      5432,
+		AdminUsername: "admin",
+		AdminPassword: "admin-test-pass",
+	}
+	svc := bizdev.MustNewService(t, cfg)
+	srv := NewServer(cfg, svc)
+	t.Cleanup(func() {
+		ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
+		defer cancel()
+		_ = srv.Shutdown(ctx)
+	})
+
+	ts := httptest.NewServer(srv.Routes())
+	defer ts.Close()
+	token := testBearerToken(t, ts.URL, cfg)
+
+	body, _ := json.Marshal(map[string]any{
+		"task_key":     "intake",
+		"project_name": "wave0-sse-ok",
+		"customer":     "北京水木东方",
+	})
+	req, _ := http.NewRequest(http.MethodPost, ts.URL+"/api/v1/tasks/run/stream", bytes.NewReader(body))
+	req.Header.Set("Content-Type", "application/json")
+	req.Header.Set("Authorization", "Bearer "+token)
+	resp, err := http.DefaultClient.Do(req)
+	if err != nil {
+		t.Fatalf("post: %v", err)
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode != http.StatusOK {
+		raw, _ := io.ReadAll(resp.Body)
+		t.Fatalf("want 200 SSE, got %d %s", resp.StatusCode, string(raw))
+	}
+	if !strings.Contains(resp.Header.Get("Content-Type"), "text/event-stream") {
+		t.Fatalf("Content-Type=%q", resp.Header.Get("Content-Type"))
+	}
+}