From 47918869e9b18acb2469684868a81631e4cb535e Mon Sep 17 00:00:00 2001 From: mika kuns Date: Mon, 10 Aug 2026 14:48:04 +0200 Subject: [PATCH] fix(worker): wire mcp config into resumed runs TaskRunner.ContinueAsync resolved a ClaudeRunConfig but never attached a per-run MCP token/config or AllowedTools the way RunAsync does, so a --resume continuation (e.g. reject-rerun with feedback) lost every mcp__claudedo_run__* tool, including AskUser. Extracted the setup into SetupMcpConfigAsync and call it from both paths, with matching cleanup in ContinueAsync's finally block. --- src/ClaudeDo.Worker/Runner/TaskRunner.cs | 52 +++++++++++------- .../Runner/ContinueAsyncExceptionTests.cs | 53 +++++++++++++++++++ 2 files changed, 87 insertions(+), 18 deletions(-) diff --git a/src/ClaudeDo.Worker/Runner/TaskRunner.cs b/src/ClaudeDo.Worker/Runner/TaskRunner.cs index e4af0422..120fec87 100644 --- a/src/ClaudeDo.Worker/Runner/TaskRunner.cs +++ b/src/ClaudeDo.Worker/Runner/TaskRunner.cs @@ -124,24 +124,7 @@ public sealed class TaskRunner var runDir = prep.RunDir!; var resolvedConfig = await ResolveConfigAsync(task, list, listConfig, null, ct); - - // Every run gets a per-run MCP identity so the agent can ask the user a - // mid-run question via AskUser. Improvement-eligible (standalone top-level) - // runs additionally get SuggestImprovement for filing out-of-scope follow-ups. - mcpToken = TaskRunTokenRegistry.GenerateToken(); - _tokens.Register(mcpToken, task.Id); - Directory.CreateDirectory(_cfg.LogRoot); - mcpConfigPath = Path.Combine(_cfg.LogRoot, $"{task.Id}_mcp.json"); - await File.WriteAllTextAsync(mcpConfigPath, BuildRunMcpConfigJson(mcpToken), ct); - - var improvementEligible = task.ParentTaskId is null && task.PlanningPhase == PlanningPhase.None; - resolvedConfig = resolvedConfig with - { - McpConfigPath = mcpConfigPath, - AllowedTools = improvementEligible - ? "mcp__claudedo_run__AskUser,mcp__claudedo_run__SuggestImprovement" - : "mcp__claudedo_run__AskUser", - }; + (mcpToken, mcpConfigPath, resolvedConfig) = await SetupMcpConfigAsync(task, resolvedConfig, ct); await _skillSeeder.SeedAsync(runDir, resolvedConfig.SkillNames, wtCtx is not null, ct); @@ -273,6 +256,9 @@ public sealed class TaskRunner await _skillSeeder.SeedAsync(runDir, resolvedConfig.SkillNames, wtCtx is not null, ct); + var (mcpToken, mcpConfigPath, mcpConfig) = await SetupMcpConfigAsync(task, resolvedConfig, ct); + resolvedConfig = mcpConfig; + try { var nextRunNumber = lastRun.RunNumber + 1; @@ -299,6 +285,36 @@ public sealed class TaskRunner _logger.LogError(ex, "Unhandled exception continuing task {TaskId}", taskId); await MarkFailed(taskId, task.Title, slot, $"Unhandled error: {ex.Message}"); } + finally + { + _tokens.Unregister(mcpToken); + try { File.Delete(mcpConfigPath); } catch { /* best effort */ } + } + } + + // Every run gets a per-run MCP identity so the agent can ask the user a mid-run question + // via AskUser. Improvement-eligible (standalone top-level) runs additionally get + // SuggestImprovement for filing out-of-scope follow-ups. Used by both a fresh run and a + // --resume continuation, which needs the same wiring or the resumed session loses every + // mcp__claudedo_run__* tool. + private async Task<(string McpToken, string McpConfigPath, ClaudeRunConfig Config)> SetupMcpConfigAsync( + TaskEntity task, ClaudeRunConfig resolvedConfig, CancellationToken ct) + { + var mcpToken = TaskRunTokenRegistry.GenerateToken(); + _tokens.Register(mcpToken, task.Id); + Directory.CreateDirectory(_cfg.LogRoot); + var mcpConfigPath = Path.Combine(_cfg.LogRoot, $"{task.Id}_mcp.json"); + await File.WriteAllTextAsync(mcpConfigPath, BuildRunMcpConfigJson(mcpToken), ct); + + var improvementEligible = task.ParentTaskId is null && task.PlanningPhase == PlanningPhase.None; + var config = resolvedConfig with + { + McpConfigPath = mcpConfigPath, + AllowedTools = improvementEligible + ? "mcp__claudedo_run__AskUser,mcp__claudedo_run__SuggestImprovement" + : "mcp__claudedo_run__AskUser", + }; + return (mcpToken, mcpConfigPath, config); } private readonly record struct RunDirResult(string? RunDir, WorktreeContext? WtCtx, string? FailureReason); diff --git a/tests/ClaudeDo.Worker.Tests/Runner/ContinueAsyncExceptionTests.cs b/tests/ClaudeDo.Worker.Tests/Runner/ContinueAsyncExceptionTests.cs index 941d0dd2..70afcd75 100644 --- a/tests/ClaudeDo.Worker.Tests/Runner/ContinueAsyncExceptionTests.cs +++ b/tests/ClaudeDo.Worker.Tests/Runner/ContinueAsyncExceptionTests.cs @@ -92,6 +92,59 @@ public sealed class ContinueAsyncExceptionTests : IDisposable Assert.Equal(TaskStatus.Failed, task.Status); } + [Fact] + public async Task ContinueAsync_WiresMcpConfigAndAllowedTools_LikeAFreshRun() + { + string listId, taskId; + + using (var ctx = _db.CreateContext()) + { + listId = Guid.NewGuid().ToString(); + ctx.Lists.Add(new ListEntity { Id = listId, Name = "L", WorkingDir = null, CreatedAt = DateTime.UtcNow }); + + taskId = Guid.NewGuid().ToString(); + ctx.Tasks.Add(new TaskEntity + { + Id = taskId, + ListId = listId, + Title = "Continue me", + Status = TaskStatus.WaitingForReview, + CreatedAt = DateTime.UtcNow, + }); + await ctx.SaveChangesAsync(); + + await new TaskRunRepository(ctx).AddAsync(new TaskRunEntity + { + Id = Guid.NewGuid().ToString(), + TaskId = taskId, + RunNumber = 1, + IsRetry = false, + Prompt = "original prompt", + SessionId = "sess-continue-mcp-test", + StartedAt = DateTime.UtcNow.AddMinutes(-5), + FinishedAt = DateTime.UtcNow.AddMinutes(-1), + ExitCode = 0, + ResultMarkdown = "first result", + }); + } + + IReadOnlyList? capturedArgs = null; + var claude = new FakeClaudeProcess((_, _, args, _, _) => + { + capturedArgs = args; + return Task.FromResult(new RunResult { ExitCode = 0, SessionId = "sess-continue-mcp-test", ResultMarkdown = "ok" }); + }); + + using var ctx2 = _db.CreateContext(); + var runner = BuildRunner(claude, ctx2); + + await runner.ContinueAsync(taskId, "please continue", "slot-1", CancellationToken.None); + + Assert.NotNull(capturedArgs); + Assert.Contains("--mcp-config", capturedArgs!); + Assert.Contains("--allowedTools", capturedArgs!); + } + private sealed class ThrowingClaudeProcess : IClaudeProcess { private readonly Exception _ex;