Agent-System-Backend/internal/runmgr/project_work_gate_test.go

package runmgr

import (
	"context"
	"sync/atomic"
	"testing"
	"time"

	"agent-system-backend/internal/profiles"
	backendruntime "agent-system-backend/internal/runtime"
	"agent-system-backend/pkg/extension"
	core "agent-system-core"
)

type projectGateRunner struct{ calls atomic.Int32 }

func (*projectGateRunner) Name() string { return "project-gate-test" }

func (r *projectGateRunner) Execute(backendruntime.RunCtx) error {
	r.calls.Add(1)
	return nil
}

type projectReplyRunner struct{}

func (*projectReplyRunner) Name() string { return "project-reply-test" }
func (*projectReplyRunner) Execute(rc backendruntime.RunCtx) error {
	return rc.SetAssistantParts([]core.Part{{Type: core.PartText, Text: "最终项目答复"}})
}

func enableProjectGateForTest(m *Manager) {
	m.SetProjectWorkGate(
		func(context.Context, string, string) (extension.ProjectWorkDecision, error) {
			return extension.ProjectWorkDecision{State: extension.ProjectWorkProceed, DecisionID: "test-decision"}, nil
		},
		func(context.Context, extension.ProjectRiskAnswer) (extension.ProjectWorkDecision, error) {
			return extension.ProjectWorkDecision{State: extension.ProjectWorkProceed, DecisionID: "test-decision"}, nil
		},
	)
}

func TestProjectWorkGateClarifiesBeforeProceeding(t *testing.T) {
	preflights, responses, questions := 0, 0, 0
	preflight := func(context.Context, string, string) (extension.ProjectWorkDecision, error) {
		preflights++
		if preflights == 1 {
			return extension.ProjectWorkDecision{
				State: extension.ProjectWorkClarify, RiskID: "risk-1", RiskRevision: "r1",
				Question: "效率低于人工是指更慢吗？", Options: []string{"效率更高且成本更低", "完成时间和成本更低", "确实允许效率更低"},
			}, nil
		}
		return extension.ProjectWorkDecision{State: extension.ProjectWorkProceed, DecisionID: "decision-after-recheck"}, nil
	}
	respond := func(_ context.Context, answer extension.ProjectRiskAnswer) (extension.ProjectWorkDecision, error) {
		responses++
		if answer.Actor != "alice" || answer.ProjectID != "project-1" || answer.RiskID != "risk-1" || answer.RiskRevision != "r1" || answer.Answer != "效率更高且成本更低" || answer.IdempotencyKey == "" {
			t.Fatalf("untrusted risk answer: %+v", answer)
		}
		return extension.ProjectWorkDecision{State: extension.ProjectWorkProceed, DecisionID: "decision-1"}, nil
	}
	ask := func(card core.HITLCard) (core.HITLResponse, error) {
		questions++
		if card.Question == "" || len(card.Choices) != 3 || !card.Blocked {
			t.Fatalf("risk decision did not become a blocking HITL card: %+v", card)
		}
		return core.HITLResponse{ChoiceID: "option-0"}, nil
	}
	decision, err := resolveProjectWorkDecision(context.Background(), "alice", "project-1", "run-1", "conversation.run", preflight, respond, ask)
	if err != nil || decision.State != extension.ProjectWorkProceed || preflights != 2 || responses != 1 || questions != 1 {
		t.Fatalf("gate did not wait for clarification: decision=%+v preflight=%d responses=%d questions=%d err=%v", decision, preflights, responses, questions, err)
	}
}

func TestProjectWorkGateNeverTreatsPendingOrMismatchedProceedAsApproval(t *testing.T) {
	for _, tc := range []struct {
		name     string
		decision extension.ProjectWorkDecision
	}{
		{name: "pending", decision: extension.ProjectWorkDecision{State: extension.ProjectWorkPending}},
		{name: "unbound proceed", decision: extension.ProjectWorkDecision{State: extension.ProjectWorkProceed}},
	} {
		t.Run(tc.name, func(t *testing.T) {
			preflight := func(context.Context, string, string) (extension.ProjectWorkDecision, error) { return tc.decision, nil }
			decision, err := resolveProjectWorkDecision(context.Background(), "alice", "project-1", "run-1", "conversation.run", preflight, nil, nil)
			if err == nil || decision.State == extension.ProjectWorkProceed {
				t.Fatalf("unsafe decision proceeded: %+v err=%v", decision, err)
			}
		})
	}
}

