From ef86a8c29b8f4f8e9b6a3848f261c041a04dced3 Mon Sep 17 00:00:00 2001
From: mika kuns <mika.kuns@fb-tuning.de>
Date: Thu, 4 Jun 2026 15:50:06 +0200
Subject: [PATCH] feat(mcp): add per-run TaskRunTokenRegistry

---
 .../Runner/TaskRunTokenRegistry.cs            | 25 +++++++++++++
 .../Runner/TaskRunTokenRegistryTests.cs       | 36 +++++++++++++++++++
 2 files changed, 61 insertions(+)
 create mode 100644 src/ClaudeDo.Worker/Runner/TaskRunTokenRegistry.cs
 create mode 100644 tests/ClaudeDo.Worker.Tests/Runner/TaskRunTokenRegistryTests.cs

diff --git a/src/ClaudeDo.Worker/Runner/TaskRunTokenRegistry.cs b/src/ClaudeDo.Worker/Runner/TaskRunTokenRegistry.cs
new file mode 100644
index 0000000..a222816
--- /dev/null
+++ b/src/ClaudeDo.Worker/Runner/TaskRunTokenRegistry.cs
@@ -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('=');
+    }
+}
diff --git a/tests/ClaudeDo.Worker.Tests/Runner/TaskRunTokenRegistryTests.cs b/tests/ClaudeDo.Worker.Tests/Runner/TaskRunTokenRegistryTests.cs
new file mode 100644
index 0000000..5bd3534
--- /dev/null
+++ b/tests/ClaudeDo.Worker.Tests/Runner/TaskRunTokenRegistryTests.cs
@@ -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);
+    }
+}