ClaudeDo/src/ClaudeDo.Worker/Runner/LiveSessionRegistry.cs

using System.Collections.Concurrent;
using ClaudeDo.Worker.Runner.Interfaces;

namespace ClaudeDo.Worker.Runner;

// Singleton in-memory registry of active live streaming sessions.
// A session's lifetime matches its associated task run; dead entries are removed by the runner.
//
// Interactive (stream-json) sessions never exit on their own — they wait on stdin — and there is
// no client-disconnect teardown, so an abandoned chat would otherwise keep its claude.exe alive
// for the worker's whole lifetime. IdleSessionReaper periodically stops sessions that have seen
// no activity past a timeout (see ReapIdleAsync); Touch() records that activity.
public sealed class LiveSessionRegistry
{
    private sealed class Entry
    {
        public required ILiveSession Session { get; init; }
        public long LastActivityTicksUtc;
    }

    private readonly ConcurrentDictionary<string, Entry> _sessions = new();

    public void Register(string taskId, ILiveSession session)
    {
        if (_sessions.TryRemove(taskId, out var existing))
        {
            // Best-effort stop of the replaced session; don't await to avoid deadlock risk.
            _ = existing.Session.StopAsync().ContinueWith(t =>
            {
                if (t.IsFaulted) { /* swallow — old session is already orphaned */ }
            }, TaskScheduler.Default);
        }
        _sessions[taskId] = new Entry { Session = session, LastActivityTicksUtc = DateTime.UtcNow.Ticks };
    }

    // Marks a session as active so the idle reaper leaves it alone. Called on every user
    // message and every output line. No-op if the session is not (yet) registered.
    public void Touch(string taskId)
    {
        if (_sessions.TryGetValue(taskId, out var entry))
            Interlocked.Exchange(ref entry.LastActivityTicksUtc, DateTime.UtcNow.Ticks);
    }

    public bool TryGet(string taskId, out ILiveSession session)
    {
        if (_sessions.TryGetValue(taskId, out var entry))
        {
            session = entry.Session;
            return true;
        }
        session = null!;
        return false;
    }

    public void Unregister(string taskId) => _sessions.TryRemove(taskId, out _);

    public async Task StopAsync(string taskId)
    {
        if (_sessions.TryRemove(taskId, out var entry))
            await entry.Session.StopAsync();
    }

    // Stops and removes every session whose last activity is older than (nowUtc - idleTimeout),
    // skipping any session with a turn in flight (an agent that's actively working, even if quiet).
    // Returns the reaped task ids.
    public async Task<IReadOnlyList<string>> ReapIdleAsync(DateTime nowUtc, TimeSpan idleTimeout)
    {
        var cutoffTicks = (nowUtc - idleTimeout).Ticks;
        List<string>? reaped = null;

        foreach (var kvp in _sessions)
        {
            var entry = kvp.Value;
            if (entry.Session.IsTurnInFlight) continue;
            if (Interlocked.Read(ref entry.LastActivityTicksUtc) > cutoffTicks) continue;

            if (_sessions.TryRemove(kvp.Key, out var removed))
            {
                try { await removed.Session.StopAsync(); }
                catch { /* already dead — leave it removed */ }
                (reaped ??= new()).Add(kvp.Key);
            }
        }

        return reaped ?? (IReadOnlyList<string>)Array.Empty<string>();
    }
}