Agent-Squad/adapters/sqlite/queue_test.go

package platform

import (
	"path/filepath"
	"sync"
	"testing"
	"time"
)

// openQueue opens a hermetic store in a temp dir and returns it with its path.
func openQueue(t *testing.T) (*Store, string) {
	t.Helper()
	path := filepath.Join(t.TempDir(), "queue.db")
	s, err := Open(path)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = s.Close() })
	return s, path
}

// newQueueRuntime registers a runtime with the given cap.
func newQueueRuntime(t *testing.T, s *Store, workspaceID, name string, cap int) Runtime {
	t.Helper()
	rt, err := s.RegisterRuntime(Runtime{
		WorkspaceID: workspaceID, Name: name, Command: "in-process",
		Protocol: "agent-native", MaxConcurrency: cap,
	})
	if err != nil {
		t.Fatal(err)
	}
	return rt
}

// newQueueAgent creates an agent with the given cap, pinned to runtimeID
// (an empty runtimeID leaves the agent unbound, claimable by any runtime).
func newQueueAgent(t *testing.T, s *Store, workspaceID, name, runtimeID string, cap int) Agent {
	t.Helper()
	agent, err := s.CreateAgent(Agent{WorkspaceID: workspaceID, Name: name, RuntimeID: runtimeID, MaxConcurrency: cap})
	if err != nil {
		t.Fatal(err)
	}
	return agent
}

// queueTask creates an issue and queues one task for agentID.
func queueTask(t *testing.T, s *Store, workspaceID, agentID, title string) Task {
	t.Helper()
	issue, err := s.CreateIssue(Issue{WorkspaceID: workspaceID, Title: title, Body: title, Status: "todo"})
	if err != nil {
		t.Fatal(err)
	}
	task, err := s.enqueue(issue, agentID, false, "")
	if err != nil {
		t.Fatal(err)
	}
	return task
}

// taskStatus reads the durable status of one task.
func taskStatus(t *testing.T, s *Store, id string) string {
	t.Helper()
	var status string
	if err := s.db.QueryRow(`SELECT status FROM task WHERE id=?`, id).Scan(&status); err != nil {
		t.Fatal(err)
	}
	return status
}

// R6: the runtime cap bounds every agent on that runtime, not one agent only.
// A per-agent cap of 4 each still yields a single in-flight task at cap 1.
func TestClaimEnforcesRuntimeCapAcrossAgents(t *testing.T) {
	for _, tc := range []struct {
		name     string
		bindToRT bool
	}{
		{name: "agents bound to the runtime", bindToRT: true},
		{name: "agents not bound to the runtime", bindToRT: false},
	} {
		t.Run(tc.name, func(t *testing.T) {
			s, _ := openQueue(t)
			ws, err := s.CreateWorkspace("lab")
			if err != nil {
				t.Fatal(err)
			}
			rt := newQueueRuntime(t, s, ws.ID, "rt", 1)
			runtimeID := ""
			if tc.bindToRT {
				runtimeID = rt.ID
			}
			first := newQueueAgent(t, s, ws.ID, "a1", runtimeID, 4)
			second := newQueueAgent(t, s, ws.ID, "a2", runtimeID, 4)
			t1 := queueTask(t, s, ws.ID, first.ID, "one")
			t2 := queueTask(t, s, ws.ID, second.ID, "two")

			claimed, ok, err := s.Claim(rt.ID, time.Now())
			if err != nil || !ok || claimed.ID != t1.ID || claimed.RuntimeID != rt.ID {
				t.Fatalf("first claim %+v ok=%v err=%v want task %s on %s", claimed, ok, err, t1.ID, rt.ID)
			}
			if _, ok, err := s.Claim(rt.ID, time.Now()); err != nil || ok {
				t.Fatalf("runtime cap 1 must refuse the other agent: ok=%v err=%v", ok, err)
			}
			if got := taskStatus(t, s, t2.ID); got != "queued" {
				t.Fatalf("second task %q must stay queued while the runtime is full", got)
			}

			if err := s.FinishTask(claimed.ID, "completed", ""); err != nil {
				t.Fatal(err)
			}
			next, ok, err := s.Claim(rt.ID, time.Now())
			if err != nil || !ok || next.ID != t2.ID {
				t.Fatalf("claim after release %+v ok=%v err=%v want task %s", next, ok, err, t2.ID)
			}
		})
	}
}

