fix(worker): claim Running before creating run resources

TaskRunner.RunAsync created the worktree (PrepareRunDirectoryAsync)
before claiming Running via StartRunningAsync. RunNow dispatches by
task id with no atomic claim of their own, so a Queued task racing
the queue picker's atomic SQL claim could hit WorktreeManager's
branch-collision self-heal, which force-removes and recreates the
winner's live worktree mid-run.

Move the claim before any resource creation and bail out immediately
when it's rejected. OverrideSlotService.RunNow also fast-rejects a
task already Running in the DB (defense in depth). RunCancellationRegistry
now refuses (and logs) a double registration instead of silently
overwriting the first CTS, so a losing dispatch's cleanup can no
longer unregister the winner's cancellation token.
This commit is contained in:
mika kuns
2026-08-06 13:33:23 +02:00
parent 0d1e3b9a6f
commit 774f9d3d13
10 changed files with 298 additions and 41 deletions
@@ -33,9 +33,15 @@ public sealed class OverrideSlotService
{
using (var context = _dbFactory.CreateDbContext())
{
var exists = await new TaskRepository(context).GetByIdAsync(taskId);
if (exists is null)
var task = await new TaskRepository(context).GetByIdAsync(taskId);
if (task is null)
throw new KeyNotFoundException($"Task '{taskId}' not found.");
// Fast-fail precheck only — not the race guard. A task sitting Queued still
// passes here and can race the queue picker's atomic claim; TaskRunner.RunAsync
// resolves that by claiming Running before creating any resources.
if (task.Status == Data.Models.TaskStatus.Running)
throw new InvalidOperationException("task is already running");
}
StartInSlot(taskId, ct => RunInSlotAsync(taskId, ct), "RunInSlotAsync failed for task {TaskId}");
@@ -10,8 +10,23 @@ namespace ClaudeDo.Worker.Queue;
public sealed class RunCancellationRegistry
{
private readonly ConcurrentDictionary<string, CancellationTokenSource> _running = new(StringComparer.Ordinal);
private readonly ILogger<RunCancellationRegistry> _logger;
public void Register(string taskId, CancellationTokenSource cts) => _running[taskId] = cts;
public RunCancellationRegistry(ILogger<RunCancellationRegistry> logger) => _logger = logger;
/// Registers the CTS driving <paramref name="taskId"/>'s run. Refuses (and logs) instead of
/// silently overwriting when a CTS is already registered for this task id — an overwrite
/// would mean a double-dispatch is in flight, and the caller that lost the race would then
/// unregister the winner's CTS out from under it, leaving CancelAsync unable to reach the
/// still-running process.
public bool Register(string taskId, CancellationTokenSource cts)
{
if (_running.TryAdd(taskId, cts)) return true;
_logger.LogWarning(
"Task {TaskId} already has a registered run; refusing to overwrite (double-dispatch?)", taskId);
return false;
}
/// Removes the registration only if <paramref name="cts"/> is still the registered
/// one — a re-run may already have registered a newer CTS under the same task id.
+26 -22
View File
@@ -87,6 +87,32 @@ public sealed class TaskRunner
attachmentPaths = attachments.Select(a => Path.Combine(_attachments.TaskDir(task.Id), a.FileName)).ToList();
}
var now = DateTime.UtcNow;
// Claim Running before creating any resources (worktree, MCP token file, ...):
// the queue picker claims Queued→Running atomically (incl. StartedAt) before
// dispatching, so only unclaimed dispatches (override slot) need to claim here.
// Claiming first means a losing double-dispatch (RunNow racing the picker for the
// same row) bails out immediately instead of creating a worktree the winner then
// has to self-heal past.
if (!alreadyClaimed)
{
var startResult = await _state.StartRunningAsync(task.Id, now, ct);
if (!startResult.Ok)
{
_logger.LogWarning("Task {TaskId} skipped: StartRunningAsync rejected ({Reason})", task.Id, startResult.Reason);
return;
}
}
else
{
// Queue-claimed dispatches skip StartRunningAsync (the atomic SQL claim in
// QueuePicker already flipped the row to Running), so it never broadcasts
// TaskUpdated for this transition. Send it here so the task-list badge flips
// live instead of staying on "Queued" until the run finishes.
await _broadcaster.TaskUpdated(task.Id);
}
await _broadcaster.TaskStarted(slot, task.Id, now);
// Determine working directory: worktree or sandbox.
var prep = await PrepareRunDirectoryAsync(task, list, ct);
if (prep.FailureReason is not null)
@@ -117,28 +143,6 @@ public sealed class TaskRunner
: "mcp__claudedo_run__AskUser",
};
var now = DateTime.UtcNow;
// The queue picker claims Queued→Running atomically (incl. StartedAt) before
// dispatching; only unclaimed dispatches (override slot) claim here.
if (!alreadyClaimed)
{
var startResult = await _state.StartRunningAsync(task.Id, now, ct);
if (!startResult.Ok)
{
_logger.LogWarning("Task {TaskId} skipped: StartRunningAsync rejected ({Reason})", task.Id, startResult.Reason);
return;
}
}
else
{
// Queue-claimed dispatches skip StartRunningAsync (the atomic SQL claim in
// QueuePicker already flipped the row to Running), so it never broadcasts
// TaskUpdated for this transition. Send it here so the task-list badge flips
// live instead of staying on "Queued" until the run finishes.
await _broadcaster.TaskUpdated(task.Id);
}
await _broadcaster.TaskStarted(slot, task.Id, now);
await _skillSeeder.SeedAsync(runDir, resolvedConfig.SkillNames, wtCtx is not null, ct);
// Build prompt: title + description + only the OPEN sub-tasks (resolved ones are dropped).