diff --git a/server/internal/service/task.go b/server/internal/service/task.go index 2900d42085c..dec5e995da9 100644 --- a/server/internal/service/task.go +++ b/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 diff --git a/server/internal/service/task_batch_claim_test.go b/server/internal/service/task_batch_claim_test.go index d9e556556c1..c431d032906 100644 --- a/server/internal/service/task_batch_claim_test.go +++ b/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) } } diff --git a/server/pkg/db/generated/agent.sql.go b/server/pkg/db/generated/agent.sql.go index 00da62d99ef..517d38dad7b 100644 --- a/server/pkg/db/generated/agent.sql.go +++ b/server/pkg/db/generated/agent.sql.go @@ -3533,18 +3533,22 @@ func (q *Queries) ListPendingTasksByRuntime(ctx context.Context, runtimeID pgtyp const listQueuedClaimCandidatesByRuntime = `-- name: ListQueuedClaimCandidatesByRuntime :many SELECT id, agent_id, issue_id, status, priority, dispatched_at, started_at, completed_at, result, error, created_at, context, runtime_id, session_id, work_dir, trigger_comment_id, chat_session_id, autopilot_run_id, attempt, max_attempts, parent_task_id, failure_reason, trigger_summary, force_fresh_session, is_leader_task, wait_reason, initiator_user_id, handoff_note, prepare_lease_expires_at, squad_id, runtime_mcp_overlay, escalation_for_task_id, fire_at, originator_user_id, runtime_connected_apps, coalesced_comment_ids, delivered_comment_ids, chat_input_task_id, chat_finalize_deferred_at, originator_source, delegated_from_task_id, retry_of_task_id, rerun_of_task_id, rule_version_id, trigger_evidence_kind, trigger_evidence_ref_id, accountable_user_id 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 ` -// 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. func (q *Queries) ListQueuedClaimCandidatesByRuntime(ctx context.Context, runtimeID pgtype.UUID) ([]AgentTaskQueue, error) { rows, err := q.db.Query(ctx, listQueuedClaimCandidatesByRuntime, runtimeID) if err != nil { @@ -3615,20 +3619,22 @@ func (q *Queries) ListQueuedClaimCandidatesByRuntime(ctx context.Context, runtim const listQueuedClaimCandidatesByRuntimes = `-- name: ListQueuedClaimCandidatesByRuntimes :many SELECT id, agent_id, issue_id, status, priority, dispatched_at, started_at, completed_at, result, error, created_at, context, runtime_id, session_id, work_dir, trigger_comment_id, chat_session_id, autopilot_run_id, attempt, max_attempts, parent_task_id, failure_reason, trigger_summary, force_fresh_session, is_leader_task, wait_reason, initiator_user_id, handoff_note, prepare_lease_expires_at, squad_id, runtime_mcp_overlay, escalation_for_task_id, fire_at, originator_user_id, runtime_connected_apps, coalesced_comment_ids, delivered_comment_ids, chat_input_task_id, chat_finalize_deferred_at, originator_source, delegated_from_task_id, retry_of_task_id, rerun_of_task_id, rule_version_id, trigger_evidence_kind, trigger_evidence_ref_id, accountable_user_id FROM agent_task_queue -WHERE runtime_id = ANY($1::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($1::uuid[]) AND status = 'queued' + ORDER BY runtime_id, agent_id, priority DESC, created_at ASC + ) candidates +) ORDER BY priority DESC, created_at ASC ` -// 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. func (q *Queries) ListQueuedClaimCandidatesByRuntimes(ctx context.Context, runtimeIds []pgtype.UUID) ([]AgentTaskQueue, error) { rows, err := q.db.Query(ctx, listQueuedClaimCandidatesByRuntimes, runtimeIds) if err != nil { diff --git a/server/pkg/db/queries/agent.sql b/server/pkg/db/queries/agent.sql index 2a45a938a4a..b0482afae46 100644 --- a/server/pkg/db/queries/agent.sql +++ b/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