Agent-System/internal/server/phase5_fault_injection_pg_test.go

package server_test

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

	"github.com/jackc/pgx/v5/pgxpool"

	"chat-agent-system/internal/bizdev"
	"chat-agent-system/internal/db"
	"chat-agent-system/internal/projectinstance"
	"chat-agent-system/internal/server"
	"chat-agent-system/internal/stream"
	"chat-agent-system/internal/taskrun"
	"chat-agent-system/internal/workflowdef"
)

// setupPGProjectReadyToAdvance 创建并激活官方 main 工作流、完成 intake 文档批准，返回可推进实例。
func setupPGProjectReadyToAdvance(t *testing.T, pool *pgxpool.Pool, handler http.Handler, srv *server.Server) (projectID, instanceID int64, creatorToken string, configRevision int) {
	t.Helper()
	ctx := context.Background()
	wfKey := workflowdef.OfficialKeyMain
	seedOfficialWorkflowManifest(t, pool, wfKey)

	creatorToken = loginTestUser(t, handler, "zhanggong", "123456")
	creatorID := fetchAuthUserID(t, handler, creatorToken)
	reviewerID := fetchAuthUserID(t, handler, loginTestUser(t, handler, "lizhuguan", "123456"))
	approverID := fetchAuthUserID(t, handler, loginTestUser(t, handler, "wangzong", "123456"))

	code := "p5-fault-" + time.Now().Format("150405.000000")
	err := pool.QueryRow(ctx, `
		INSERT INTO projects (code, name, status, creator_user_id, lifecycle_status, active_workflow_key, workflow_selection_version, workflow_plan_revision)
		VALUES ($1, $2, 'active', $3, 'planning', '', 1, 0) RETURNING id
	`, code, "Phase5 Fault Test "+code, creatorID).Scan(&projectID)
	if err != nil {
		t.Fatalf("insert project: %v", err)
	}
	ensureProjectMembers(t, pool, projectID, map[int64]string{
		creatorID:  "owner",
		reviewerID: "member",
		approverID: "member",
	})
	t.Cleanup(func() { cleanupProject(t, pool, projectID) })

	reqDefault := httptest.NewRequest(http.MethodGet, fmt.Sprintf("/api/v1/projects/%d/workflow-plan/default?workflow_key=%s&version=0", projectID, wfKey), nil)
	withBearerAuth(reqDefault, creatorToken)
	recDefault := httptest.NewRecorder()
	handler.ServeHTTP(recDefault, reqDefault)
	if recDefault.Code != http.StatusOK {
		t.Fatalf("default plan: %d %s", recDefault.Code, recDefault.Body.String())
	}
	var plan projectinstance.WorkflowPlan
	_ = json.Unmarshal(recDefault.Body.Bytes(), &plan)
	for i := range plan.Stages {
		for j := range plan.Stages[i].DocTasks {
			plan.Stages[i].DocTasks[j].Duties = projectinstance.DocDuty{
				AuthorUserID: creatorID, ReviewerUserID: reviewerID, ApproverUserID: approverID,
			}
		}
	}
	saveBody, _ := json.Marshal(projectinstance.SavePlanInput{Plan: plan, ExpectedPlanRevision: 0, ExpectedSelectionVersion: 1})
	reqSave := httptest.NewRequest(http.MethodPut, fmt.Sprintf("/api/v1/projects/%d/workflow-plan", projectID), bytes.NewReader(saveBody))
	withBearerAuth(reqSave, creatorToken)
	recSave := httptest.NewRecorder()
	handler.ServeHTTP(recSave, reqSave)
	if recSave.Code != http.StatusOK {
		t.Fatalf("save plan: %d %s", recSave.Code, recSave.Body.String())
	}

	reqPreview := httptest.NewRequest(http.MethodPost, fmt.Sprintf("/api/v1/projects/%d/activation-preview", projectID), nil)
	withBearerAuth(reqPreview, creatorToken)
	recPreview := httptest.NewRecorder()
	handler.ServeHTTP(recPreview, reqPreview)
	if recPreview.Code != http.StatusOK {
		t.Fatalf("activation preview: %d %s", recPreview.Code, recPreview.Body.String())
	}
	var preview projectinstance.ActivationPreview
	_ = json.Unmarshal(recPreview.Body.Bytes(), &preview)

	activateBody, _ := json.Marshal(projectinstance.ActivateInput{
		ConfirmationToken: preview.ConfirmationToken, ExpectedProjectVersion: preview.ExpectedProjectVersion,
		ExpectedPlanRevision: preview.ExpectedPlanRevision,
	})
	reqActivate := httptest.NewRequest(http.MethodPost, fmt.Sprintf("/api/v1/projects/%d/workflow-instances/activate", projectID), bytes.NewReader(activateBody))
	withBearerAuth(reqActivate, creatorToken)
	recActivate := httptest.NewRecorder()
	handler.ServeHTTP(recActivate, reqActivate)
	if recActivate.Code != http.StatusOK {
		t.Fatalf("activate: %d %s", recActivate.Code, recActivate.Body.String())
	}
	var activated projectinstance.ActivateResult
	_ = json.Unmarshal(recActivate.Body.Bytes(), &activated)
	instanceID = activated.Instance.ID
	configRevision = activated.Instance.ConfigRevision

	docID := fmt.Sprintf("doc_p5_fault_%d", projectID)
	docEngine := bizdev.Snapshot(srv.BizdevService()).Gov
	_, _ = docEngine.CreateOrGetDocument(ctx, docID, projectID, "bd_intake", "Intake", "V0.1", "", creatorID, "A", reviewerID, "R", approverID, "P")
	_ = docEngine.MarkDocumentContentUpdated(ctx, docID, "hash_p5_fault", false)
	_, _ = docEngine.SubmitForReview(ctx, docID, creatorID, "author", "Author")
	_, _ = docEngine.ReviewApprove(ctx, docID, reviewerID, "reviewer", "Reviewer", "ok")
	_, _ = docEngine.ApproverApprove(ctx, docID, approverID, "approver", "Approver", "ok")
	return projectID, instanceID, creatorToken, configRevision
}

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

	handler, srv := setupPGGovernanceServer(t, pool)
	_, instanceID, creatorToken, configRevision := setupPGProjectReadyToAdvance(t, pool, handler, srv)

	advanceBody, _ := json.Marshal(projectinstance.AdvanceStageInput{
		StageKey: "V0.1_triage", ExpectedConfigRevision: configRevision,
	})
	url := fmt.Sprintf("/api/v1/workflow-instances/%d/stages/V0.1_triage/advance", instanceID)

	var wg sync.WaitGroup
	results := make([]int, 2)
	for i := 0; i < 2; i++ {
		wg.Add(1)
		go func(idx int) {
			defer wg.Done()
			req := httptest.NewRequest(http.MethodPost, url, bytes.NewReader(advanceBody))
			withBearerAuth(req, creatorToken)
			rec := httptest.NewRecorder()
			handler.ServeHTTP(rec, req)
			results[idx] = rec.Code
		}(i)
	}
	wg.Wait()

	okCount, conflictCount := 0, 0
	for _, code := range results {
		switch code {
		case http.StatusOK:
			okCount++
		case http.StatusConflict, http.StatusBadRequest, http.StatusUnprocessableEntity:
			conflictCount++
		}
	}
	if okCount != 1 {
		t.Fatalf("expected exactly one 200 advance, got codes=%v", results)
	}
	if conflictCount < 1 {
		t.Fatalf("expected second advance to fail with conflict/gate error, got codes=%v", results)
	}
}

