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

package runmgr

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"agent-system-backend/internal/store"
	core "agent-system-core"
)

// 协作存储 U5：双实例语义（KTD4/KTD5）。两个独立 Manager（A/B）共享同一
// PG schema（等价两真实进程的存储视图；进程级隔离由 CI 双进程验收与部署
// 合同覆盖——本测试验证所有权/围栏/命令的库级正确性）。
// 场景：同会话争抢、跨实例取消、旧 owner 迟到结算被拒。

func pgSchemaStore(t *testing.T) *store.Store {
	t.Helper()
	dsn := os.Getenv("AGENT_PG_TEST_DSN")
	if dsn == "" {
		t.Skip("AGENT_PG_TEST_DSN 未设置：真库多实例语义由 CI 必跑 job 覆盖")
	}
	schema := "u5_" + strings.ToLower(store.NewID("t")) + "_" + fmt.Sprint(os.Getpid())
	admin, err := sql.Open("pgx", dsn)
	if err != nil {
		t.Fatal(err)
	}
	if _, err := admin.Exec(`CREATE EXTENSION IF NOT EXISTS pg_trgm`); err != nil {
		admin.Close()
		t.Fatalf("pg_trgm: %v", err)
	}
	if _, err := admin.Exec("CREATE SCHEMA " + schema); err != nil {
		admin.Close()
		t.Fatalf("schema: %v", err)
	}
	t.Cleanup(func() {
		admin.Exec("DROP SCHEMA IF EXISTS " + schema + " CASCADE")
		admin.Close()
	})
	st, err := store.OpenPostgres(store.AppendPGRuntimeParam(dsn, "search_path", schema+",public"), filepath.Join(t.TempDir(), "art"))
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { st.Close() })
	return st
}

func newOwnedManager(t *testing.T, st *store.Store, tag string) *Manager {
	t.Helper()
	mgr := New(st, nil, nil, nil, nil, noopRunner{}, noopRunner{}, noopRunner{}, noopRunner{})
	mgr.SetInstanceTag(tag, "127.0.0.1:0")
	t.Cleanup(mgr.Shutdown)
	return mgr
}

// 同会话争抢：A 占用后 B 的 ClaimConversation 被 409 拒（库级租约裁决）。
func TestMultiInstanceConversationContention(t *testing.T) {
	st := pgSchemaStore(t)
	a := newOwnedManager(t, st, "instA")
	b := newOwnedManager(t, st, "instB")
	if err := a.ClaimConversation("cnv_race"); err != nil {
		t.Fatalf("A claim: %v", err)
	}
	defer a.ReleaseConversation("cnv_race")
	if err := b.ClaimConversation("cnv_race"); err == nil || !errors.Is(err, ErrConversationBusy) {
		t.Fatalf("B claim should be busy: %v", err)
	}
	// A 释放后 B 可占。
	a.ReleaseConversation("cnv_race")
	if err := b.ClaimConversation("cnv_race"); err != nil {
		t.Fatalf("B claim after release: %v", err)
	}
}

