Agent-System/internal/server/task_finalize_test.go

package server

import (
	"context"
	"os"
	"path/filepath"
	"testing"
	"time"

	"chat-agent-system/internal/bizdev"
	"chat-agent-system/internal/config"
	"chat-agent-system/internal/docgov"
	"chat-agent-system/internal/projectworkflow"
	"chat-agent-system/internal/squad"
	"chat-agent-system/internal/stream"
	"chat-agent-system/internal/taskrun"
)

func TestFinalizeTaskRun_ReplacesPlaceholderContent(t *testing.T) {
	cfg := &config.Config{Env: "test", DataDir: t.TempDir(), AgentBackend: "stub"}
	svc := bizdev.NewService(cfg)
	srv := NewServer(cfg, svc)

	workDir := t.TempDir()
	_ = os.WriteFile(filepath.Join(workDir, "output.md"), []byte("# 最终交付\n\n正文内容。"), 0o644)

	ctxMgr := svc.ContextManager()
	asstID, err := ctxMgr.RecordUserMessageWithID(context.Background(), "sess-finalize", 0, "assistant", "任务执行中…", "TRIGGER_TASK", "", map[string]any{
		"task_status": "running",
	})
	if err != nil {
		t.Fatalf("RecordUserMessageWithID: %v", err)
	}

	pwSvc := projectworkflow.NewService(projectworkflow.NewStore(nil), nil)
	_, _ = pwSvc.CreateWorkflowRun(context.Background(), &projectworkflow.WorkflowRun{
		RunID:     "run_finalize_test",
		ProjectID: 56,
		TaskKey:   "intake",
		Status:    "running",
	})

	srv.pwSvc = pwSvc
	srv.taskRuns = taskrun.NewManager(pwSvc, ctxMgr)
	srv.initTaskRunExecutors()

	record := &taskrun.RunRecord{
		RunID:          "run_finalize_test",
		SessionID:      "sess-finalize",
		ProjectID:      56,
		TaskKey:        "intake",
		AssistantMsgID: asstID,
	}

	status, finErr := srv.finalizeTaskRun(context.Background(), record, &bizdev.TaskRunResult{
		Success:       true,
		ArtifactTitle: "用户初始需求收集表",
		WorkDir:       workDir,
		TaskKey:       "intake",
		QualityReport: &squad.QualityReport{Passed: false, Score: 40, Violations: []string{"不做展开"}},
		DocRecord:     &docgov.DocumentRecord{DocID: "doc_proj56_intake", ContentSHA256: "abc"},
	})
	if finErr != nil {
		t.Fatalf("finalizeTaskRun: %v", finErr)
	}
	if status != "quality_failed" {
		t.Fatalf("status=%q want quality_failed", status)
	}

	msg, err := ctxMgr.GetMessageByID(context.Background(), asstID)
	if err != nil {
		t.Fatalf("GetMessageByID: %v", err)
	}
	if msg.Content == "任务执行中…" {
		t.Fatalf("content still placeholder: %q", msg.Content)
	}
	if msg.Metadata["task_status"] != "quality_failed" {
		t.Fatalf("metadata task_status=%v", msg.Metadata["task_status"])
	}
	artifact, ok := msg.Metadata["artifact"].(map[string]any)
	if !ok || artifact["doc_id"] != "doc_proj56_intake" {
		t.Fatalf("missing artifact doc_id in metadata: %#v", msg.Metadata["artifact"])
	}

	run, err := pwSvc.GetWorkflowRun(context.Background(), "run_finalize_test")
	if err != nil {
		t.Fatalf("GetWorkflowRun: %v", err)
	}
	if run.OutputDocID != "doc_proj56_intake" {
		t.Fatalf("OutputDocID=%q", run.OutputDocID)
	}
	if run.OutputSummary == "" {
		t.Fatal("OutputSummary empty")
	}
}

func TestAsyncTaskRun_FinalizesAfterWorkerCompletes(t *testing.T) {
	cfg := &config.Config{Env: "test", DataDir: t.TempDir(), AgentBackend: "stub"}
	svc := bizdev.NewService(cfg)
	pwSvc := projectworkflow.NewService(projectworkflow.NewStore(nil), nil)
	ctxMgr := svc.ContextManager()

	workDir := t.TempDir()
	_ = os.WriteFile(filepath.Join(workDir, "output.md"), []byte("# async finalize\n"), 0o644)

	srv := &Server{cfg: cfg, svc: svc, pwSvc: pwSvc, taskRuns: taskrun.NewManager(pwSvc, ctxMgr)}
	srv.taskRuns.SetExecutor(
		func(_ context.Context, _ string, _ bizdev.TaskRunRequest, _ *stream.Streamer) (*bizdev.TaskRunResult, error) {
			return &bizdev.TaskRunResult{
				Success:       true,
				ArtifactTitle: "用户初始需求收集表",
				WorkDir:       workDir,
				TaskKey:       "intake",
				QualityReport: &squad.QualityReport{Passed: true, Score: 95},
				DocRecord:     &docgov.DocumentRecord{DocID: "doc_proj1_intake"},
			}, nil
		},
		nil,
		nil,
	)
	srv.taskRuns.SetFinalizer(func(ctx context.Context, record *taskrun.RunRecord, result *bizdev.TaskRunResult) (string, error) {
		return srv.finalizeTaskRun(ctx, record, result)
	})

	asstID, _ := ctxMgr.RecordUserMessageWithID(context.Background(), "sess-async", 0, "assistant", "任务执行中…", "TRIGGER_TASK", "", nil)
	runID := "run_async_finalize"
	_, _ = pwSvc.CreateWorkflowRun(context.Background(), &projectworkflow.WorkflowRun{
		RunID: runID, ProjectID: 1, TaskKey: "intake", Status: "running",
	})

	record := &taskrun.RunRecord{
		RunID: runID, SessionID: "sess-async", ProjectID: 1, TaskKey: "intake", AssistantMsgID: asstID,
	}
	_, err := srv.taskRuns.CreateAndStart(context.Background(), runID, record, bizdev.TaskRunRequest{TaskKey: "intake", ProjectID: 1}, bizdev.StageRunRequest{}, "")
	if err != nil {
		t.Fatalf("CreateAndStart: %v", err)
	}

	deadline := time.Now().Add(5 * time.Second)
	for time.Now().Before(deadline) {
		rec, getErr := srv.taskRuns.GetRun(runID)
		if getErr == nil && rec.Status != "running" {
			msg, _ := ctxMgr.GetMessageByID(context.Background(), asstID)
			if msg != nil && msg.Content != "任务执行中…" {
				return
			}
			t.Fatalf("run finished with status=%s but content=%q", rec.Status, msg.Content)
		}
		time.Sleep(50 * time.Millisecond)
	}
	t.Fatal("worker did not finish")
}