From 54244ed97a2731faa759fd6ed3719ba7e7040993 Mon Sep 17 00:00:00 2001 From: weichen Date: Thu, 17 Sep 2026 16:42:56 +0000 Subject: [PATCH] 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 --- internal/server/sse.go | 61 +++++++-- internal/server/sse_wave0_test.go | 199 ++++++++++++++++++++++++++++++ 2 files changed, 248 insertions(+), 12 deletions(-) create mode 100644 internal/server/sse_wave0_test.go diff --git a/internal/server/sse.go b/internal/server/sse.go index 16d5daa8..c51730af 100644 --- a/internal/server/sse.go +++ b/internal/server/sse.go @@ -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 diff --git a/internal/server/sse_wave0_test.go b/internal/server/sse_wave0_test.go new file mode 100644 index 00000000..ff59e36b --- /dev/null +++ b/internal/server/sse_wave0_test.go @@ -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")) + } +}