From 194ce58a72a23e66dba7d12e8159856518ade2f4 Mon Sep 17 00:00:00 2001 From: mika kuns Date: Wed, 5 Aug 2026 10:49:55 +0200 Subject: [PATCH] feat(worker): add wait_for_task_change MCP tool Replaces the list handler's Start-Sleep + blind get_task poll (Phase 3 of the merge-helper prompt) with a blocking MCP tool that returns as soon as a task leaves Queued/Running, or times out. Implemented as an async DB poll (short-lived DbContext, 500ms delay, no held connection) rather than hooking HubBroadcaster, keeping the existing broadcast callers untouched. timeoutSeconds is clamped server-side to 170s, under the list handler's 200s MCP_TOOL_TIMEOUT. --- src/ClaudeDo.Data/PromptFiles.cs | 2 +- src/ClaudeDo.Worker/CLAUDE.md | 1 + .../External/TaskWaitMcpTools.cs | 83 +++++++++++ src/ClaudeDo.Worker/Program.cs | 2 + .../External/TaskWaitMcpToolsTests.cs | 129 ++++++++++++++++++ 5 files changed, 216 insertions(+), 1 deletion(-) create mode 100644 src/ClaudeDo.Worker/External/TaskWaitMcpTools.cs create mode 100644 tests/ClaudeDo.Worker.Tests/External/TaskWaitMcpToolsTests.cs diff --git a/src/ClaudeDo.Data/PromptFiles.cs b/src/ClaudeDo.Data/PromptFiles.cs index d2551f76..682f7806 100644 --- a/src/ClaudeDo.Data/PromptFiles.cs +++ b/src/ClaudeDo.Data/PromptFiles.cs @@ -270,7 +270,7 @@ public static class PromptFiles - Running or WaitingForChildren → leave it; only poll. - WaitingForReview → leave it; it goes straight to Phase 4. - Poll get_task until every task has left Queued and Running — WaitingForReview on success, Failed on error. Report progress as tasks land; do not poll silently for minutes. + Call wait_for_task_change with the ids of every task still Queued or Running (timeoutSeconds up to 170) instead of sleeping and polling get_task yourself. It returns as soon as any of them leaves Queued/Running — WaitingForReview on success, Failed on error — or reports timedOut if none did. Report progress as tasks land, then call it again with whatever ids are still Queued/Running until none remain. ## Phase 4 — Review and merge One task at a time, in the order the brief lists them. diff --git a/src/ClaudeDo.Worker/CLAUDE.md b/src/ClaudeDo.Worker/CLAUDE.md index 480252d8..6bf33fe6 100644 --- a/src/ClaudeDo.Worker/CLAUDE.md +++ b/src/ClaudeDo.Worker/CLAUDE.md @@ -42,6 +42,7 @@ Interfaces (e.g. `IQueueWaker`, `IPrimeClock`, `ITaskStateService`) live in an ` - `AgentMcpTools` — `ListAgents` - `LifecycleMcpTools` — `ResetFailedTask` - `AppSettingsMcpTools` — `GetAppSettings` (read-only; includes `MaxParallelExecutions`) + - `TaskWaitMcpTools` — `WaitForTaskChange(taskIds, timeoutSeconds = 60)`: blocks until any given task leaves `Queued`/`Running`, or times out; returns immediately for a task already outside `Queued`/`Running` (unknown ids reported as status `"NotFound"`, also immediate). Implemented as an async DB poll (short-lived `DbContext` per check, 500ms delay between checks, no held connection or busy loop) rather than hooking `HubBroadcaster` — kept deliberately isolated so it can't regress the existing broadcast callers. `timeoutSeconds` is clamped server-side to `TaskWaitMcpTools.MaxTimeoutSeconds` (170s), comfortably under the list handler's `MCP_TOOL_TIMEOUT` (200s, see `InteractiveLaunchSpecService`) so the tool reports `timedOut: true` instead of racing the client's own abort. Replaces the list handler's old "sleep + poll get_task in a loop" Phase 3 instruction (`PromptFiles.MergeHelperDefault`). - `AttachmentMcpTools` — `AddTaskAttachment(taskId, fileName, textContent?|base64Content?)`, `ListTaskAttachments`, `RemoveTaskAttachment`. Re-attaching the same fileName overwrites; add/remove refuse on a Running task. - `ExternalMcpService` also exposes two daily-prep tools: - `GetDailyPrepCandidates` — returns Idle, non-blocked tasks in a git repo NOT excluded by `AppSettings.ReportExcludedPaths` and not already `IsMyDay`, plus the current Idle MyDay tasks and `maxTasks` (= `DailyPrepMaxTasks`). Repo-exclusion logic lives in the `DailyPrepFilter` helper (same file). diff --git a/src/ClaudeDo.Worker/External/TaskWaitMcpTools.cs b/src/ClaudeDo.Worker/External/TaskWaitMcpTools.cs new file mode 100644 index 00000000..0d4bf2db --- /dev/null +++ b/src/ClaudeDo.Worker/External/TaskWaitMcpTools.cs @@ -0,0 +1,83 @@ +using System.ComponentModel; +using ClaudeDo.Data; +using Microsoft.EntityFrameworkCore; +using ModelContextProtocol.Server; +using TaskStatus = ClaudeDo.Data.Models.TaskStatus; + +namespace ClaudeDo.Worker.External; + +public sealed record TaskStatusChangeDto(string TaskId, string Status); +public sealed record WaitForTaskChangeResult(IReadOnlyList Changed, bool TimedOut); + +[McpServerToolType] +public sealed class TaskWaitMcpTools +{ + // InteractiveLaunchSpecService sets MCP_TOOL_TIMEOUT=200000ms for the list handler + // session; this cap leaves a ~30s margin so the tool itself reports TimedOut instead + // of racing the client's own abort. + internal const int MaxTimeoutSeconds = 170; + private static readonly TimeSpan PollInterval = TimeSpan.FromMilliseconds(500); + + private readonly IDbContextFactory _dbFactory; + + public TaskWaitMcpTools(IDbContextFactory dbFactory) + { + _dbFactory = dbFactory; + } + + [McpServerTool, Description( + "Blocks until at least one of the given tasks leaves Queued/Running, or until timeoutSeconds elapses " + + "(clamped server-side to 170s). Returns immediately if any task is already outside Queued/Running " + + "when called (an unknown id is reported as status \"NotFound\" and counts as changed). Use this instead " + + "of polling get_task in a loop. Result: { changed: [{ taskId, status }], timedOut }.")] + public async Task WaitForTaskChange( + string[] taskIds, int timeoutSeconds = 60, CancellationToken cancellationToken = default) + { + if (taskIds.Length == 0) + throw new ArgumentException("taskIds must not be empty.", nameof(taskIds)); + + var timeout = TimeSpan.FromSeconds(Math.Clamp(timeoutSeconds, 1, MaxTimeoutSeconds)); + using var timeoutCts = new CancellationTokenSource(timeout); + using var linked = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, timeoutCts.Token); + + try + { + while (true) + { + var changed = await CheckOnceAsync(taskIds, linked.Token); + if (changed.Count > 0) + return new WaitForTaskChangeResult(changed, TimedOut: false); + + await Task.Delay(PollInterval, linked.Token); + } + } + catch (OperationCanceledException) when (!cancellationToken.IsCancellationRequested) + { + return new WaitForTaskChangeResult(Array.Empty(), TimedOut: true); + } + } + + private async Task> CheckOnceAsync(string[] taskIds, CancellationToken ct) + { + await using var ctx = await _dbFactory.CreateDbContextAsync(ct); + var rows = await ctx.Tasks + .AsNoTracking() + .Where(t => taskIds.Contains(t.Id)) + .Select(t => new { t.Id, t.Status }) + .ToListAsync(ct); + + var byId = rows.ToDictionary(r => r.Id, r => r.Status); + var result = new List(); + foreach (var id in taskIds) + { + if (!byId.TryGetValue(id, out var status)) + { + result.Add(new TaskStatusChangeDto(id, "NotFound")); + continue; + } + if (status != TaskStatus.Queued && status != TaskStatus.Running) + result.Add(new TaskStatusChangeDto(id, status.ToString())); + } + return result; + } +} diff --git a/src/ClaudeDo.Worker/Program.cs b/src/ClaudeDo.Worker/Program.cs index 5acfe65e..6aa0fdef 100644 --- a/src/ClaudeDo.Worker/Program.cs +++ b/src/ClaudeDo.Worker/Program.cs @@ -294,6 +294,7 @@ if (cfg.ExternalMcpPort > 0) externalBuilder.Services.AddScoped(); externalBuilder.Services.AddScoped(); externalBuilder.Services.AddScoped(); + externalBuilder.Services.AddScoped(); externalBuilder.Services.AddSingleton(app.Services.GetRequiredService()); externalBuilder.Services.AddScoped(); externalBuilder.Services.AddScoped(); @@ -308,6 +309,7 @@ if (cfg.ExternalMcpPort > 0) .WithTools() .WithTools() .WithTools() + .WithTools() .WithTools(); externalBuilder.WebHost.UseUrls($"http://127.0.0.1:{cfg.ExternalMcpPort}"); diff --git a/tests/ClaudeDo.Worker.Tests/External/TaskWaitMcpToolsTests.cs b/tests/ClaudeDo.Worker.Tests/External/TaskWaitMcpToolsTests.cs new file mode 100644 index 00000000..17069cb8 --- /dev/null +++ b/tests/ClaudeDo.Worker.Tests/External/TaskWaitMcpToolsTests.cs @@ -0,0 +1,129 @@ +using System.Diagnostics; +using ClaudeDo.Data; +using ClaudeDo.Data.Models; +using ClaudeDo.Data.Repositories; +using ClaudeDo.Worker.External; +using ClaudeDo.Worker.Tests.Infrastructure; +using TaskStatus = ClaudeDo.Data.Models.TaskStatus; + +namespace ClaudeDo.Worker.Tests.External; + +public sealed class TaskWaitMcpToolsTests : IDisposable +{ + private readonly DbFixture _db = new(); + private readonly ClaudeDoDbContext _ctx; + private readonly TaskRepository _tasks; + private readonly ListRepository _lists; + + public TaskWaitMcpToolsTests() + { + _ctx = _db.CreateContext(); + _tasks = new TaskRepository(_ctx); + _lists = new ListRepository(_ctx); + } + + public void Dispose() { _ctx.Dispose(); _db.Dispose(); } + + private TaskWaitMcpTools BuildSut() => new(_db.CreateFactory()); + + private async Task SeedTaskAsync(TaskStatus status) + { + var listId = Guid.NewGuid().ToString(); + await _lists.AddAsync(new ListEntity { Id = listId, Name = "L", CreatedAt = DateTime.UtcNow }); + var task = new TaskEntity + { + Id = Guid.NewGuid().ToString(), ListId = listId, Title = "t", + Status = status, CreatedAt = DateTime.UtcNow, CommitType = "chore", + }; + await _tasks.AddAsync(task); + return task; + } + + [Fact] + public async Task WaitForTaskChange_AlreadyOutOfQueuedRunning_ReturnsImmediately() + { + var task = await SeedTaskAsync(TaskStatus.WaitingForReview); + var sut = BuildSut(); + var sw = Stopwatch.StartNew(); + + var result = await sut.WaitForTaskChange([task.Id], timeoutSeconds: 30, CancellationToken.None); + + sw.Stop(); + Assert.False(result.TimedOut); + Assert.Equal(task.Id, Assert.Single(result.Changed).TaskId); + Assert.Equal("WaitingForReview", result.Changed[0].Status); + Assert.True(sw.Elapsed < TimeSpan.FromSeconds(2), $"took {sw.Elapsed}"); + } + + [Fact] + public async Task WaitForTaskChange_UnknownId_ReturnsImmediatelyAsNotFound() + { + var sut = BuildSut(); + var sw = Stopwatch.StartNew(); + + var result = await sut.WaitForTaskChange(["missing-id"], timeoutSeconds: 30, CancellationToken.None); + + sw.Stop(); + Assert.False(result.TimedOut); + Assert.Equal("NotFound", Assert.Single(result.Changed).Status); + Assert.True(sw.Elapsed < TimeSpan.FromSeconds(2), $"took {sw.Elapsed}"); + } + + [Fact] + public async Task WaitForTaskChange_StatusChangesWhileWaiting_ReturnsBeforeTimeout() + { + var task = await SeedTaskAsync(TaskStatus.Running); + var sut = BuildSut(); + var sw = Stopwatch.StartNew(); + + var waitTask = sut.WaitForTaskChange([task.Id], timeoutSeconds: 10, CancellationToken.None); + + await Task.Delay(150); + // Simulate the status change a broadcast would announce, via a separate context + // (mirrors what TaskStateService does from a different scope/process). + await using (var writeCtx = _db.CreateContext()) + { + var writeRepo = new TaskRepository(writeCtx); + var loaded = await writeRepo.GetByIdAsync(task.Id); + loaded!.Status = TaskStatus.Done; + await writeRepo.UpdateAsync(loaded); + } + + var result = await waitTask; + sw.Stop(); + + Assert.False(result.TimedOut); + Assert.Equal("Done", Assert.Single(result.Changed).Status); + Assert.True(sw.Elapsed < TimeSpan.FromSeconds(5), $"took {sw.Elapsed}"); + } + + [Fact] + public async Task WaitForTaskChange_NoChange_TimesOut() + { + var task = await SeedTaskAsync(TaskStatus.Running); + var sut = BuildSut(); + var sw = Stopwatch.StartNew(); + + var result = await sut.WaitForTaskChange([task.Id], timeoutSeconds: 1, CancellationToken.None); + + sw.Stop(); + Assert.True(result.TimedOut); + Assert.Empty(result.Changed); + Assert.True(sw.Elapsed >= TimeSpan.FromMilliseconds(900), $"took {sw.Elapsed}"); + } + + [Fact] + public async Task WaitForTaskChange_EmptyTaskIds_Throws() + { + var sut = BuildSut(); + await Assert.ThrowsAsync(() => + sut.WaitForTaskChange([], timeoutSeconds: 5, CancellationToken.None)); + } + + [Fact] + public void MaxTimeoutSeconds_StaysComfortablyUnderMcpToolTimeout() + { + // InteractiveLaunchSpecService sets MCP_TOOL_TIMEOUT=200000ms for the list handler. + Assert.True(TaskWaitMcpTools.MaxTimeoutSeconds < 200); + } +}