// The agent cap holds when two runtimes compete for the same agent.
func TestClaimEnforcesAgentCapAcrossRuntimes(t *testing.T) {
	s, _ := openQueue(t)
	ws, err := s.CreateWorkspace("lab")
	if err != nil {
		t.Fatal(err)
	}
	rtA := newQueueRuntime(t, s, ws.ID, "rt-a", 2)
	rtB := newQueueRuntime(t, s, ws.ID, "rt-b", 2)
	agent := newQueueAgent(t, s, ws.ID, "solo", "", 1)
	t1 := queueTask(t, s, ws.ID, agent.ID, "one")
	t2 := queueTask(t, s, ws.ID, agent.ID, "two")

	claimed, ok, err := s.Claim(rtA.ID, time.Now())
	if err != nil || !ok || claimed.ID != t1.ID {
		t.Fatalf("rt-a claim %+v ok=%v err=%v want task %s", claimed, ok, err, t1.ID)
	}
	if _, ok, err := s.Claim(rtB.ID, time.Now()); err != nil || ok {
		t.Fatalf("agent cap 1 must bound the second runtime: ok=%v err=%v", ok, err)
	}
	if got := taskStatus(t, s, t2.ID); got != "queued" {
		t.Fatalf("second task %q must stay queued while the agent is busy", got)
	}
}

// Concurrent claims never overshoot the runtime cap or an agent cap.
func TestClaimConcurrentNeverExceedsCaps(t *testing.T) {
	s, _ := openQueue(t)
	ws, err := s.CreateWorkspace("lab")
	if err != nil {
		t.Fatal(err)
	}
	rt := newQueueRuntime(t, s, ws.ID, "rt", 2)
	tight := newQueueAgent(t, s, ws.ID, "tight", rt.ID, 1)
	loose := newQueueAgent(t, s, ws.ID, "loose", rt.ID, 8)
	other := newQueueAgent(t, s, ws.ID, "other", rt.ID, 8)
	for _, agent := range []Agent{tight, loose, other} {
		queueTask(t, s, ws.ID, agent.ID, "task for "+agent.Name)
	}

	const workers = 12
	var wg sync.WaitGroup
	var mu sync.Mutex
	claimed := map[string]string{} // task id -> agent id
	failures := []string{}
	for i := 0; i < workers; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			task, ok, err := s.Claim(rt.ID, time.Now())
			mu.Lock()
			defer mu.Unlock()
			switch {
			case err != nil:
				failures = append(failures, err.Error())
			case ok:
				claimed[task.ID] = task.AgentID
			}
		}()
	}
	wg.Wait()

	if len(failures) > 0 {
		t.Fatalf("concurrent claims failed: %v", failures)
	}
	if len(claimed) == 0 {
		t.Fatal("no claim landed")
	}
	if len(claimed) > rt.MaxConcurrency {
		t.Fatalf("runtime cap %d exceeded: %v", rt.MaxConcurrency, claimed)
	}
	perAgent := map[string]int{}
	for _, agentID := range claimed {
		perAgent[agentID]++
	}
	caps := map[string]int{tight.ID: tight.MaxConcurrency, loose.ID: loose.MaxConcurrency, other.ID: other.MaxConcurrency}
	for agentID, n := range perAgent {
		if cap, known := caps[agentID]; !known {
			t.Fatalf("unknown agent %s claimed", agentID)
		} else if n > cap {
			t.Fatalf("agent cap %d exceeded by %d for %s: %v", cap, n, agentID, claimed)
		}
	}
	if got := perAgent[tight.ID]; got > 1 {
		t.Fatalf("agent cap 1 exceeded by %d: %v", got, claimed)
	}
}

// A completed or lease-expired task gives its slot back.
func TestClaimCapacityFreedByReleaseAndReclaim(t *testing.T) {
	s, _ := openQueue(t)
	ws, err := s.CreateWorkspace("lab")
	if err != nil {
		t.Fatal(err)
	}
	rt := newQueueRuntime(t, s, ws.ID, "rt", 1)
	agent := newQueueAgent(t, s, ws.ID, "solo", rt.ID, 1)
	t1 := queueTask(t, s, ws.ID, agent.ID, "one")
	t2 := queueTask(t, s, ws.ID, agent.ID, "two")

	claimed, ok, err := s.Claim(rt.ID, time.Now())
	if err != nil || !ok || claimed.ID != t1.ID {
		t.Fatalf("first claim %+v ok=%v err=%v want task %s", claimed, ok, err, t1.ID)
	}
	if _, ok, err := s.Claim(rt.ID, time.Now()); err != nil || ok {
		t.Fatalf("the agent is at its cap: ok=%v err=%v", ok, err)
	}

	if err := s.FinishTask(claimed.ID, "completed", ""); err != nil {
		t.Fatal(err)
	}
	second, ok, err := s.Claim(rt.ID, time.Now())
	if err != nil || !ok || second.ID != t2.ID {
		t.Fatalf("claim after release %+v ok=%v err=%v want task %s", second, ok, err, t2.ID)
	}

	n, err := s.RecoverOrphans(time.Now().Add(10 * time.Minute))
	if err != nil || n != 1 {
		t.Fatalf("reclaim %d %v", n, err)
	}
	if got := taskStatus(t, s, t2.ID); got != "queued" {
		t.Fatalf("reclaimed task %q must be queued again", got)
	}
	third, ok, err := s.Claim(rt.ID, time.Now())
	if err != nil || !ok || third.ID != t2.ID || third.Attempts != 1 {
		t.Fatalf("claim after reclaim %+v ok=%v err=%v want task %s with one attempt", third, ok, err, t2.ID)
	}
}

