Agent-System-Backend/internal/store/regeneration_account_write_test.go

package store

import (
	"errors"
	"testing"

	core "agent-system-core"
)

func TestPrepareRegenerationAccountCheckedFencesAllWrites(t *testing.T) {
	st := openTestStore(t)
	for _, account := range []struct {
		owner string
		admin bool
	}{{"admin", true}, {"alice", false}, {"bob", false}} {
		if err := st.CreateAccount(account.owner, "", "", account.admin); err != nil {
			t.Fatal(err)
		}
	}
	project, err := st.CreateProject("alice", "regeneration")
	if err != nil {
		t.Fatal(err)
	}
	conv, err := st.CreateConversation("alice", project.ID, "thread", "")
	if err != nil {
		t.Fatal(err)
	}
	run := seedRun(t, st, conv, core.RunCompleted)
	target := core.Message{ConversationID: conv.ID, RunID: run.ID, Role: core.RoleAssistant}
	sibling := core.Message{ConversationID: conv.ID, RunID: run.ID, Role: core.RoleAssistant}
	for _, msg := range []*core.Message{&target, &sibling} {
		if err := st.AppendMessage(msg); err != nil {
			t.Fatal(err)
		}
	}
	assertUnchanged := func() {
		t.Helper()
		msgs, err := st.ListMessages(conv.ID)
		if err != nil {
			t.Fatal(err)
		}
		if len(msgs) != 2 || msgs[0].Excluded || msgs[1].Excluded || msgs[0].VariantGroup != "" {
			t.Fatalf("rejected regeneration changed messages: %+v", msgs)
		}
		if _, _, ok := st.PeekRegenPending(conv.ID); ok {
			t.Fatal("rejected regeneration created pending marker")
		}
	}
	if _, _, err := st.PrepareRegenerationAccountChecked("bob", conv.ID, target.ID, false); !errors.Is(err, ErrNotFound) {
		t.Fatalf("foreign actor: %v", err)
	}
	assertUnchanged()
	account, err := st.GetAccountByOwner("alice")
	if err != nil {
		t.Fatal(err)
	}
	if _, err := st.SetAccountEnabled("admin", account.ID, account.StateVersion, false); err != nil {
		t.Fatal(err)
	}
	if _, _, err := st.PrepareRegenerationAccountChecked("alice", conv.ID, target.ID, false); !errors.Is(err, ErrAccountForbidden) {
		t.Fatalf("disabled actor: %v", err)
	}
	assertUnchanged()
	account, err = st.GetAccountByOwner("alice")
	if err != nil {
		t.Fatal(err)
	}
	if _, err := st.SetAccountEnabled("admin", account.ID, account.StateVersion, true); err != nil {
		t.Fatal(err)
	}
	group, index, err := st.PrepareRegenerationAccountChecked("alice", conv.ID, target.ID, false)
	if err != nil || group != "var_"+target.ID || index != 1 {
		t.Fatalf("active regeneration: group=%q index=%d err=%v", group, index, err)
	}
	msgs, err := st.ListMessages(conv.ID)
	if err != nil || len(msgs) != 2 || !msgs[0].Excluded || !msgs[1].Excluded || msgs[0].VariantGroup != group {
		t.Fatalf("active regeneration missed message state: %+v %v", msgs, err)
	}
	if pendingGroup, pendingIndex, ok := st.PeekRegenPending(conv.ID); !ok || pendingGroup != group || pendingIndex != index {
		t.Fatalf("active regeneration marker: group=%q index=%d ok=%v", pendingGroup, pendingIndex, ok)
	}
	if err := st.ClearRegenPendingAfterFailedStart("alice", conv.ID, group, index, false); err != nil {
		t.Fatal(err)
	}
	if _, _, ok := st.PeekRegenPending(conv.ID); ok {
		t.Fatal("failed-start cleanup left pending marker")
	}
	group, index, err = st.PrepareRegenerationAccountChecked("alice", conv.ID, target.ID, false)
	if err != nil {
		t.Fatal(err)
	}
	if err := st.ClearRegenPendingAfterFailedStart("bob", conv.ID, group, index, false); !errors.Is(err, ErrNotFound) {
		t.Fatalf("foreign marker cleanup: %v", err)
	}
	account, err = st.GetAccountByOwner("alice")
	if err != nil {
		t.Fatal(err)
	}
	if _, err := st.SetAccountEnabled("admin", account.ID, account.StateVersion, false); err != nil {
		t.Fatal(err)
	}
	if err := st.ClearRegenPendingAfterFailedStart("alice", conv.ID, group, index, false); err != nil {
		t.Fatalf("disabled marker safety cleanup: %v", err)
	}
	if _, _, ok := st.PeekRegenPending(conv.ID); ok {
		t.Fatal("failed-start cleanup left stale marker after disable")
	}
}

