Agent-Squad/squad/backend_test.go

package squad

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

type labeledExec struct {
	label string
	perm  string
	tools string
	delay time.Duration
	enter func()
	leave func()
}

func (e *labeledExec) Name() string { return e.label }

func (e *labeledExec) Execute(ctx context.Context, req ExecRequest) (*ExecResult, error) {
	if e.enter != nil {
		e.enter()
	}
	if e.leave != nil {
		defer e.leave()
	}
	if e.delay > 0 {
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		case <-time.After(e.delay):
		}
	}
	return &ExecResult{Success: true, Logs: e.label + "|" + e.perm + "|" + e.tools}, nil
}

func TestConfigIdentityChangesWithBashAndTools(t *testing.T) {
	a := RoleSpec{Role: "r", Model: "m", BashPermission: BashPermissionAllow}
	d := RoleSpec{Role: "r", Model: "m", BashPermission: BashPermissionDeny}
	if ConfigIdentity(a) == ConfigIdentity(d) {
		t.Fatal("allow and deny must differ")
	}
	t1 := RoleSpec{Role: "r", Model: "m", Tools: []string{"a", "b"}}
	t2 := RoleSpec{Role: "r", Model: "m", Tools: []string{"b", "a"}}
	if ConfigIdentity(t1) != ConfigIdentity(t2) {
		t.Fatal("tool order must not change identity")
	}
	t3 := RoleSpec{Role: "r", Model: "m", Tools: []string{"a", "c"}}
	if ConfigIdentity(t1) == ConfigIdentity(t3) {
		t.Fatal("tool set change must change identity")
	}
	same := RoleSpec{Role: "r", Model: "m", BashPermission: BashPermissionAllow, MaxConcurrentTasks: 1}
	bump := RoleSpec{Role: "r", Model: "m", BashPermission: BashPermissionAllow, MaxConcurrentTasks: 8}
	if ConfigIdentity(same) != ConfigIdentity(bump) {
		t.Fatal("MaxConcurrentTasks must not affect ConfigIdentity")
	}
}

func TestResolveAllowDenyTighteningAndLoosening(t *testing.T) {
	var builds atomic.Int32
	r := NewBackendResolver()
	r.Register("cap", func(role RoleSpec) (Executor, error) {
		builds.Add(1)
		return &labeledExec{label: "cap", perm: role.BashPermission}, nil
	})

	allow := RoleSpec{Role: "writer", Backend: "cap", Model: "m1", BashPermission: BashPermissionAllow}
	e1, _, err := r.Resolve(allow, "")
	if err != nil {
		t.Fatal(err)
	}
	res, _ := e1.Execute(context.Background(), ExecRequest{})
	if res.Logs != "cap|allow|" {
		t.Fatalf("want allow executor, got %q", res.Logs)
	}

	deny := RoleSpec{Role: "writer", Backend: "cap", Model: "m1", BashPermission: BashPermissionDeny}
	e2, _, err := r.Resolve(deny, "")
	if err != nil {
		t.Fatal(err)
	}
	res, _ = e2.Execute(context.Background(), ExecRequest{})
	if res.Logs != "cap|deny|" {
		t.Fatalf("allow→deny must take effect on next Resolve, got %q", res.Logs)
	}
	if e1 == e2 {
		t.Fatal("tightened config must not reuse cached allow executor")
	}

	e3, _, err := r.Resolve(allow, "")
	if err != nil {
		t.Fatal(err)
	}
	res, _ = e3.Execute(context.Background(), ExecRequest{})
	if res.Logs != "cap|allow|" {
		t.Fatalf("deny→allow must take effect, got %q", res.Logs)
	}
	if builds.Load() < 2 {
		t.Fatalf("factory must rebuild on permission change, builds=%d", builds.Load())
	}
}

func TestResolveFactoryReplacementTakesEffect(t *testing.T) {
	r := NewBackendResolver()
	r.Register("cap", func(role RoleSpec) (Executor, error) {
		return &labeledExec{label: "v1", perm: role.BashPermission}, nil
	})
	role := RoleSpec{Role: "writer", Backend: "cap", Model: "m", BashPermission: BashPermissionAllow}
	e1, _, err := r.Resolve(role, "")
	if err != nil {
		t.Fatal(err)
	}
	res, _ := e1.Execute(context.Background(), ExecRequest{})
	if res.Logs != "v1|allow|" {
		t.Fatalf("got %q", res.Logs)
	}

	r.Register("cap", func(role RoleSpec) (Executor, error) {
		return &labeledExec{label: "v2", perm: role.BashPermission}, nil
	})
	e2, _, err := r.Resolve(role, "")
	if err != nil {
		t.Fatal(err)
	}
	res, _ = e2.Execute(context.Background(), ExecRequest{})
	if res.Logs != "v2|allow|" {
		t.Fatalf("factory replacement must take effect, got %q", res.Logs)
	}
}

