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 (failInsertStore) ListActiveBySession(context.Context, string) ([]*taskrun.RunRecord, error) { return nil, nil } func (failInsertStore) LatestBySession(context.Context, string) (*taskrun.RunRecord, error) { return nil, 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.setOverrideTaskRuns(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")) } }