// 旧 owner 迟到结算被拒（fencing）：A 持 epoch N；租约到期 B 接管（N+1）；
// A 的 SettleRunFenced 被围栏拒绝、run 不变；B 的结算成功。
func TestMultiInstanceStaleOwnerSettleFenced(t *testing.T) {
	st := pgSchemaStore(t)
	conv := recoveryConversation(t, st)
	// A 以 1 秒租约占用并开 run。
	eA, err := st.TryClaimExecution(string(conv.ID), "instA", "", 1)
	if err != nil {
		t.Fatal(err)
	}
	run := recoveryRun(t, st, conv, core.RunRunning)
	time.Sleep(1100 * time.Millisecond) // 租约到期
	// B 接管。
	eB, err := st.TryClaimExecution(string(conv.ID), "instB", "", 60)
	if err != nil || eB != eA+1 {
		t.Fatalf("B takeover: epoch=%d (want %d) err=%v", eB, eA+1, err)
	}
	// A 迟到结算：fenced，run 不变（仍 running）。
	runA := run
	runA.Status = core.RunCompleted
	updated, fenced, err := st.SettleRunFenced(&runA, nil, "instA", eA)
	if err != nil || updated || !fenced {
		t.Fatalf("A settle: updated=%v fenced=%v err=%v（应 fenced 拒写）", updated, fenced, err)
	}
	after, _ := st.GetRunByID(run.ID)
	if after.Status != core.RunRunning {
		t.Fatalf("fenced settle 不得改写 run：status=%s", after.Status)
	}
	// B 结算成功。
	runB := run
	runB.Status = core.RunCompleted
	updated, fenced, err = st.SettleRunFenced(&runB, nil, "instB", eB)
	if err != nil || !updated || fenced {
		t.Fatalf("B settle: updated=%v fenced=%v err=%v", updated, fenced, err)
	}
}

// 跨实例取消（KTD5）：run 由 A 执行；B 侧投递取消命令；A 的消费循环领取并
// 取消（本地句柄路径）；命令恰被消费一次。
func TestMultiInstanceCrossCancelViaCommand(t *testing.T) {
	st := pgSchemaStore(t)
	a := newOwnedManager(t, st, "instA")
	conv := recoveryConversation(t, st)
	if err := a.ClaimConversation(string(conv.ID)); err != nil {
		t.Fatal(err)
	}
	defer a.ReleaseConversation(string(conv.ID))
	run := recoveryRun(t, st, conv, core.RunRunning)
	a.execMu.Lock()
	a.execEpoch[string(conv.ID)], _ = st.TryClaimExecution(string(conv.ID), "instA", run.ID, 60)
	a.execMu.Unlock()

	// B 侧投递取消命令（等价 server 的跨实例路径）。
	cmdID := store.NewID("cmd")
	if err := st.EnqueueCommand(cmdID, "cancel", run.ID, string(conv.ID), 0,
		CancelCommandPayload{Owner: conv.Owner, RunID: run.ID}); err != nil {
		t.Fatal(err)
	}
	// A 的单轮消费：认领并执行（run 非本实例活跃 → 库级取消收口）。
	a.consumeCommandsOnce()
	state, _, err := st.CommandState(cmdID)
	if err != nil || state != "consumed" {
		t.Fatalf("command state=%s err=%v（应 consumed）", state, err)
	}
	after, _ := st.GetRunByID(run.ID)
	if after.Status != core.RunCancelled {
		t.Fatalf("run status=%s（应 cancelled）", after.Status)
	}
	// 重复消费不重复执行（命令已 consumed，单轮无动作）。
	a.consumeCommandsOnce()
	after2, _ := st.GetRunByID(run.ID)
	if after2.Status != core.RunCancelled {
		t.Fatalf("重复消费后 status=%s", after2.Status)
	}
}

// 非执行实例不吃命令：命令的会话执行权在 B，A 的消费循环应跳过。
func TestCommandOnlyConsumedByOwner(t *testing.T) {
	st := pgSchemaStore(t)
	a := newOwnedManager(t, st, "instA")
	conv := recoveryConversation(t, st)
	run := recoveryRun(t, st, conv, core.RunRunning)
	if _, err := st.TryClaimExecution(string(conv.ID), "instB", run.ID, 60); err != nil {
		t.Fatal(err)
	}
	cmdID := store.NewID("cmd")
	if err := st.EnqueueCommand(cmdID, "cancel", run.ID, string(conv.ID), 0, nil); err != nil {
		t.Fatal(err)
	}
	a.consumeCommandsOnce()
	state, _, _ := st.CommandState(cmdID)
	if state != "pending" {
		t.Fatalf("非执行实例不得消费命令：state=%s", state)
	}
	_ = context.Background()
}