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

package store

import (
	"errors"
	"path/filepath"
	"testing"

	core "agent-system-core"
)

type executionAccountWriter interface {
	TryClaimExecutionAccountChecked(actor, conversationID, instanceID, runID string, leaseSec int, allowImplicitDefault bool) (int64, error)
	RenewExecutionAccountChecked(actor, conversationID, instanceID string, epoch int64, leaseSec int, allowImplicitDefault bool) bool
	TryAcquireExecutionSlotAccountChecked(actor, runID string, budget ExecutionBudget, allowImplicitDefault bool) error
}

func executionWriter(t *testing.T, st *Store) executionAccountWriter {
	t.Helper()
	w, ok := any(st).(executionAccountWriter)
	if !ok {
		t.Fatal("Store is missing checked execution lease/admission APIs")
	}
	return w
}

func executionFenceFixture(t *testing.T) (*Store, core.Conversation, core.Run) {
	t.Helper()
	st, err := OpenSQLite(filepath.Join(t.TempDir(), "execution.db"))
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = st.Close() })
	for _, owner := range []string{"alice", "bob"} {
		if err := st.EnsureUser(owner, ""); err != nil {
			t.Fatal(err)
		}
	}
	p, err := st.CreateProject("alice", "p")
	if err != nil {
		t.Fatal(err)
	}
	conv, err := st.CreateConversation("alice", p.ID, "c", "")
	if err != nil {
		t.Fatal(err)
	}
	run := core.Run{ID: "run-exec-account", Owner: "alice", ConversationID: conv.ID, ProjectID: p.ID, Status: core.RunPending}
	if err := st.CreateRun(&run); err != nil {
		t.Fatal(err)
	}
	return st, conv, run
}

func TestExecutionAccountFenceClaimRejectsDisabledAndForeignActor(t *testing.T) {
	st, conv, _ := executionFenceFixture(t)
	w := executionWriter(t, st)
	if _, err := w.TryClaimExecutionAccountChecked("bob", conv.ID, "instance", "", 60, false); !errors.Is(err, ErrNotFound) {
		t.Fatalf("foreign claim: got %v, want ErrNotFound", err)
	}
	if _, err := st.DB().Exec(`UPDATE users SET enabled=FALSE WHERE owner='alice'`); err != nil {
		t.Fatal(err)
	}
	if _, err := w.TryClaimExecutionAccountChecked("alice", conv.ID, "instance", "", 60, false); !errors.Is(err, ErrAccountForbidden) {
		t.Fatalf("disabled claim: got %v, want ErrAccountForbidden", err)
	}
	if _, err := st.GetExecution(conv.ID); !errors.Is(err, ErrNotFound) {
		t.Fatalf("rejected claim left execution lease: %v", err)
	}
}

func TestExecutionAccountFenceRenewDoesNotExtendAfterDisable(t *testing.T) {
	st, conv, run := executionFenceFixture(t)
	w := executionWriter(t, st)
	epoch, err := w.TryClaimExecutionAccountChecked("alice", conv.ID, "instance", run.ID, 60, false)
	if err != nil || epoch != 1 {
		t.Fatalf("initial claim: epoch=%d err=%v", epoch, err)
	}
	before, err := st.GetExecution(conv.ID)
	if err != nil {
		t.Fatal(err)
	}
	if _, err := st.DB().Exec(`UPDATE users SET enabled=FALSE WHERE owner='alice'`); err != nil {
		t.Fatal(err)
	}
	if w.RenewExecutionAccountChecked("alice", conv.ID, "instance", epoch, 180, false) {
		t.Fatal("disabled actor renewed execution")
	}
	after, err := st.GetExecution(conv.ID)
	if err != nil || after.LeaseUntil != before.LeaseUntil || after.Epoch != before.Epoch {
		t.Fatalf("disabled renew changed lease: before=%+v after=%+v err=%v", before, after, err)
	}
}

func TestExecutionAccountFenceSlotRejectsDisabledForeignRunAndBudgetOff(t *testing.T) {
	st, _, run := executionFenceFixture(t)
	w := executionWriter(t, st)
	budget := ExecutionBudget{Global: 2, PerUser: 1}
	if err := w.TryAcquireExecutionSlotAccountChecked("bob", run.ID, budget, false); !errors.Is(err, ErrNotFound) {
		t.Fatalf("foreign run slot: got %v, want ErrNotFound", err)
	}
	if _, err := st.DB().Exec(`UPDATE users SET enabled=FALSE WHERE owner='alice'`); err != nil {
		t.Fatal(err)
	}
	if err := w.TryAcquireExecutionSlotAccountChecked("alice", run.ID, budget, false); !errors.Is(err, ErrAccountForbidden) {
		t.Fatalf("disabled slot: got %v, want ErrAccountForbidden", err)
	}
	if err := w.TryAcquireExecutionSlotAccountChecked("alice", run.ID, ExecutionBudget{}, false); err != nil {
		t.Fatalf("off budget should remain a no-write fast path: %v", err)
	}
	if n, _, err := st.ActiveExecutionSlots(); err != nil || n != 0 {
		t.Fatalf("rejected/off slot created %d leases: %v", n, err)
	}
}

func TestExecutionAccountFenceImplicitLocalStopsAtRecordedDisable(t *testing.T) {
	st, err := OpenSQLite(filepath.Join(t.TempDir(), "execution.db"))
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = st.Close() })
	w := executionWriter(t, st)
	p, err := st.CreateProject(OwnerDefault, "local")
	if err != nil {
		t.Fatal(err)
	}
	conv, err := st.CreateConversation(OwnerDefault, p.ID, "c", "")
	if err != nil {
		t.Fatal(err)
	}
	run := core.Run{ID: "run-local-exec", Owner: OwnerDefault, ConversationID: conv.ID, ProjectID: p.ID, Status: core.RunPending}
	if err := st.CreateRun(&run); err != nil {
		t.Fatal(err)
	}
	epoch, err := w.TryClaimExecutionAccountChecked(OwnerDefault, conv.ID, "instance", run.ID, 60, true)
	if err != nil || epoch != 1 {
		t.Fatalf("implicit local claim: epoch=%d err=%v", epoch, err)
	}
	if err := w.TryAcquireExecutionSlotAccountChecked(OwnerDefault, run.ID, ExecutionBudget{Global: 1}, true); err != nil {
		t.Fatalf("implicit local slot: %v", err)
	}
	if err := st.EnsureUser(OwnerDefault, ""); err != nil {
		t.Fatal(err)
	}
	if _, err := st.DB().Exec(`UPDATE users SET enabled=FALSE WHERE owner=?`, OwnerDefault); err != nil {
		t.Fatal(err)
	}
	if w.RenewExecutionAccountChecked(OwnerDefault, conv.ID, "instance", epoch, 180, true) {
		t.Fatal("recorded disabled local account renewed lease")
	}
	if err := w.TryAcquireExecutionSlotAccountChecked(OwnerDefault, run.ID, ExecutionBudget{Global: 1}, true); !errors.Is(err, ErrAccountForbidden) {
		t.Fatalf("recorded disabled local account reacquired slot: %v", err)
	}
}