feat(mcp): add per-run TaskRunTokenRegistry

This commit is contained in:
mika kuns
2026-06-04 15:50:06 +02:00
parent da23b6cd3a
commit ef86a8c29b
2 changed files with 61 additions and 0 deletions

View File

@@ -0,0 +1,25 @@
using System.Collections.Concurrent;
using System.Security.Cryptography;
namespace ClaudeDo.Worker.Runner;
// In-memory per-run MCP identity store. A task run mints a token, registers it here,
// and tears it down when the run ends. Kept out of the DB on purpose: a run that
// outlives a Worker restart is already dead (StaleTaskRecovery flips it to Failed).
public sealed class TaskRunTokenRegistry
{
private readonly ConcurrentDictionary<string, string> _tokenToTaskId = new();
public void Register(string token, string taskId) => _tokenToTaskId[token] = taskId;
public bool TryResolve(string token, out string taskId)
{
if (_tokenToTaskId.TryGetValue(token, out var id)) { taskId = id; return true; }
taskId = string.Empty;
return false;
}
public void Unregister(string token) => _tokenToTaskId.TryRemove(token, out _);
public static string GenerateToken()
{
var bytes = RandomNumberGenerator.GetBytes(32);
return Convert.ToBase64String(bytes).Replace('+', '-').Replace('/', '_').TrimEnd('=');
}
}

View File

@@ -0,0 +1,36 @@
using ClaudeDo.Worker.Runner;
using Xunit;
namespace ClaudeDo.Worker.Tests.Runner;
public sealed class TaskRunTokenRegistryTests
{
[Fact]
public void Register_then_resolve_returns_taskId()
{
var reg = new TaskRunTokenRegistry();
reg.Register("tok", "task-1");
Assert.True(reg.TryResolve("tok", out var id));
Assert.Equal("task-1", id);
}
[Fact]
public void Unregister_removes_token()
{
var reg = new TaskRunTokenRegistry();
reg.Register("tok", "task-1");
reg.Unregister("tok");
Assert.False(reg.TryResolve("tok", out _));
}
[Fact]
public void GenerateToken_is_urlsafe_and_unique()
{
var a = TaskRunTokenRegistry.GenerateToken();
var b = TaskRunTokenRegistry.GenerateToken();
Assert.NotEqual(a, b);
Assert.DoesNotContain('+', a);
Assert.DoesNotContain('/', a);
Assert.DoesNotContain('=', a);
}
}