Agent-System/internal/server/sse_wave0_test.go

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.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"))
	}
}