func TestPhase5_TaskRunHeartbeatSurvivesOrphanScan_PG(t *testing.T) {
	pool := openPGPoolWithMigrations(t)
	if pool == nil {
		return
	}
	defer pool.Close()
	if !taskrun.LeaseColumnsReady(context.Background(), pool) {
		t.Skip("lease columns not ready")
	}

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

	runID := "run_hb_survive_" + time.Now().Format("150405.000000")
	store := taskrun.NewPGStore(pool)
	done := make(chan struct{})
	srv.TaskRuns().SetExecutor(
		func(ctx context.Context, _ string, _ bizdev.TaskRunRequest, _ *stream.Streamer) (*bizdev.TaskRunResult, error) {
			<-done
			return &bizdev.TaskRunResult{Success: true, TaskKey: "intake"}, nil
		},
		nil, nil,
	)
	srv.TaskRuns().SetFinalizer(func(ctx context.Context, record *taskrun.RunRecord, result *bizdev.TaskRunResult) (string, error) {
		return "completed", nil
	})

	record := &taskrun.RunRecord{
		RunID: runID, SessionID: "sess_hb", UserID: userID, TaskKey: "intake", Status: "running",
	}
	_, err := srv.TaskRuns().CreateAndStart(context.Background(), runID, record, bizdev.TaskRunRequest{TaskKey: "intake"}, bizdev.StageRunRequest{}, "")
	if err != nil {
		t.Fatalf("CreateAndStart: %v", err)
	}
	t.Cleanup(func() { close(done) })

	time.Sleep(150 * time.Millisecond)
	if n, err := store.MarkOrphanedRunning(context.Background()); err != nil {
		t.Fatalf("orphan scan: %v", err)
	} else if n > 0 {
		rec, _ := store.GetRun(context.Background(), runID)
		if rec != nil && rec.Status == "interrupted" {
			t.Fatal("active heartbeat run should not be marked interrupted")
		}
	}

	rec, err := store.GetRun(context.Background(), runID)
	if err != nil {
		t.Fatalf("GetRun: %v", err)
	}
	if rec.Status != "running" {
		t.Fatalf("expected running with fresh lease, got %s", rec.Status)
	}
}

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

	ctx := context.Background()
	migDir := db.ResolveMigrationsDir()
	v1, err := db.CurrentVersion(ctx, pool)
	if err != nil {
		t.Fatalf("version: %v", err)
	}
	expected, _ := db.ExpectedMigrationCount(migDir)
	if int(v1) < expected {
		t.Fatalf("schema %d behind %d files", v1, expected)
	}

	// 模拟备份恢复后再次 Up（空操作应幂等）
	if err := db.RunMigrations(ctx, pool, migDir); err != nil {
		t.Fatalf("restore migration: %v", err)
	}
	v2, _ := db.CurrentVersion(ctx, pool)
	if v2 != v1 {
		t.Fatalf("migration version changed after idempotent restore: %d -> %d", v1, v2)
	}

	var taskRunsExists bool
	_ = pool.QueryRow(ctx, `
		SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = 'task_runs')
	`).Scan(&taskRunsExists)
	if !taskRunsExists {
		t.Fatal("task_runs missing after restore simulation")
	}
}