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:
@@ -124,24 +124,7 @@ public sealed class TaskRunner
|
|||||||
var runDir = prep.RunDir!;
|
var runDir = prep.RunDir!;
|
||||||
|
|
||||||
var resolvedConfig = await ResolveConfigAsync(task, list, listConfig, null, ct);
|
var resolvedConfig = await ResolveConfigAsync(task, list, listConfig, null, ct);
|
||||||
|
(mcpToken, mcpConfigPath, resolvedConfig) = await SetupMcpConfigAsync(task, resolvedConfig, 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",
|
|
||||||
};
|
|
||||||
|
|
||||||
await _skillSeeder.SeedAsync(runDir, resolvedConfig.SkillNames, wtCtx is not null, 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);
|
await _skillSeeder.SeedAsync(runDir, resolvedConfig.SkillNames, wtCtx is not null, ct);
|
||||||
|
|
||||||
|
var (mcpToken, mcpConfigPath, mcpConfig) = await SetupMcpConfigAsync(task, resolvedConfig, ct);
|
||||||
|
resolvedConfig = mcpConfig;
|
||||||
|
|
||||||
try
|
try
|
||||||
{
|
{
|
||||||
var nextRunNumber = lastRun.RunNumber + 1;
|
var nextRunNumber = lastRun.RunNumber + 1;
|
||||||
@@ -299,6 +285,36 @@ public sealed class TaskRunner
|
|||||||
_logger.LogError(ex, "Unhandled exception continuing task {TaskId}", taskId);
|
_logger.LogError(ex, "Unhandled exception continuing task {TaskId}", taskId);
|
||||||
await MarkFailed(taskId, task.Title, slot, $"Unhandled error: {ex.Message}");
|
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);
|
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);
|
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 sealed class ThrowingClaudeProcess : IClaudeProcess
|
||||||
{
|
{
|
||||||
private readonly Exception _ex;
|
private readonly Exception _ex;
|
||||||
|
|||||||
Reference in New Issue
Block a user