func TestPrepareRegenerationAccountCheckedUsesExistingVariantGroup(t *testing.T) {
	st := openTestStore(t)
	if err := st.CreateAccount("alice", "", "", false); err != nil {
		t.Fatal(err)
	}
	project, err := st.CreateProject("alice", "regeneration")
	if err != nil {
		t.Fatal(err)
	}
	conv, err := st.CreateConversation("alice", project.ID, "thread", "")
	if err != nil {
		t.Fatal(err)
	}
	target := core.Message{ConversationID: conv.ID, Role: core.RoleAssistant}
	later := core.Message{ConversationID: conv.ID, Role: core.RoleAssistant}
	for _, msg := range []*core.Message{&target, &later} {
		if err := st.AppendMessage(msg); err != nil {
			t.Fatal(err)
		}
	}
	if err := st.SetMessageFlags(target.ID, false, "variants", 1); err != nil {
		t.Fatal(err)
	}
	if err := st.SetMessageFlags(later.ID, false, "variants", 3); err != nil {
		t.Fatal(err)
	}
	group, index, err := st.PrepareRegenerationAccountChecked("alice", conv.ID, target.ID, false)
	if err != nil || group != "variants" || index != 4 {
		t.Fatalf("existing group: %q index=%d err=%v", group, index, err)
	}
	if pendingGroup, pendingIndex, ok := st.PeekRegenPending(conv.ID); !ok || pendingGroup != group || pendingIndex != index {
		t.Fatalf("next group index: group=%q index=%d ok=%v", pendingGroup, pendingIndex, ok)
	}
}

func TestFailedStartCleanupDoesNotRemoveNewerSameGroupMarker(t *testing.T) {
	st := openTestStore(t)
	if err := st.CreateAccount("alice", "", "", false); err != nil {
		t.Fatal(err)
	}
	project, err := st.CreateProject("alice", "regeneration")
	if err != nil {
		t.Fatal(err)
	}
	conv, err := st.CreateConversation("alice", project.ID, "thread", "")
	if err != nil {
		t.Fatal(err)
	}
	target := core.Message{ConversationID: conv.ID, Role: core.RoleAssistant}
	if err := st.AppendMessage(&target); err != nil {
		t.Fatal(err)
	}
	group, first, err := st.PrepareRegenerationAccountChecked("alice", conv.ID, target.ID, false)
	if err != nil {
		t.Fatal(err)
	}
	group2, second, err := st.PrepareRegenerationAccountChecked("alice", conv.ID, target.ID, false)
	if err != nil || group2 != group || second <= first {
		t.Fatalf("newer marker failed to advance: %q %d -> %q %d: %v", group, first, group2, second, err)
	}
	if err := st.ClearRegenPendingAfterFailedStart("alice", conv.ID, group, first, false); err != nil {
		t.Fatal(err)
	}
	if gotGroup, gotIndex, ok := st.PeekRegenPending(conv.ID); !ok || gotGroup != group || gotIndex != second {
		t.Fatalf("old cleanup removed newer marker: %q %d %v", gotGroup, gotIndex, ok)
	}
}