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