mika kuns
194ce58a72
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.
84 lines
3.4 KiB
C#
84 lines
3.4 KiB
C#
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<TaskStatusChangeDto> 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<ClaudeDoDbContext> _dbFactory;
|
|
|
|
public TaskWaitMcpTools(IDbContextFactory<ClaudeDoDbContext> 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<WaitForTaskChangeResult> 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<TaskStatusChangeDto>(), TimedOut: true);
|
|
}
|
|
}
|
|
|
|
private async Task<IReadOnlyList<TaskStatusChangeDto>> 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<TaskStatusChangeDto>();
|
|
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;
|
|
}
|
|
}
|