Fix team task claim queue stalls

server/internal/service/task.go

@@ -2262,11 +2262,12 @@ func (s *TaskService) RequeueTaskAfterClaimFailure(ctx context.Context, task db.
 //     the task out of `queued`, which the empty-queued cache cannot represent;
 //  3. short-circuit runtimes whose empty-claim verdict is cached, sampling the
 //     invalidation version for the rest BEFORE the candidate SELECT;
-//  4. list queued candidates across the non-empty set (one SELECT);
+//  4. list queued candidate agents across the non-empty set (one SELECT);
 //  5. mark still-empty runtimes so their next idle poll skips Postgres;
-//  6. claim per distinct agent via ClaimTask (unchanged — preserves the
+//  6. claim per candidate agent via ClaimTask (unchanged — preserves the
 //     per-(issue, agent) serialization, the agent concurrency cap, and every
-//     dispatch side effect) until maxTasks is reached.
+//     dispatch side effect), repeating while the agent still has capacity and
+//     maxTasks has not been reached.
 //
 // The returned slice contains both reclaimed and freshly-claimed tasks, each
 // already carrying its runtime_id so the daemon routes it to the matching
@@ -2387,47 +2388,58 @@ func (s *TaskService) ClaimTasksForRuntimes(ctx context.Context, runtimeIDs []pg
 		}
 	}
 
-	// 6. Claim per distinct agent (unchanged path → same per-(issue, agent)
+	// 6. Claim per candidate agent (unchanged path → same per-(issue, agent)
 	// serialization, capacity cap, and dispatch side effects) until maxTasks is