func TestRunContextUsesHITLForProjectRiskBeforeContinuing(t *testing.T) {
	rc, st := newExecEnv(t)
	rc.run.ID = "run-risk-gate"
	rc.run.Status = core.RunRunning
	if err := st.CreateRun(&rc.run); err != nil {
		t.Fatal(err)
	}
	rc.h = &runHandle{hitlCh: make(chan hitlAnswer, 1)}
	rc.h.hitlCh <- hitlAnswer{resp: core.HITLResponse{ChoiceID: "option-0"}}
	rc.mgr.SetProjectWorkAuthorization(func(context.Context, string, core.ID) (bool, error) { return true, nil })
	responded, preflights := false, 0
	rc.mgr.SetProjectWorkGate(
		func(context.Context, string, string) (extension.ProjectWorkDecision, error) {
			preflights++
			if preflights > 1 {
				return extension.ProjectWorkDecision{State: extension.ProjectWorkProceed, DecisionID: "decision-after-recheck"}, nil
			}
			return extension.ProjectWorkDecision{State: extension.ProjectWorkClarify, RiskID: "risk-1", RiskRevision: "r1", Question: "效率具体指什么？", Options: []string{"效率更高"}}, nil
		},
		func(_ context.Context, answer extension.ProjectRiskAnswer) (extension.ProjectWorkDecision, error) {
			responded = answer.Answer == "效率更高"
			return extension.ProjectWorkDecision{State: extension.ProjectWorkProceed, DecisionID: "decision-1"}, nil
		},
	)
	if err := rc.requireProjectWorkDecision("conversation.run"); err != nil || !responded || preflights != 2 {
		t.Fatalf("run did not recheck after committed human clarification: responded=%t preflights=%d err=%v", responded, preflights, err)
	}
	cards, err := st.ActiveHITLCards("alice", rc.run.ConversationID)
	if err != nil || len(cards) != 1 || cards[0].Question != "效率具体指什么？" {
		t.Fatalf("clarification card was not persisted: %+v %v", cards, err)
	}
}

func TestConversationRunWaitsForProjectSourceAnalysisBeforeRunner(t *testing.T) {
	rc, st := newExecEnv(t)
	conv, err := st.GetConversation("alice", rc.run.ConversationID)
	if err != nil {
		t.Fatal(err)
	}
	message := core.Message{ConversationID: conv.ID, Role: core.RoleUser, Parts: []core.Part{{Type: core.PartText, Text: "要求效率和成本都低于人工"}}}
	if err := st.AppendMessageWithRequest(&message, "risk-request"); err != nil {
		t.Fatal(err)
	}
	runner := &projectGateRunner{}
	rc.mgr.runner = runner
	rc.mgr.profiles = profiles.New(st)
	rc.mgr.active = map[string]*runHandle{}
	rc.mgr.SetProjectWorkAuthorization(func(context.Context, string, core.ID) (bool, error) { return true, nil })
	ready := atomic.Bool{}
	rc.mgr.SetProjectWorkGate(
		func(context.Context, string, string) (extension.ProjectWorkDecision, error) {
			if !ready.Load() {
				return extension.ProjectWorkDecision{State: extension.ProjectWorkPending}, nil
			}
			return extension.ProjectWorkDecision{State: extension.ProjectWorkProceed, DecisionID: "decision-1"}, nil
		},
		func(context.Context, extension.ProjectRiskAnswer) (extension.ProjectWorkDecision, error) {
			return extension.ProjectWorkDecision{State: extension.ProjectWorkProceed, DecisionID: "decision-1"}, nil
		},
	)
	run, err := rc.mgr.StartConversationRun(conv, message, nil, nil)
	if err != nil {
		t.Fatal(err)
	}
	time.Sleep(80 * time.Millisecond)
	if runner.calls.Load() != 0 {
		t.Fatal("model runner started while project sources were still pending")
	}
	linked, err := st.MessageByID(string(message.ID))
	if err != nil || linked.RunID != run.ID {
		t.Fatalf("committed user source was not bound before preflight: %+v %v", linked, err)
	}
	ready.Store(true)
	deadline := time.Now().Add(2 * time.Second)
	for runner.calls.Load() == 0 && time.Now().Before(deadline) {
		time.Sleep(10 * time.Millisecond)
	}
	if runner.calls.Load() != 1 {
		t.Fatalf("runner calls=%d, want one after committed proceed", runner.calls.Load())
	}
}

