Agent-System/internal/server/task_run_idempotent_sse_test.go

package server

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"chat-agent-system/internal/bizdev"
	"chat-agent-system/internal/docgov"
	"chat-agent-system/internal/stream"
)

func TestCreateTaskRun_SameTurnIdempotentReturnsExistingDespiteBusyLease(t *testing.T) {
	srv := newServerWithMemoryIAM(t, nil) // FIND-SEC-02: 业务行为用例走显式内存授权
	ctx := t.Context()
	token := loginTestAdmin(t, srv)
	userID := int64(1)

	// 慢执行器：保持 run 在内存中可被 FindByIdempotencyKey 命中
	block := make(chan struct{})
	var firstRunID string
	t.Cleanup(func() {
		close(block)
		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer cancel()
		if firstRunID != "" {
			_ = srv.taskRuns().WaitRunDone(ctx, firstRunID)
		} else if srv.taskRuns() != nil {
			_ = srv.taskRuns().WaitIdle(ctx)
		}
	})
	srv.taskRuns().SetExecutor(
		func(ctx context.Context, _ string, _ bizdev.TaskRunRequest, _ *stream.Streamer) (*bizdev.TaskRunResult, error) {
			select {
			case <-ctx.Done():
				return &bizdev.TaskRunResult{Success: false}, ctx.Err()
			case <-block:
				return &bizdev.TaskRunResult{Success: true}, nil
			}
		},
		nil,
		func(ctx context.Context, _ string, _ bizdev.StageRunRequest, _ *stream.Streamer) (*bizdev.TaskRunResult, error) {
			return &bizdev.TaskRunResult{Success: true}, nil
		},
	)

	docID := "doc_proj143_intake"
	if _, err := srv.gov().CreateOrGetDocument(
		ctx, docID, 143, "intake", "用户初始需求收集表", "V0.1_triage", "",
		userID, "编制人", userID, "审核人", userID, "批准人",
	); err != nil {
		t.Fatalf("create document: %v", err)
	}

	body1, _ := json.Marshal(createTaskRunRequest{
		TaskRunRequestDTO: TaskRunRequestDTO{TaskKey: "intake", ProjectID: 143, ProjectName: "幂等项目", DocID: docID},
		SessionID:         "sess-idem-busy",
		ClientTurnID:      "turn-same-1",
		UserMessage:       "请编制首轮摸底",
	})
	req1 := httptest.NewRequest(http.MethodPost, "/api/v1/tasks/runs", bytes.NewReader(body1))
	withBearerAuth(req1, token)
	w1 := httptest.NewRecorder()
	srv.Routes().ServeHTTP(w1, req1)
	if w1.Code != http.StatusAccepted {
		t.Fatalf("first create: %d %s", w1.Code, w1.Body.String())
	}
	var first createTaskRunResponse
	if err := json.Unmarshal(w1.Body.Bytes(), &first); err != nil {
		t.Fatalf("decode first: %v", err)
	}
	if first.RunID == "" {
		t.Fatal("missing run_id")
	}
	firstRunID = first.RunID

	hist1, err := srv.ctxMgr().GetChatHistory(ctx, "sess-idem-busy", 50, 0)
	if err != nil {
		t.Fatalf("history1: %v", err)
	}
	if len(hist1) != 2 {
		t.Fatalf("after first create history=%d want 2", len(hist1))
	}

	// 活跃小队租约：模拟长任务仍占文档（旧路径会 409）
	if _, err := srv.gov().AcquireLease(ctx, docgov.LeaseAcquireRequest{
		DocID: docID, HolderType: docgov.LeaseHolderSquad, HolderID: first.RunID, TTL: time.Minute,
	}); err != nil {
		t.Fatalf("acquire squad lease: %v", err)
	}
	t.Cleanup(func() {
		_, _ = srv.gov().ReleaseLease(context.Background(), docID, string(docgov.LeaseHolderSquad), first.RunID)
	})

	body2, _ := json.Marshal(createTaskRunRequest{
		TaskRunRequestDTO: TaskRunRequestDTO{TaskKey: "intake", ProjectID: 143, ProjectName: "幂等项目", DocID: docID},
		SessionID:         "sess-idem-busy",
		ClientTurnID:      "turn-same-1",
		UserMessage:       "请编制首轮摸底",
	})
	req2 := httptest.NewRequest(http.MethodPost, "/api/v1/tasks/runs", bytes.NewReader(body2))
	withBearerAuth(req2, token)
	w2 := httptest.NewRecorder()
	srv.Routes().ServeHTTP(w2, req2)
	if w2.Code != http.StatusAccepted {
		t.Fatalf("same-turn retry want 202, got %d %s", w2.Code, w2.Body.String())
	}
	var second createTaskRunResponse
	_ = json.Unmarshal(w2.Body.Bytes(), &second)
	if second.RunID != first.RunID {
		t.Fatalf("want same run %s, got %s", first.RunID, second.RunID)
	}

	hist2, err := srv.ctxMgr().GetChatHistory(ctx, "sess-idem-busy", 50, 0)
	if err != nil {
		t.Fatalf("history2: %v", err)
	}
	if len(hist2) != 2 {
		t.Fatalf("idempotent retry must not append chat; history=%d", len(hist2))
	}

	// 不同 turn：仍应 DOC_BUSY 409
	body3, _ := json.Marshal(createTaskRunRequest{
		TaskRunRequestDTO: TaskRunRequestDTO{TaskKey: "intake", ProjectID: 143, ProjectName: "幂等项目", DocID: docID},
		SessionID:         "sess-idem-busy",
		ClientTurnID:      "turn-other-2",
		UserMessage:       "另一轮摸底",
	})
	req3 := httptest.NewRequest(http.MethodPost, "/api/v1/tasks/runs", bytes.NewReader(body3))
	withBearerAuth(req3, token)
	w3 := httptest.NewRecorder()
	srv.Routes().ServeHTTP(w3, req3)
	if w3.Code != http.StatusConflict {
		t.Fatalf("other turn want 409, got %d %s", w3.Code, w3.Body.String())
	}
	var busy map[string]any
	_ = json.Unmarshal(w3.Body.Bytes(), &busy)
	if busy["error_code"] != "DOC_BUSY" {
		t.Fatalf("want DOC_BUSY, got %+v", busy)
	}
}

func TestSSEHeartbeat_ExtendsPastServerWriteTimeout(t *testing.T) {
	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)
	}
	srv := &http.Server{
		Handler:      mux,
		WriteTimeout: 200 * time.Millisecond, // 短于流时长；无 extend 会断
	}
	go func() { _ = srv.Serve(ln) }()
	t.Cleanup(func() { _ = srv.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)
	}
	body := string(raw)
	if !strings.Contains(body, "heartbeat") {
		t.Fatalf("expected heartbeats, body=%q", body)
	}
	if !strings.Contains(body, "event: done") {
		t.Fatalf("stream should survive WriteTimeout via deadline extend; body=%q", body)
	}
}