Agent-System/internal/server/phase5_finalize_fault_pg_test.go

package server_test

import (
	"context"
	"errors"
	"testing"
	"time"

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

func TestPhase5_FinalizeFaultInjection_NoDoneEvent(t *testing.T) {
	m := taskrun.NewManager(nil, nil)
	m.SetExecutor(
		func(_ context.Context, _ string, _ bizdev.TaskRunRequest, _ *stream.Streamer) (*bizdev.TaskRunResult, error) {
			return &bizdev.TaskRunResult{Success: true, TaskKey: "intake"}, nil
		},
		nil, nil,
	)
	m.SetFinalizer(func(_ context.Context, _ *taskrun.RunRecord, _ *bizdev.TaskRunResult) (string, error) {
		return "", errors.New("injected finalize db failure")
	})

	runID := "run_finalize_fault"
	record := &taskrun.RunRecord{RunID: runID, Status: "running", TaskKey: "intake"}
	_, _ = m.CreateAndStart(context.Background(), runID, record, bizdev.TaskRunRequest{TaskKey: "intake"}, bizdev.StageRunRequest{}, "")

	deadline := time.Now().Add(3 * time.Second)
	for time.Now().Before(deadline) {
		rec, err := m.GetRun(runID)
		if err == nil && rec.Status == "finalizing_failed" {
			return
		}
		time.Sleep(20 * time.Millisecond)
	}
	t.Fatal("expected finalizing_failed after injected finalize fault")
}

func TestPhase5_RollingRestartReplay_PG(t *testing.T) {
	pool := openPGPoolWithMigrations(t)
	if pool == nil {
		return
	}
	defer pool.Close()

	handler, srv := setupPGGovernanceServer(t, pool)
	token := loginTestUser(t, handler, "zhanggong", "123456")
	userID := fetchAuthUserID(t, handler, token)

	store := taskrun.NewPGStore(pool)
	runID := "run_rolling_" + time.Now().Format("150405.000000")
	srv.TaskRuns().SetStore(store)
	srv.TaskRuns().SetExecutor(
		func(_ context.Context, _ string, _ bizdev.TaskRunRequest, st *stream.Streamer) (*bizdev.TaskRunResult, error) {
			if st != nil {
				st.EmitPhase("roll", "step", "agent", 1, 1, 10)
			}
			return &bizdev.TaskRunResult{Success: true, TaskKey: "intake"}, nil
		},
		nil, nil,
	)
	srv.TaskRuns().SetFinalizer(func(_ context.Context, _ *taskrun.RunRecord, _ *bizdev.TaskRunResult) (string, error) {
		return "completed", nil
	})

	projectID := seedProjectForRun(t, pool, "run-roll", userID)
	record := &taskrun.RunRecord{
		RunID: runID, SessionID: "sess_roll", ProjectID: projectID, UserID: userID, TaskKey: "intake", Status: "running",
	}
	_, _ = srv.TaskRuns().CreateAndStart(context.Background(), runID, record, bizdev.TaskRunRequest{TaskKey: "intake"}, bizdev.StageRunRequest{}, "")
	t.Cleanup(func() {
		_, _ = pool.Exec(context.Background(), `DELETE FROM task_run_events WHERE run_id = $1`, runID)
		_, _ = pool.Exec(context.Background(), `DELETE FROM task_runs WHERE run_id = $1`, runID)
	})

	deadline := time.Now().Add(8 * time.Second)
	for time.Now().Before(deadline) {
		rec, err := store.GetRun(context.Background(), runID)
		if err == nil && rec.Status == "completed" {
			break
		}
		time.Sleep(50 * time.Millisecond)
	}

	// 模拟三次进程重启：新 Manager 仅挂载 Store
	for i := 0; i < 3; i++ {
		restarted := taskrun.NewManager(srv.ProjectWorkflowService(), bizdev.Snapshot(srv.BizdevService()).Ctx)
		restarted.SetStore(store)
		rec, err := restarted.GetRun(runID)
		if err != nil {
			t.Fatalf("restart %d GetRun: %v", i, err)
		}
		if rec.Status != "completed" {
			t.Fatalf("restart %d status=%s", i, rec.Status)
		}
		_, replay, err := restarted.Subscribe(runID, 0)
		if err != nil {
			t.Fatalf("restart %d Subscribe: %v", i, err)
		}
		if len(replay) == 0 {
			t.Fatalf("restart %d expected replay events", i)
		}
	}

	_ = handler
	_ = token
}