func TestResolveToolConfigChangeTakesEffect(t *testing.T) {
	r := NewBackendResolver()
	r.Register("cap", func(role RoleSpec) (Executor, error) {
		tools := ""
		if len(role.Tools) > 0 {
			tools = role.Tools[0]
		}
		return &labeledExec{label: "cap", tools: tools}, nil
	})
	base := RoleSpec{Role: "writer", Backend: "cap", Model: "m", Tools: []string{"read"}}
	e1, _, err := r.Resolve(base, "")
	if err != nil {
		t.Fatal(err)
	}
	res, _ := e1.Execute(context.Background(), ExecRequest{})
	if res.Logs != "cap||read" {
		t.Fatalf("got %q", res.Logs)
	}
	next := RoleSpec{Role: "writer", Backend: "cap", Model: "m", Tools: []string{"write"}}
	e2, _, err := r.Resolve(next, "")
	if err != nil {
		t.Fatal(err)
	}
	res, _ = e2.Execute(context.Background(), ExecRequest{})
	if res.Logs != "cap||write" {
		t.Fatalf("tool-config change must take effect, got %q", res.Logs)
	}
}

func TestResolveUnchangedConfigReusesCache(t *testing.T) {
	var builds atomic.Int32
	r := NewBackendResolver()
	r.Register("cap", func(role RoleSpec) (Executor, error) {
		builds.Add(1)
		return &labeledExec{label: "cap", perm: role.BashPermission}, nil
	})
	role := RoleSpec{Role: "writer", Backend: "cap", Model: "m", BashPermission: BashPermissionAsk}
	e1, _, err := r.Resolve(role, "")
	if err != nil {
		t.Fatal(err)
	}
	e2, _, err := r.Resolve(role, "")
	if err != nil {
		t.Fatal(err)
	}
	if e1 != e2 {
		t.Fatal("unchanged config must reuse cached executor")
	}
	if builds.Load() != 1 {
		t.Fatalf("want 1 factory call, got %d", builds.Load())
	}
}

func TestResolveParallelFirstDoesNotDuplicate(t *testing.T) {
	var builds atomic.Int32
	r := NewBackendResolver()
	r.Register("cap", func(role RoleSpec) (Executor, error) {
		builds.Add(1)
		time.Sleep(20 * time.Millisecond) // widen race window
		return &labeledExec{label: "cap"}, nil
	})
	role := RoleSpec{Role: "writer", Backend: "cap", Model: "m", MaxConcurrentTasks: 2}

	const n = 32
	var wg sync.WaitGroup
	got := make([]Executor, n)
	errs := make([]error, n)
	wg.Add(n)
	for i := 0; i < n; i++ {
		i := i
		go func() {
			defer wg.Done()
			got[i], _, errs[i] = r.Resolve(role, "")
		}()
	}
	wg.Wait()
	for i, err := range errs {
		if err != nil {
			t.Fatalf("resolve %d: %v", i, err)
		}
	}
	for i := 1; i < n; i++ {
		if got[i] != got[0] {
			t.Fatalf("parallel first Resolve must share one executor")
		}
	}
	if builds.Load() != 1 {
		t.Fatalf("want single factory construction, got %d", builds.Load())
	}
}