func TestProjectAssistantTextIsNotStreamedBeforeFinalRiskDecision(t *testing.T) {
	rc, _ := newExecEnv(t)
	rc.run.ID = "run-buffered-project-reply"
	enableProjectGateForTest(rc.mgr)
	if err := rc.Emit(core.EvMessageDelta, core.MessageDeltaPayload{Index: 0, Text: "未经风险检查的正文"}); err != nil {
		t.Fatal(err)
	}
	rc.mgr.events.mu.Lock()
	buffered := len(rc.mgr.events.buf[string(rc.run.ConversationID)])
	rc.mgr.events.mu.Unlock()
	if buffered != 0 {
		t.Fatalf("project reply leaked %d message delta(s) before final decision", buffered)
	}
}

func TestFinalProjectReplyWaitsForRiskDecisionBeforeCommit(t *testing.T) {
	rc, st := newExecEnv(t)
	conv, err := st.GetConversation("alice", rc.run.ConversationID)
	if err != nil {
		t.Fatal(err)
	}
	message := core.Message{ConversationID: conv.ID, Role: core.RoleUser, Parts: []core.Part{{Type: core.PartText, Text: "项目问题"}}}
	if err := st.AppendMessageWithRequest(&message, "final-risk-request"); err != nil {
		t.Fatal(err)
	}
	rc.mgr.runner = &projectReplyRunner{}
	rc.mgr.profiles = profiles.New(st)
	rc.mgr.active = map[string]*runHandle{}
	rc.mgr.SetProjectWorkAuthorization(func(context.Context, string, core.ID) (bool, error) { return true, nil })
	finalReady := atomic.Bool{}
	rc.mgr.SetProjectWorkGate(
		func(_ context.Context, _ string, action string) (extension.ProjectWorkDecision, error) {
			if action == "assistant.final" && !finalReady.Load() {
				return extension.ProjectWorkDecision{State: extension.ProjectWorkPending}, nil
			}
			return extension.ProjectWorkDecision{State: extension.ProjectWorkProceed, DecisionID: "decision-" + action}, nil
		},
		func(context.Context, extension.ProjectRiskAnswer) (extension.ProjectWorkDecision, error) {
			return extension.ProjectWorkDecision{State: extension.ProjectWorkProceed, DecisionID: "decision-risk"}, nil
		},
	)
	if _, err := rc.mgr.StartConversationRun(conv, message, nil, nil); err != nil {
		t.Fatal(err)
	}
	time.Sleep(80 * time.Millisecond)
	messages, err := st.ListMessages(string(conv.ID))
	if err != nil || len(messages) != 1 {
		t.Fatalf("assistant reply committed before final risk decision: %+v %v", messages, err)
	}
	finalReady.Store(true)
	deadline := time.Now().Add(2 * time.Second)
	for time.Now().Before(deadline) {
		messages, err = st.ListMessages(string(conv.ID))
		if err == nil && len(messages) == 2 {
			break
		}
		time.Sleep(10 * time.Millisecond)
	}
	if len(messages) != 2 || messages[1].Role != core.RoleAssistant || messages[1].Parts[0].Text != "最终项目答复" {
		t.Fatalf("approved final reply was not committed once: %+v", messages)
	}
}