From ca3fb156ddb7aad2883b09f5bb96c7af3f83149f Mon Sep 17 00:00:00 2001 From: Eve Date: Mon, 13 Jul 2026 20:54:43 +0800 Subject: [PATCH] fix(batch-claim): return partial success instead of dropping committed claims (MUL-4257) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Sol-Boy review blocker: ClaimTasksForRuntimes reclaims (step 2) and claims per agent (step 6) in independent transactions, but a step-4 candidate-SELECT error or a mid-loop ClaimTask error did 'return nil, err' — discarding tasks already committed as dispatched. The handler 500s; the daemon sees a definite (non- uncertain) 500 and HTTP-falls-back, claiming a SECOND batch into the same free slots while the first batch waits for stale reclaim — the double-claim this PR removes. - Both error paths now prefer partial success: if any task has already committed (claimed non-empty), return it (nil error) so the handler finalizes and returns 200; the errored candidates stay queued for the next poll. The remaining error is logged. Only a genuinely empty result still returns the error (safe: no committed claim to lose, HTTP fallback just re-fails). Regression (internal/service, DB-backed, fault-injected): - PartialSuccessOnSecondAgentClaimFailure: fail the 2nd ClaimTask's Begin → the first agent's committed task is returned, not dropped. - PartialSuccessOnCandidateQueryFailureAfterReclaim: a stale dispatched task is reclaimed, then the candidate SELECT fails → the reclaimed task is returned. Co-authored-by: multica-agent --- server/internal/service/task.go | 23 +++ .../service/task_batch_claim_partial_test.go | 142 ++++++++++++++++++ 2 files changed, 165 insertions(+) create mode 100644 server/internal/service/task_batch_claim_partial_test.go diff --git a/server/internal/service/task.go b/server/internal/service/task.go index 44ea9edce3..3208d256c7 100644 --- a/server/internal/service/task.go +++ b/server/internal/service/task.go @@ -1782,6 +1782,19 @@ func (s *TaskService) ClaimTasksForRuntimes(ctx context.Context, runtimeIDs []pg // 4. One candidate SELECT across the non-empty set. candidates, err := s.Queries.ListQueuedClaimCandidatesByRuntimes(ctx, nonEmpty) if err != nil { + // Steps 2/6 commit reclaimed/claimed tasks in their own transactions, + // so `claimed` may already hold tasks dispatched server-side. Dropping + // them with a 500 makes the daemon HTTP-fall-back and claim a SECOND + // batch into the same free slots (the first batch then waits for stale + // reclaim) — the same double-claim this PR set out to remove + // (MUL-4257). Prefer partial success: hand back what committed so the + // handler finalizes and returns it; the errored candidates stay queued + // for the next poll. + if len(claimed) > 0 { + slog.Error("batch claim: candidate query failed after partial success; returning claimed tasks to avoid loss", + "error", err, "claimed", len(claimed)) + return claimed, nil + } return nil, fmt.Errorf("list queued claim candidates: %w", err) } @@ -1816,6 +1829,16 @@ func (s *TaskService) ClaimTasksForRuntimes(ctx context.Context, runtimeIDs []pg 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) } if task == nil { diff --git a/server/internal/service/task_batch_claim_partial_test.go b/server/internal/service/task_batch_claim_partial_test.go new file mode 100644 index 0000000000..e5d7475c27 --- /dev/null +++ b/server/internal/service/task_batch_claim_partial_test.go @@ -0,0 +1,142 @@ +package service + +import ( + "context" + "fmt" + "strings" + "sync" + "testing" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" + "github.com/jackc/pgx/v5/pgtype" + "github.com/multica-ai/multica/server/internal/events" + "github.com/multica-ai/multica/server/internal/util" + db "github.com/multica-ai/multica/server/pkg/db/generated" +) + +// failNthBeginTxStarter fails the Nth Begin so a specific ClaimTask inside the +// batch loop errors AFTER earlier ones have committed. +type failNthBeginTxStarter struct { + inner TxStarter + failOn int + mu sync.Mutex + count int +} + +func (s *failNthBeginTxStarter) Begin(ctx context.Context) (pgx.Tx, error) { + s.mu.Lock() + s.count++ + n := s.count + s.mu.Unlock() + if n == s.failOn { + return nil, fmt.Errorf("injected begin failure #%d", n) + } + return s.inner.Begin(ctx) +} + +// candidateFailDBTX passes every statement through EXCEPT the batch candidate +// SELECT, which it fails — simulating a candidate query error after step-2 +// reclaims have already committed. +type candidateFailDBTX struct{ inner db.DBTX } + +func (f candidateFailDBTX) Exec(ctx context.Context, sql string, args ...interface{}) (pgconn.CommandTag, error) { + return f.inner.Exec(ctx, sql, args...) +} + +func (f candidateFailDBTX) Query(ctx context.Context, sql string, args ...interface{}) (pgx.Rows, error) { + if strings.Contains(sql, "ListQueuedClaimCandidatesByRuntimes") { + return nil, fmt.Errorf("injected candidate query failure") + } + return f.inner.Query(ctx, sql, args...) +} + +func (f candidateFailDBTX) QueryRow(ctx context.Context, sql string, args ...interface{}) pgx.Row { + return f.inner.QueryRow(ctx, sql, args...) +} + +// TestClaimTasksForRuntimes_PartialSuccessOnSecondAgentClaimFailure is the +// MUL-4257 review regression: ClaimTask runs per agent in its own transaction, +// so when a later agent's claim fails the already-committed (dispatched) tasks +// must be RETURNED, not dropped with an error. Dropping them would 500 the +// handler and make the daemon HTTP-fall-back and double-claim the same slots. +func TestClaimTasksForRuntimes_PartialSuccessOnSecondAgentClaimFailure(t *testing.T) { + ctx := context.Background() + pool := newTaskClaimRacePool(t) + tx := &failNthBeginTxStarter{inner: pool, failOn: 2} // 1st agent claims, 2nd errors + svc := NewTaskService(db.New(pool), tx, nil, events.New()) + + rt1, rt2 := batchClaimFixture(t, ctx, pool) + ids := []pgtype.UUID{util.MustParseUUID(rt1), util.MustParseUUID(rt2)} + + claimed, err := svc.ClaimTasksForRuntimes(ctx, ids, 5) + if err != nil { + t.Fatalf("expected partial success (nil error), got %v", err) + } + if len(claimed) != 1 { + t.Fatalf("claimed %d tasks, want 1 (the first agent's committed task, not dropped)", len(claimed)) + } + + // The one returned task must actually be dispatched in the DB (committed), + // proving we returned a real claim rather than a phantom. + var status string + if err := pool.QueryRow(ctx, `SELECT status FROM agent_task_queue WHERE id = $1`, util.UUIDToString(claimed[0].ID)).Scan(&status); err != nil { + t.Fatalf("read status: %v", err) + } + if status != "dispatched" { + t.Fatalf("returned task status = %s, want dispatched", status) + } +} + +// TestClaimTasksForRuntimes_PartialSuccessOnCandidateQueryFailureAfterReclaim +// covers the other partial-commit path: step-2 reclaims a stale dispatched task +// (committed), then the step-4 candidate SELECT errors. The reclaimed task must +// be returned, not dropped — same double-claim risk otherwise (MUL-4257). +func TestClaimTasksForRuntimes_PartialSuccessOnCandidateQueryFailureAfterReclaim(t *testing.T) { + ctx := context.Background() + pool := newTaskClaimRacePool(t) + svc := NewTaskService(db.New(candidateFailDBTX{inner: pool}), pool, nil, events.New()) + + rt1, rt2 := batchClaimFixture(t, ctx, pool) + + // Seed a stale dispatched task on rt1 (never started, dispatched long ago, + // no live prepare lease) so step-2 reclaim picks it up. + var workspaceID, ownerID, agentID string + if err := pool.QueryRow(ctx, `SELECT workspace_id, owner_id FROM agent_runtime WHERE id = $1`, rt1).Scan(&workspaceID, &ownerID); err != nil { + t.Fatalf("load runtime: %v", err) + } + if err := pool.QueryRow(ctx, `SELECT id FROM agent WHERE runtime_id = $1 LIMIT 1`, rt1).Scan(&agentID); err != nil { + t.Fatalf("load agent: %v", err) + } + var issueID string + if err := pool.QueryRow(ctx, ` + INSERT INTO issue (workspace_id, title, status, priority, creator_id, creator_type, number, position) + VALUES ($1, 'stale dispatched', 'in_progress', 'none', $2, 'member', 900123, 99) + RETURNING id`, workspaceID, ownerID).Scan(&issueID); err != nil { + t.Fatalf("create issue: %v", err) + } + var staleTaskID string + if err := pool.QueryRow(ctx, ` + INSERT INTO agent_task_queue (agent_id, runtime_id, issue_id, status, priority, context, dispatched_at, started_at, prepare_lease_expires_at) + VALUES ($1, $2, $3, 'dispatched', 0, '{}'::jsonb, now() - interval '1 hour', NULL, NULL) + RETURNING id`, agentID, rt1, issueID).Scan(&staleTaskID); err != nil { + t.Fatalf("create stale dispatched task: %v", err) + } + t.Cleanup(func() { + c := context.Background() + pool.Exec(c, `DELETE FROM agent_task_queue WHERE id = $1`, staleTaskID) + pool.Exec(c, `DELETE FROM issue WHERE id = $1`, issueID) + }) + + ids := []pgtype.UUID{util.MustParseUUID(rt1), util.MustParseUUID(rt2)} + claimed, err := svc.ClaimTasksForRuntimes(ctx, ids, 5) + if err != nil { + t.Fatalf("expected partial success (nil error) despite candidate query failure, got %v", err) + } + if len(claimed) != 1 { + t.Fatalf("claimed %d tasks, want 1 (the reclaimed task, not dropped)", len(claimed)) + } + if util.UUIDToString(claimed[0].ID) != staleTaskID { + t.Fatalf("returned task = %s, want the reclaimed stale task %s", util.UUIDToString(claimed[0].ID), staleTaskID) + } +}