func TestSharedGateAcrossVersionsAndLimiterChange(t *testing.T) {
	var (
		mu      sync.Mutex
		current int
		peak    int
	)
	hold := make(chan struct{})
	entered := make(chan struct{}, 8)
	trackEnter := func() {
		mu.Lock()
		current++
		if current > peak {
			peak = current
		}
		mu.Unlock()
		entered <- struct{}{}
		<-hold
	}
	trackLeave := func() {
		mu.Lock()
		current--
		mu.Unlock()
	}

	r := NewBackendResolver()
	r.Register("cap", func(role RoleSpec) (Executor, error) {
		return &labeledExec{
			label: "cap",
			perm:  role.BashPermission,
			enter: trackEnter,
			leave: trackLeave,
		}, nil
	})

	allow := RoleSpec{Role: "writer", Backend: "cap", Model: "m", BashPermission: BashPermissionAllow, MaxConcurrentTasks: 1}
	eOld, _, err := r.Resolve(allow, "")
	if err != nil {
		t.Fatal(err)
	}
	errCh := make(chan error, 4)
	go func() { _, err := eOld.Execute(context.Background(), ExecRequest{}); errCh <- err }()
	<-entered // old holds the sole slot

	deny := RoleSpec{Role: "writer", Backend: "cap", Model: "m", BashPermission: BashPermissionDeny, MaxConcurrentTasks: 1}
	eNew, _, err := r.Resolve(deny, "")
	if err != nil {
		t.Fatal(err)
	}
	if eOld == eNew {
		t.Fatal("version bump must produce a new executor")
	}

	secondDone := make(chan struct{})
	go func() {
		_, err := eNew.Execute(context.Background(), ExecRequest{})
		errCh <- err
		close(secondDone)
	}()
	select {
	case <-secondDone:
		t.Fatal("new-version execute must wait on shared role cap while old is in-flight")
	case <-time.After(40 * time.Millisecond):
	}

	// Raising MaxConcurrentTasks must update the shared gate (not spawn a bypassing limiter).
	deny2 := RoleSpec{Role: "writer", Backend: "cap", Model: "m", BashPermission: BashPermissionDeny, MaxConcurrentTasks: 2}
	eLim, _, err := r.Resolve(deny2, "")
	if err != nil {
		t.Fatal(err)
	}
	go func() { _, err := eLim.Execute(context.Background(), ExecRequest{}); errCh <- err }()

	// After setMax(2): waiting eNew can enter (active=2). eLim must still wait.
	select {
	case <-entered: // eNew
	case <-time.After(time.Second):
		t.Fatal("limiter raise must admit the waiting new-version task")
	}
	select {
	case <-entered:
		t.Fatal("third task must not bypass shared cap after limiter change")
	case <-time.After(40 * time.Millisecond):
	}

	close(hold)
	for i := 0; i < 3; i++ {
		if err := <-errCh; err != nil {
			t.Fatal(err)
		}
	}
	mu.Lock()
	p := peak
	mu.Unlock()
	if p > 2 {
		t.Fatalf("shared gate must cap total in-flight across versions; peak=%d", p)
	}
}

func TestInFlightOldVersionCannotExceedCap(t *testing.T) {
	var (
		mu      sync.Mutex
		current int
		peak    int
	)
	enter := func() {
		mu.Lock()
		current++
		if current > peak {
			peak = current
		}
		mu.Unlock()
	}
	leave := func() {
		mu.Lock()
		current--
		mu.Unlock()
	}

	r := NewBackendResolver()
	r.Register("cap", func(role RoleSpec) (Executor, error) {
		return &labeledExec{
			label: "cap",
			perm:  role.BashPermission,
			delay: 80 * time.Millisecond,
			enter: enter,
			leave: leave,
		}, nil
	})

	v1 := RoleSpec{Role: "writer", Backend: "cap", Model: "m", BashPermission: BashPermissionAllow, MaxConcurrentTasks: 2}
	e1, _, err := r.Resolve(v1, "")
	if err != nil {
		t.Fatal(err)
	}
	v2 := RoleSpec{Role: "writer", Backend: "cap", Model: "m", BashPermission: BashPermissionDeny, MaxConcurrentTasks: 2}
	e2, _, err := r.Resolve(v2, "")
	if err != nil {
		t.Fatal(err)
	}

	var wg sync.WaitGroup
	// 2 from old version + 2 from new version; shared cap=2 ⇒ peak must stay ≤2.
	for _, e := range []Executor{e1, e1, e2, e2} {
		wg.Add(1)
		e := e
		go func() {
			defer wg.Done()
			_, _ = e.Execute(context.Background(), ExecRequest{})
		}()
	}
	wg.Wait()

	mu.Lock()
	p := peak
	mu.Unlock()
	if p > 2 {
		t.Fatalf("in-flight old+new must not exceed role cap; peak=%d", p)
	}
}

func TestRegisterInstanceStubSemantics(t *testing.T) {
	r := NewBackendResolver()
	e, name, err := r.Resolve(RoleSpec{Role: "x"}, "")
	if err != nil {
		t.Fatal(err)
	}
	if name != "stub" {
		t.Fatalf("want stub, got %s", name)
	}
	if e.Name() != "stub" {
		t.Fatalf("want stub executor, got %s", e.Name())
	}
	custom := &labeledExec{label: "stub-custom"}
	r.RegisterInstance("stub", custom)
	e2, _, err := r.Resolve(RoleSpec{Role: "x"}, "")
	if err != nil {
		t.Fatal(err)
	}
	if e2.Name() != "stub-custom" {
		t.Fatalf("RegisterInstance stub replace must take effect, got %s", e2.Name())
	}
}