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.
This commit is contained in:
mika kuns
2026-08-10 14:48:04 +02:00
parent 6a2a19cc9e
commit 47918869e9
2 changed files with 87 additions and 18 deletions
+34 -18
View File
@@ -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);
@@ -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<string>? 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;