// A capacity read that fails refuses the claim; it never allows it.
func TestClaimRefusesWhenReadFails(t *testing.T) {
	t.Run("unreadable store", func(t *testing.T) {
		s, path := openQueue(t)
		ws, err := s.CreateWorkspace("lab")
		if err != nil {
			t.Fatal(err)
		}
		rt := newQueueRuntime(t, s, ws.ID, "rt", 2)
		agent := newQueueAgent(t, s, ws.ID, "a", rt.ID, 2)
		task := queueTask(t, s, ws.ID, agent.ID, "one")
		if err := s.Close(); err != nil {
			t.Fatal(err)
		}
		if claimed, ok, err := s.Claim(rt.ID, time.Now()); ok {
			t.Fatalf("an unreadable store handed out %+v (err=%v)", claimed, err)
		}
		again, err := Open(path)
		if err != nil {
			t.Fatal(err)
		}
		defer again.Close()
		if got := taskStatus(t, again, task.ID); got != "queued" {
			t.Fatalf("task %q after a failed capacity read", got)
		}
	})

	t.Run("agent lookup fails", func(t *testing.T) {
		s, _ := openQueue(t)
		ws, err := s.CreateWorkspace("lab")
		if err != nil {
			t.Fatal(err)
		}
		rt := newQueueRuntime(t, s, ws.ID, "rt", 2)
		agent := newQueueAgent(t, s, ws.ID, "a", rt.ID, 2)
		task := queueTask(t, s, ws.ID, agent.ID, "one")
		if _, err := s.db.Exec(`DROP TABLE agent`); err != nil {
			t.Fatal(err)
		}
		claimed, ok, err := s.Claim(rt.ID, time.Now())
		if ok || err != nil {
			t.Fatalf("a failed agent read must skip the candidate: %+v ok=%v err=%v", claimed, ok, err)
		}
		if _, err := s.db.Exec(`CREATE TABLE agent (id TEXT PRIMARY KEY, runtime_id TEXT)`); err != nil {
			t.Fatal(err)
		}
		if got := taskStatus(t, s, task.ID); got != "queued" {
			t.Fatalf("task %q after a skipped candidate", got)
		}
	})

	// The guarded count runs inside the atomic claim: when it cannot be read
	// the statement must fail rather than hand the task out.
	t.Run("guard count fails", func(t *testing.T) {
		s, _ := openQueue(t)
		ws, err := s.CreateWorkspace("lab")
		if err != nil {
			t.Fatal(err)
		}
		rt := newQueueRuntime(t, s, ws.ID, "rt", 2)
		agent := newQueueAgent(t, s, ws.ID, "a", rt.ID, 2)
		task := queueTask(t, s, ws.ID, agent.ID, "one")
		if _, err := s.db.Exec(`DROP TABLE agent`); err != nil {
			t.Fatal(err)
		}
		lease := time.Now().UTC().Add(5 * time.Minute).Format(time.RFC3339Nano)
		res, err := s.db.Exec(claimSQL, rt.ID, lease, now(), task.ID, agent.ID, 2, rt.ID, rt.ID, rt.MaxConcurrency)
		if err == nil {
			n, _ := res.RowsAffected()
			t.Fatalf("a failed guard count claimed %d task(s)", n)
		}
		if _, err := s.db.Exec(`CREATE TABLE agent (id TEXT PRIMARY KEY, runtime_id TEXT)`); err != nil {
			t.Fatal(err)
		}
		if got := taskStatus(t, s, task.ID); got != "queued" {
			t.Fatalf("task %q after a failed guard count", got)
		}
	})
}