feat(mcp): add per-run TaskRunTokenRegistry
This commit is contained in:
25
src/ClaudeDo.Worker/Runner/TaskRunTokenRegistry.cs
Normal file
25
src/ClaudeDo.Worker/Runner/TaskRunTokenRegistry.cs
Normal 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('=');
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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);
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user