-	// reached.
-	triedAgents := make(map[string]struct{}, len(candidates))
-	for i := range candidates {
-		if len(claimed) >= maxTasks {
-			break
-		}
-		agentKey := util.UUIDToString(candidates[i].AgentID)
-		if _, tried := triedAgents[agentKey]; tried {
-			continue
-		}
-		triedAgents[agentKey] = struct{}{}
+	// reached. The candidate query is already collapsed to one row per
+	// (runtime, agent), so a team/squad leader with a large queue does not force
+	// an unbounded scan. Iterate in rounds so every candidate agent gets a first
+	// chance before any one agent fills extra slots with its remaining capacity.
+	exhaustedAgents := make(map[string]struct{}, len(candidates))
+	for len(claimed) < maxTasks {
+		progress := false
+		for i := range candidates {
+			if len(claimed) >= maxTasks {
+				break
+			}
+			agentKey := util.UUIDToString(candidates[i].AgentID)
+			if _, exhausted := exhaustedAgents[agentKey]; exhausted {
+				continue
+			}
 
-		task, err := s.ClaimTask(ctx, candidates[i].AgentID)
-		if err != nil {
-			// Each ClaimTask commits in its own transaction, so earlier
-			// iterations (and step-2 reclaims) are already dispatched
-			// server-side. Returning nil here would drop them and force the
-			// daemon to double-claim via HTTP fallback (MUL-4257). Return the
-			// partial batch instead; the failed agent's task stays queued.
-			if len(claimed) > 0 {
-				slog.Error("batch claim: claim task failed after partial success; returning claimed tasks to avoid loss",
-					"error", err, "claimed", len(claimed))
-				return claimed, nil
+			task, err := s.ClaimTask(ctx, candidates[i].AgentID)
+			if err != nil {
+				// Each ClaimTask commits in its own transaction, so earlier
+				// iterations (and step-2 reclaims) are already dispatched
+				// server-side. Returning nil here would drop them and force the
+				// daemon to double-claim via HTTP fallback (MUL-4257). Return the
+				// partial batch instead; the failed agent's task stays queued.
+				if len(claimed) > 0 {
+					slog.Error("batch claim: claim task failed after partial success; returning claimed tasks to avoid loss",
+						"error", err, "claimed", len(claimed))
+					return claimed, nil
+				}
+				return nil, fmt.Errorf("claim task: %w", err)
 			}
-			return nil, fmt.Errorf("claim task: %w", err)
-		}
-		if task == nil {
-			continue
+			if task == nil {
+				exhaustedAgents[agentKey] = struct{}{}
+				continue
+			}
+			// ClaimAgentTask selects by agent only; guard that the claimed task
+			// belongs to a runtime this daemon hosts. An agent with a
+			// higher-priority queued task on ANOTHER daemon's runtime could
+			// otherwise be dispatched here and dropped — matching the singular
+			// path's runtime_id guard. Such a stray dispatch is recovered by the
+			// reclaim path on the owning daemon's next poll.
+			if _, ok := runtimeInSet[util.UUIDToString(task.RuntimeID)]; !ok {
+				exhaustedAgents[agentKey] = struct{}{}
+				continue
+			}
+			claimed = append(claimed, *task)
+			progress = true
 		}
-		// ClaimAgentTask selects by agent only; guard that the claimed task
-		// belongs to a runtime this daemon hosts. An agent with a
-		// higher-priority queued task on ANOTHER daemon's runtime could
-		// otherwise be dispatched here and dropped — matching the singular
-		// path's runtime_id guard. Such a stray dispatch is recovered by the
-		// reclaim path on the owning daemon's next poll.
-		if _, ok := runtimeInSet[util.UUIDToString(task.RuntimeID)]; !ok {
-			continue
+		if !progress {
+			break
 		}
-		claimed = append(claimed, *task)
 	}
 
 	return claimed, nil

server/internal/service/task_batch_claim_test.go

@@ -95,10 +95,9 @@ func batchClaimFixture(t *testing.T, ctx context.Context, pool *pgxpool.Pool) (r
 }
 
 // TestClaimTasksForRuntimes_MultiRuntimeDrain verifies the machine-level batch
-// claim (MUL-4257): a single call claims across all runtimes, one task per
-// agent per call (matching the singular path's dedup), routes each task to its
-// runtime, respects a subsequent drain, and reports empty once nothing is
-// queued.
+// claim (MUL-4257): a single call claims across all runtimes, lets one agent
+// fill multiple free slots up to its max_concurrent_tasks, routes each task to
+// its runtime, and reports empty once nothing is queued.
 func TestClaimTasksForRuntimes_MultiRuntimeDrain(t *testing.T) {
 	ctx := context.Background()
 	pool := newTaskClaimRacePool(t)
@@ -107,42 +106,59 @@ func TestClaimTasksForRuntimes_MultiRuntimeDrain(t *testing.T) {
 	rt1, rt2 := batchClaimFixture(t, ctx, pool)
 	ids := []pgtype.UUID{util.MustParseUUID(rt1), util.MustParseUUID(rt2)}
 
-	// Call 1: one task per agent (agent1→rt1, agent2→rt2) => 2 tasks, one per runtime.
+	// Call 1: agent1 has two different-issue tasks and max_concurrent_tasks=5,
+	// so both may dispatch in the same batch; agent2 contributes one more.
 	got1, err := svc.ClaimTasksForRuntimes(ctx, ids, 5)
 	if err != nil {
 		t.Fatalf("call1: %v", err)
 	}
-	if len(got1) != 2 {
-		t.Fatalf("call1 claimed %d tasks, want 2", len(got1))
+	if len(got1) != 3 {
+		t.Fatalf("call1 claimed %d tasks, want 3", len(got1))
 	}
 	seen := map[string]int{}
 	for _, task := range got1 {
 		seen[util.UUIDToString(task.RuntimeID)]++
 	}
-	if seen[rt1] != 1 || seen[rt2] != 1 {
-		t.Fatalf("call1 runtime distribution = %v, want one task each for rt1/rt2", seen)
+	if seen[rt1] != 2 || seen[rt2] != 1 {
+		t.Fatalf("call1 runtime distribution = %v, want two tasks for rt1 and one for rt2", seen)
 	}
 
-	// Call 2: agent1 still has a second queued task (different issue, capacity 5);
-	// agent2 is drained => exactly 1 task, on rt1.
+	// Call 2: everything dispatched => empty.
 	got2, err := svc.ClaimTasksForRuntimes(ctx, ids, 5)
 	if err != nil {
 		t.Fatalf("call2: %v", err)
 	}
-	if len(got2) != 1 {
-		t.Fatalf("call2 claimed %d tasks, want 1", len(got2))
-	}
-	if util.UUIDToString(got2[0].RuntimeID) != rt1 {
-		t.Fatalf("call2 claimed runtime = %s, want rt1", util.UUIDToString(got2[0].RuntimeID))
+	if len(got2) != 0 {
+		t.Fatalf("call2 claimed %d tasks, want 0", len(got2))
 	}
+}
+
+// TestListQueuedClaimCandidatesByRuntimes_CollapsesTeamLeaderQueue pins the
+// GitHub #3166 shape: many team/squad-assigned tasks point at the same leader
+// agent and runtime. Candidate listing must stay bounded by runtime/agent, not
+// by the number of queued tasks, or the daemon's short claim request can time
+// out before the service even starts dispatching.
+func TestListQueuedClaimCandidatesByRuntimes_CollapsesTeamLeaderQueue(t *testing.T) {
+	ctx := context.Background()
+	pool := newTaskClaimRacePool(t)
+	queries := db.New(pool)
+
+	rt1, rt2 := batchClaimFixture(t, ctx, pool)
+	ids := []pgtype.UUID{util.MustParseUUID(rt1), util.MustParseUUID(rt2)}
 
-	// Call 3: everything dispatched => empty.
-	got3, err := svc.ClaimTasksForRuntimes(ctx, ids, 5)
+	got, err := queries.ListQueuedClaimCandidatesByRuntimes(ctx, ids)
 	if err != nil {
-		t.Fatalf("call3: %v", err)
+		t.Fatalf("list candidates: %v", err)
 	}
-	if len(got3) != 0 {
-		t.Fatalf("call3 claimed %d tasks, want 0", len(got3))
+	if len(got) != 2 {
+		t.Fatalf("candidate count = %d, want 2 (one per runtime/agent, not one per queued task)", len(got))
+	}
+	seen := map[string]int{}
+	for _, task := range got {
+		seen[util.UUIDToString(task.RuntimeID)]++
+	}
+	if seen[rt1] != 1 || seen[rt2] != 1 {
+		t.Fatalf("candidate runtime distribution = %v, want one representative for each runtime", seen)
 	}
 }
 

server/pkg/db/queries/agent.sql

@@ -1034,16 +1034,20 @@ WHERE runtime_id = $1 AND status IN ('queued', 'dispatched')
 ORDER BY priority DESC, created_at ASC;
 
 -- name: ListQueuedClaimCandidatesByRuntime :many
--- Returns rows the runtime can attempt to claim. Status is restricted to
--- 'queued' (in contrast to ListPendingTasksByRuntime which also includes
--- 'dispatched') because dispatched rows are by definition already owned
--- and cannot be re-claimed — including them in the candidate list pads
--- the result with rows that always lose the per-(issue, agent) race in
--- ClaimAgentTask, wasting CPU and a SELECT every poll cycle when the
--- runtime is busy on a long-running task. Backed by the partial index
--- idx_agent_task_queue_claim_candidates so the warm path is cheap.
+-- Returns one queued representative per agent on the runtime. The service uses
+-- these rows only to decide which agents should be attempted; ClaimAgentTask
+-- re-selects the exact runnable task while enforcing capacity and serialization.
+-- Collapsing here keeps a team/squad leader with a large queue from forcing the
+-- claim endpoint to read every queued row before it can dispatch work.
 SELECT * FROM agent_task_queue
-WHERE runtime_id = $1 AND status = 'queued'
+WHERE id IN (
+    SELECT id FROM (
+        SELECT DISTINCT ON (agent_id) id, priority, created_at
+        FROM agent_task_queue
+        WHERE runtime_id = $1 AND status = 'queued'
+        ORDER BY agent_id, priority DESC, created_at ASC
+    ) candidates
+)
 ORDER BY priority DESC, created_at ASC;
 
 -- name: PromoteDueDeferredTasksForRuntime :many
@@ -1055,18 +1059,20 @@ WHERE runtime_id = @runtime_id
 RETURNING *;
 
 -- name: ListQueuedClaimCandidatesByRuntimes :many
--- Batch variant of ListQueuedClaimCandidatesByRuntime (MUL-4257): returns
--- queued claim candidates across every runtime_id in the input set in ONE round
--- trip, so a daemon can list candidates for all of its runtimes with a single
--- query instead of one per runtime. Ordering matches the singular query
--- (priority, then FIFO) so the batch claim loop keeps the same fairness. The
--- runtime_id filter is served by the partial index
--- idx_agent_task_queue_claim_candidates; the cross-runtime ORDER BY still needs
--- a sort step (each runtime's slice is index-ordered, but merging several
--- runtimes' rows into one priority/FIFO order is not). The per-machine
--- candidate set is small, so this is cheap in practice.
+-- Batch variant of ListQueuedClaimCandidatesByRuntime (MUL-4257): returns one
+-- queued representative per (runtime, agent) across the input set. This keeps
+-- the request bounded by active agents rather than queued tasks, while still
+-- preserving runtime-level empty-cache bookkeeping and global priority/FIFO
+-- ordering among the representatives.
 SELECT * FROM agent_task_queue
-WHERE runtime_id = ANY(@runtime_ids::uuid[]) AND status = 'queued'
+WHERE id IN (
+    SELECT id FROM (
+        SELECT DISTINCT ON (runtime_id, agent_id) id, priority, created_at
+        FROM agent_task_queue
+        WHERE runtime_id = ANY(@runtime_ids::uuid[]) AND status = 'queued'
+        ORDER BY runtime_id, agent_id, priority DESC, created_at ASC
+    ) candidates
+)
 ORDER BY priority DESC, created_at ASC;
 
 -- name: PromoteDueDeferredTasksForRuntimes :many