Agent-System/internal/app/query_owner_test.go

package app

import (
	"context"
	"errors"
	"strings"
	"testing"

	"chat-agent-system/internal/app/intentflow"
)

func TestQueryOwnerRejectsGenerate(t *testing.T) {
	q := NewQueryOwner(nil, nil)
	_, err := q.Execute(context.Background(), intentflow.StepExec{
		ActionKey: "task.run",
		Principal: intentflow.Principal{UserID: 1, ProjectID: 2, SessionID: "s", Environment: "dev"},
	})
	if !errors.Is(err, intentflow.ErrPlanActionForbidden) {
		t.Fatalf("got %v", err)
	}
}

func TestQueryOwnerReceiptIncludesAuthScopeDigest(t *testing.T) {
	accessOK := func(context.Context, int64, int64) error { return nil }
	q := NewQueryOwner(accessOK, nil)
	exec := intentflow.StepExec{
		ActionKey: "regulation.query", IdempotencyKey: "idem-q1", ActionDigest: "d1",
		SafeSummary: "查 GB 9706",
		Principal:   intentflow.Principal{UserID: 7, ProjectID: 42, SessionID: "sess", Environment: "dev"},
	}
	rec, err := q.Execute(context.Background(), exec)
	if err != nil {
		t.Fatal(err)
	}
	scope := AuthScopeDigest(exec)
	if !strings.Contains(rec.ResultDigest, "scope="+scope) {
		t.Fatalf("receipt missing auth scope digest: %s", rec.ResultDigest)
	}
	if !strings.HasPrefix(rec.ResultRef, "query:regulation.query:") {
		t.Fatalf("result ref must not be fake: %s", rec.ResultRef)
	}
	if !strings.Contains(rec.ResultDigest, "recover=") {
		t.Fatalf("receipt missing recover marker: %s", rec.ResultDigest)
	}
	if q.ExecCount(exec.IdempotencyKey) != 1 {
		t.Fatalf("exec count = %d", q.ExecCount(exec.IdempotencyKey))
	}
	again, err := q.Execute(context.Background(), exec)
	if err != nil {
		t.Fatal(err)
	}
	if again.ResultDigest != rec.ResultDigest {
		t.Fatal("idempotent replay must return same receipt")
	}
	if q.ExecCount(exec.IdempotencyKey) != 1 {
		t.Fatalf("replay must not re-exec, count=%d", q.ExecCount(exec.IdempotencyKey))
	}
}

func TestQueryOwnerReauthAndVersionCheck(t *testing.T) {
	denied := NewQueryOwner(func(context.Context, int64, int64) error {
		return errors.New("revoked")
	}, nil)
	_, err := denied.Execute(context.Background(), intentflow.StepExec{
		ActionKey: "regulation.query", IdempotencyKey: "k", ActionDigest: "d",
		Principal: intentflow.Principal{UserID: 1, ProjectID: 2, SessionID: "s", Environment: "dev"},
	})
	if !errors.Is(err, intentflow.ErrPlanActionForbidden) {
		t.Fatalf("revoked access must fail closed, got %v", err)
	}

	q := NewQueryOwner(func(context.Context, int64, int64) error { return nil }, nil)
	q.RequiredVersion = map[string]string{"regulation.query": "v2"}
	_, err = q.Execute(context.Background(), intentflow.StepExec{
		ActionKey: "regulation.query", IdempotencyKey: "k2", ActionDigest: "d",
		ExpectedResourceVersion: "v1",
		Principal:               intentflow.Principal{UserID: 1, ProjectID: 2, SessionID: "s", Environment: "dev"},
	})
	if !errors.Is(err, intentflow.ErrPlanActionVersionConflict) {
		t.Fatalf("want version conflict, got %v", err)
	}
}

func TestConfirmZeroWritesBeforeDispatch(t *testing.T) {
	intentflow.SetFeatureFlagsForTest(intentflow.FeatureFlags{
		PlanTier: intentflow.TierConfirm, ProjectAllowlist: []int64{42},
	})
	t.Cleanup(intentflow.ResetFeatureFlagsForTest)
	query := NewQueryOwner(func(context.Context, int64, int64) error { return nil }, nil)
	task := NewTaskRunAdapter()
	p := intentflow.NewPlans(&intentflow.SerialUoW{}, intentflow.NewMemoryStore(), intentflow.Options{
		Owner: intentflow.OwnerMux{Query: intentflow.QueryGuard{Inner: query}, Task: task},
	})
	ctx := context.Background()
	prin := intentflow.Principal{UserID: 7, Environment: "dev", ProjectID: 42, SessionID: "sess-i1"}
	created, err := p.HandleTurn(ctx, intentflow.TurnInput{
		Principal: prin, Text: "查一下相关法规，然后帮我生成报价单",
		ClientTurnID: "turn-i1", RequestDigest: "d-i1",
	})
	if err != nil || created == nil || created.Plan == nil {
		t.Fatalf("create: %+v %v", created, err)
	}
	if created.Plan.Status != intentflow.PlanAwaitingConfirmation {
		t.Fatalf("status %s", created.Plan.Status)
	}
	if task.StartCount() != 0 {
		t.Fatal("generate must not start before confirm")
	}
	// QueryOwner has no receipts until DispatchEligible runs.
	probe := intentflow.StepExec{
		ActionKey: "regulation.query", IdempotencyKey: "probe-before", ActionDigest: "x",
		Principal: prin,
	}
	if _, ok, _ := query.LookupReceipt(ctx, probe.IdempotencyKey); ok {
		t.Fatal("query must not have receipts before confirm")
	}
	confirmed, err := p.DecidePlan(ctx, intentflow.DecisionCommand{
		Principal: prin, PlanID: created.Plan.PlanID, ExpectedRevision: created.Plan.Revision,
		DecisionNonce: "i1-n", DecisionDigest: "i1-d", Kind: intentflow.DecisionConfirm,
	})
	if err != nil {
		t.Fatal(err)
	}
	if task.StartCount() != 0 {
		t.Fatal("DecidePlan alone must not start generate")
	}
	view, err := p.DispatchEligible(ctx, "i1-worker", confirmed.PlanID)
	if err != nil {
		t.Fatal(err)
	}
	if view == nil {
		t.Fatal("expected dispatched plan")
	}
	if task.StartCount() != 1 {
		t.Fatalf("generate starts = %d", task.StartCount())
	}
}