Agent-System/internal/server/multi_intent_u8_piagent_pg_test.go

package server_test

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

	"chat-agent-system/internal/app"
	"chat-agent-system/internal/app/intentflow"
	"chat-agent-system/internal/app/uow"
	"chat-agent-system/internal/bizdev"
	"chat-agent-system/internal/events"
	"chat-agent-system/internal/projectinstance"
	"chat-agent-system/internal/projectworkflow"
	"chat-agent-system/internal/taskrun"
	"chat-agent-system/internal/workflowdef"
)

// TestMultiIntentU8_RealPI_ConfirmGenerate_PG closes the U8 real-PI evidence gap:
// confirm-before-start, live pi-agent generate via TaskRunAdapter, idempotent replay,
// revoke, and PG plan recovery after a fresh Plans wiring (restart simulation).
func TestMultiIntentU8_RealPI_ConfirmGenerate_PG(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping slow U8 real-PI PG E2E in short mode")
	}
	if !piAgentRunnerHealthy() {
		t.Skip("pi-agent runner not available at http://127.0.0.1:18865/health")
	}
	if os.Getenv("PI_AGENT_RUNNER_URL") == "" {
		t.Setenv("PI_AGENT_RUNNER_URL", "http://127.0.0.1:18865")
	}

	pool := openPGPoolWithMigrations(t)
	if pool == nil {
		return
	}
	defer pool.Close()

	ctx := context.Background()
	handler, srv := setupPGPiAgentServer(t, pool)

	seedOfficialWorkflowManifest(t, pool, workflowdef.OfficialKeyMain)

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

	createBody, _ := json.Marshal(projectworkflow.CreateProjectInput{
		Name:         fmt.Sprintf("U8 Wave2 PI %d", time.Now().UnixNano()%1_000_000),
		DeviceType:   "Ultrasonic Device",
		IntendedUse:  "Tissue resection",
		TargetMarket: "NMPA Class III",
	})
	reqCreate := httptest.NewRequest(http.MethodPost, "/api/v1/projects", bytes.NewReader(createBody))
	withBearerAuth(reqCreate, creatorToken)
	recCreate := httptest.NewRecorder()
	handler.ServeHTTP(recCreate, reqCreate)
	if recCreate.Code != http.StatusCreated {
		t.Fatalf("create project: %d %s", recCreate.Code, recCreate.Body.String())
	}
	var proj projectworkflow.Project
	if err := json.Unmarshal(recCreate.Body.Bytes(), &proj); err != nil {
		t.Fatalf("decode project: %v", err)
	}
	t.Cleanup(func() { cleanupProject(t, pool, proj.ID) })
	projectID := proj.ID

	reqDefault := httptest.NewRequest(http.MethodGet, fmt.Sprintf("/api/v1/projects/%d/workflow-plan/default?workflow_key=%s&version=0", projectID, workflowdef.OfficialKeyMain), 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 wfPlan projectinstance.WorkflowPlan
	_ = json.Unmarshal(recDefault.Body.Bytes(), &wfPlan)
	ensureProjectMembers(t, pool, projectID, map[int64]string{
		creatorID:  "owner",
		reviewerID: "member",
		approverID: "member",
	})
	for i := range wfPlan.Stages {
		for j := range wfPlan.Stages[i].DocTasks {
			wfPlan.Stages[i].DocTasks[j].Duties = projectinstance.DocDuty{
				AuthorUserID: creatorID, ReviewerUserID: reviewerID, ApproverUserID: approverID,
			}
		}
	}
	saveBody, _ := json.Marshal(projectinstance.SavePlanInput{Plan: wfPlan, 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)
	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())
	}

	sessionID := fmt.Sprintf("sess-u8-wave2-%d", projectID)
	sessBody, _ := json.Marshal(map[string]any{
		"session_id": sessionID,
		"project_id": projectID,
		"user_id":    creatorID,
		"title":      "U8 Wave2 intent plan",
	})
	reqSess := httptest.NewRequest(http.MethodPost, "/api/v1/chat/sessions", bytes.NewReader(sessBody))
	withBearerAuth(reqSess, creatorToken)
	recSess := httptest.NewRecorder()
	handler.ServeHTTP(recSess, reqSess)
	if recSess.Code != http.StatusOK && recSess.Code != http.StatusCreated {
		t.Fatalf("upsert chat session: %d %s", recSess.Code, recSess.Body.String())
	}

	intentflow.SetFeatureFlagsForTest(intentflow.FeatureFlags{PlanTier: intentflow.TierConfirm})
	t.Cleanup(intentflow.ResetFeatureFlagsForTest)

	taskOwner := app.NewTaskRunAdapter()
	var runID string
	taskReq := bizdev.TaskRunRequest{
		TaskKey:      "intake",
		ProjectName:  proj.Name,
		DeviceType:   proj.DeviceType,
		IntendedUse:  "Tissue resection",
		TargetMarket: "NMPA Class III",
		Backend:      "pi-agent",
		ProjectID:    projectID,
		UserID:       creatorID,
		AuthorID:     creatorID,
		ReviewerID:   reviewerID,
		ApproverID:   approverID,
		SessionID:    sessionID,
		Fields: map[string]string{
			"customer_name": "U8 Wave2 真实 PI 客户",
			"device_type":   proj.DeviceType,
			"intended_use":  "Tissue resection",
			"target_market": "NMPA Class III",
		},
		Materials: []bizdev.MaterialInput{{
			Filename: "project_brief.md",
			Content:  []byte("# U8 Wave2\n\n- 产品：超声软组织手术系统\n- 目标：确认后真实 PI generate\n"),
		}},
	}
	taskOwner.Starter = func(ctx context.Context, exec intentflow.StepExec) (string, error) {
		if existing, err := srv.TaskRuns().FindByIdempotencyKey(ctx, exec.IdempotencyKey); err != nil {
			return "", err
		} else if existing != nil && existing.RunID != "" {
			runID = existing.RunID
			return existing.RunID, nil
		}
		rec := &taskrun.RunRecord{
			SessionID:      exec.Principal.SessionID,
			ProjectID:      exec.Principal.ProjectID,
			UserID:         exec.Principal.UserID,
			IdempotencyKey: exec.IdempotencyKey,
			TaskKey:        "intake",
			Status:         "pending",
		}
		req := taskReq
		req.UserID = exec.Principal.UserID
		req.ProjectID = exec.Principal.ProjectID
		req.SessionID = exec.Principal.SessionID
		created, err := srv.TaskRuns().CreateAndStart(ctx, "", rec, req, bizdev.StageRunRequest{}, "")
		if err != nil {
			return "", err
		}
		runID = created.RunID
		return created.RunID, nil
	}

	plans := intentflow.NewPlans(uow.NewPG(pool), intentflow.NewRepository(), intentflow.Options{
		Owner: intentflow.OwnerMux{
			Query: intentflow.QueryGuard{Inner: intentflow.NewFakeActionOwner()},
			Task:  taskOwner,
		},
		Events: events.PGSink{},
	})
	prin := intentflow.Principal{
		UserID: creatorID, Environment: "test", ProjectID: projectID, SessionID: sessionID,
	}
	created, err := plans.CreatePlan(ctx, intentflow.CreateCommand{
		Principal: prin, ClientTurnID: fmt.Sprintf("u8-turn-%d", projectID), RequestDigest: "u8-wave2-d1",
		WritePlan: true, ExecutionMode: intentflow.ModeLive,
		Steps: []intentflow.StepDraft{
			{ActionKey: "regulation.query", ActionDigest: "u8-q", ExecutionMode: intentflow.ModeLive},
			{ActionKey: "task.run", ActionDigest: "u8-w", ExecutionMode: intentflow.ModeLive},
		},
	})
	if err != nil || created == nil {
		t.Fatalf("create plan: %+v %v", created, err)
	}
	t.Cleanup(func() {
		_, _ = pool.Exec(context.Background(), `DELETE FROM intent_plan_decisions WHERE plan_id=$1`, created.PlanID)
		_, _ = pool.Exec(context.Background(), `DELETE FROM intent_plan_steps WHERE plan_id=$1`, created.PlanID)
		_, _ = pool.Exec(context.Background(), `DELETE FROM intent_plans WHERE plan_id=$1`, created.PlanID)
	})
	if created.Status != intentflow.PlanAwaitingConfirmation {
		t.Fatalf("status before confirm = %s", created.Status)
	}
	if taskOwner.StartCount() != 0 {
		t.Fatal("generate must not start before confirm")
	}
	t.Logf("✓ [1. awaiting] plan_id=%s revision=%d starts=0", created.PlanID, created.Revision)

	confirmed, err := plans.DecidePlan(ctx, intentflow.DecisionCommand{
		Principal: prin, PlanID: created.PlanID, ExpectedRevision: created.Revision,
		DecisionNonce: fmt.Sprintf("u8-n-%d", projectID), DecisionDigest: "u8-d",
		Kind: intentflow.DecisionConfirm,
	})
	if err != nil {
		t.Fatal(err)
	}
	if taskOwner.StartCount() != 0 {
		t.Fatal("confirm alone must not start generate")
	}

	piStart := time.Now()
	view, err := plans.AdvanceEligible(ctx, "u8-wave2-worker", confirmed.PlanID)
	if err != nil {
		t.Fatalf("advance: %v", err)
	}
	if view == nil {
		t.Fatal("expected advanced plan view")
	}
	if taskOwner.StartCount() != 1 {
		t.Fatalf("generate starts = %d, plan=%+v", taskOwner.StartCount(), view)
	}
	if runID == "" {
		t.Fatal("starter did not capture run id")
	}
	waitCtx, cancel := context.WithTimeout(ctx, 45*time.Minute)
	defer cancel()
	if err := srv.TaskRuns().WaitRunDone(waitCtx, runID); err != nil {
		t.Fatalf("wait PI run %s: %v", runID, err)
	}
	doneRec, err := srv.TaskRuns().GetRun(runID)
	if err != nil || doneRec == nil {
		t.Fatalf("get run: %v", err)
	}
	if doneRec.Status != taskrun.StatusCompleted {
		t.Fatalf("PI run status=%s want %s", doneRec.Status, taskrun.StatusCompleted)
	}
	t.Logf("✓ [2. real PI generate] run_id=%s status=%s elapsed=%s", runID, doneRec.Status, time.Since(piStart).Round(time.Second))

	_, err = plans.AdvanceEligible(ctx, "u8-wave2-worker", confirmed.PlanID)
	if err != nil {
		t.Fatal(err)
	}
	if taskOwner.StartCount() != 1 {
		t.Fatalf("replay must not restart generate, starts=%d", taskOwner.StartCount())
	}
	t.Log("✓ [3. idempotent replay] starts still 1")

	taskOwner.Denied = true
	_, err = taskOwner.Execute(ctx, intentflow.StepExec{
		ActionKey: "task.run", IdempotencyKey: "post-revoke", ActionDigest: "x",
		Principal: prin,
	})
	if err != intentflow.ErrPlanActionForbidden {
		t.Fatalf("revoke want forbidden, got %v", err)
	}
	t.Log("✓ [4. revoke] TaskRunAdapter denied")

	restarted := intentflow.NewPlans(uow.NewPG(pool), intentflow.NewRepository(), intentflow.Options{
		Owner: intentflow.OwnerMux{
			Query: intentflow.QueryGuard{Inner: intentflow.NewFakeActionOwner()},
			Task:  app.NewTaskRunAdapter(),
		},
		Events: events.PGSink{},
	})
	reloaded, err := restarted.GetPlan(ctx, prin, created.PlanID)
	if err != nil || reloaded == nil {
		t.Fatalf("reload plan after restart wiring: %v", err)
	}
	var taskOK bool
	for _, s := range reloaded.Steps {
		if s.ActionKey == "task.run" && s.Status == intentflow.StepSucceeded {
			taskOK = true
			break
		}
	}
	if !taskOK {
		t.Fatalf("restart recovery missing succeeded task.run: %+v", reloaded.Steps)
	}
	t.Logf("✓ [5. PG restart recovery] plan_id=%s status=%s", reloaded.PlanID, reloaded.Status)

	srv.SetIntentPlansForTest(restarted)
	reqGet := httptest.NewRequest(http.MethodGet, fmt.Sprintf("%s?project_id=%d&session_id=%s",
		fmt.Sprintf("/api/v1/intent-plans/%s", created.PlanID), projectID, sessionID), nil)
	withBearerAuth(reqGet, creatorToken)
	recGet := httptest.NewRecorder()
	handler.ServeHTTP(recGet, reqGet)
	if recGet.Code != http.StatusOK {
		t.Fatalf("GET plan after re-login: %d %s", recGet.Code, recGet.Body.String())
	}
	t.Logf("✓ [6. HTTP refresh] GET plan %d bytes", recGet.Body.Len())
}