Files
ClaudeDo/tests/ClaudeDo.Worker.Tests/Queue/OverrideSlotServiceTests.cs
T
mika kuns 109e85da83 fix(worker): honor RunCancellationRegistry.Register's return value at both dispatch sites
Register(taskId, cts) already refuses (and logs) a double-registration, but
both call sites discarded the bool and dispatched anyway under an
unregistered CTS. If the loser then unregistered the winner's CTS during
its own cleanup, TryCancel could silently no-op against a live process.

- OverrideSlotService.StartInSlot now fails RunNow/ContinueTask loudly
  (throws) when it loses the registration race instead of registering
  over — or silently proceeding despite losing to — the queue picker.
- QueueService's picker loop retries registration briefly (the loser's own
  claim-attempt-then-unregister resolves fast) before dispatching; if
  registration never resolves it marks the already-claimed row Failed
  instead of running it unregistered.
- RunCancellationRegistry.Unregister already had compare-and-remove
  semantics (TryRemove(KeyValuePair)), so a loser's cleanup could not have
  removed the winner's CTS once registration correctly failed.

Added regression tests exercising the real registry through both dispatch
paths: RunNow losing the registration race throws without disturbing the
winner, the picker's retry succeeds and TryCancel reaches the live run when
the loser unregisters in time, and the picker fails the task instead of
running unregistered when it never does.
2026-08-06 14:34:48 +02:00

127 lines
5.7 KiB
C#

using ClaudeDo.Data;
using ClaudeDo.Data.Models;
using ClaudeDo.Data.Repositories;
using ClaudeDo.Worker.Config;
using ClaudeDo.Worker.Hub;
using ClaudeDo.Worker.Queue;
using ClaudeDo.Worker.Runner;
using ClaudeDo.Worker.Tests.Infrastructure;
using Microsoft.Extensions.Logging.Abstractions;
using TaskStatus = ClaudeDo.Data.Models.TaskStatus;
namespace ClaudeDo.Worker.Tests.Queue;
/// OverrideSlotService.RunNow's own precheck — distinct from QueueService.EnsureNotInQueueSlot,
/// which only guards a task already tracked in its in-memory queue-slot dict. This covers a
/// task the DB already shows Running (e.g. the queue picker's atomic claim landed) regardless
/// of local slot bookkeeping. The real race-safety net is TaskRunner.RunAsync claiming Running
/// before creating any resources; this precheck just fails fast for the common case.
public sealed class OverrideSlotServiceTests : IDisposable
{
private readonly DbFixture _db = new();
private readonly string _tempDir;
private readonly WorkerConfig _cfg;
public OverrideSlotServiceTests()
{
_tempDir = Path.Combine(Path.GetTempPath(), $"cd_override_{Guid.NewGuid():N}");
Directory.CreateDirectory(_tempDir);
_cfg = new WorkerConfig { SandboxRoot = _tempDir, LogRoot = _tempDir };
}
public void Dispose() { _db.Dispose(); try { Directory.Delete(_tempDir, true); } catch { } }
private OverrideSlotService BuildService(RunCancellationRegistry? runCancels = null)
{
var dbFactory = _db.CreateFactory();
var state = TaskStateServiceBuilder.Build(dbFactory).State;
var wt = new WorktreeManager(new ClaudeDo.Data.Git.GitService(), dbFactory, _cfg, NullLogger<WorktreeManager>.Instance);
var runner = new TaskRunner(new FakeClaudeProcess(), dbFactory, new HubBroadcaster(new CapturingHubContext()), wt,
new ClaudeArgsBuilder(), _cfg, NullLogger<TaskRunner>.Instance, state, new TaskRunTokenRegistry(),
new AttachmentStore(), new FakeSessionSkillSeeder(), new FakeTranscriptUsageReader());
runCancels ??= new RunCancellationRegistry(NullLogger<RunCancellationRegistry>.Instance);
return new OverrideSlotService(dbFactory, runner, NullLogger<OverrideSlotService>.Instance, runCancels);
}
[Fact]
public async Task RunNow_TaskAlreadyRunningInDb_ThrowsWithoutDispatching()
{
string listId = Guid.NewGuid().ToString(), taskId = Guid.NewGuid().ToString();
using (var ctx = _db.CreateContext())
{
ctx.Lists.Add(new ListEntity { Id = listId, Name = "L", CreatedAt = DateTime.UtcNow });
ctx.Tasks.Add(new TaskEntity
{
Id = taskId, ListId = listId, Title = "T", Status = TaskStatus.Running,
StartedAt = DateTime.UtcNow, CreatedAt = DateTime.UtcNow,
});
await ctx.SaveChangesAsync();
}
var service = BuildService();
var ex = await Assert.ThrowsAsync<InvalidOperationException>(() => service.RunNow(taskId));
Assert.Contains("already running", ex.Message);
Assert.Null(service.CurrentSlot);
}
[Fact]
public async Task RunNow_TaskQueued_DoesNotThrow()
{
string listId = Guid.NewGuid().ToString(), taskId = Guid.NewGuid().ToString();
using (var ctx = _db.CreateContext())
{
ctx.Lists.Add(new ListEntity { Id = listId, Name = "L", CreatedAt = DateTime.UtcNow });
ctx.Tasks.Add(new TaskEntity
{
Id = taskId, ListId = listId, Title = "T", Status = TaskStatus.Queued,
CreatedAt = DateTime.UtcNow,
});
await ctx.SaveChangesAsync();
}
var service = BuildService();
// The dispatch itself races with FakeClaudeProcess's immediate completion in the
// background, so this only asserts the precheck doesn't reject a claimable task.
await service.RunNow(taskId);
}
// Regression for the RunCancellationRegistry.Register-return double-dispatch bug: the
// queue picker's atomic Queued->Running claim can land between RunNow's DB precheck and
// its registry registration. If the queue side already holds the registration, RunNow
// must fail loudly instead of silently registering over it (or being ignored and then
// unregistering the winner's CTS during its own cleanup).
[Fact]
public async Task RunNow_LosesRegistrationRaceToQueue_ThrowsAndLeavesWinnersCtsUntouched()
{
string listId = Guid.NewGuid().ToString(), taskId = Guid.NewGuid().ToString();
using (var ctx = _db.CreateContext())
{
ctx.Lists.Add(new ListEntity { Id = listId, Name = "L", CreatedAt = DateTime.UtcNow });
ctx.Tasks.Add(new TaskEntity
{
Id = taskId, ListId = listId, Title = "T", Status = TaskStatus.Queued,
CreatedAt = DateTime.UtcNow,
});
await ctx.SaveChangesAsync();
}
var runCancels = new RunCancellationRegistry(NullLogger<RunCancellationRegistry>.Instance);
var service = BuildService(runCancels);
// Simulate the queue picker having already won the registration race for this task id.
using var winnerCts = new CancellationTokenSource();
Assert.True(runCancels.Register(taskId, winnerCts));
var ex = await Assert.ThrowsAsync<InvalidOperationException>(() => service.RunNow(taskId));
Assert.Contains("lost the double-dispatch race", ex.Message);
Assert.Null(service.CurrentSlot);
// The loser must not have cancelled or displaced the winner's registration.
Assert.False(winnerCts.IsCancellationRequested);
Assert.True(runCancels.TryCancel(taskId));
Assert.True(winnerCts.IsCancellationRequested);
}
}