ClaudeDo/src/ClaudeDo.Worker/Runner/StreamingClaudeSession.cs

using System.Text.Json;
using ClaudeDo.Worker.Runner.Interfaces;

namespace ClaudeDo.Worker.Runner;

public sealed class StreamingClaudeSession : ILiveSession
{
    private readonly IClaudeStreamTransport _transport;
    private readonly Func<string, Task> _onLine;
    private readonly ILogger<StreamingClaudeSession> _logger;

    private readonly SemaphoreSlim _sendLock = new(1, 1);
    private volatile bool _isTurnInFlight;
    private readonly Queue<string> _pending = new();

    public bool IsTurnInFlight => _isTurnInFlight;

    public StreamingClaudeSession(
        IClaudeStreamTransport transport,
        Func<string, Task> onLine,
        ILogger<StreamingClaudeSession> logger)
    {
        _transport = transport;
        _onLine = onLine;
        _logger = logger;
    }

    public async Task StartAsync(
        IReadOnlyList<string> args,
        string workingDirectory,
        string firstPrompt,
        CancellationToken ct)
    {
        _transport.LineReceived += HandleLineAsync;
        await _transport.StartAsync(args, workingDirectory, ct);
        await SendTurnAsync(firstPrompt, ct);
    }

    private async Task HandleLineAsync(string line)
    {
        try { await _onLine(line); }
        catch (Exception ex) { _logger.LogWarning(ex, "onLine callback threw"); }

        bool isResult;
        try
        {
            using var doc = JsonDocument.Parse(line);
            isResult = doc.RootElement.TryGetProperty("type", out var typeProp)
                && typeProp.GetString() == "result";
        }
        catch { isResult = false; }

        if (!isResult) return;

        // Turn ended — flush one queued message if available.
        await _sendLock.WaitAsync();
        try
        {
            _isTurnInFlight = false;
            if (_pending.Count > 0)
            {
                var next = _pending.Dequeue();
                await SendTurnAsync(next, CancellationToken.None);
            }
        }
        finally
        {
            _sendLock.Release();
        }
    }

    public async Task SendUserMessageAsync(string text, CancellationToken ct)
    {
        await _sendLock.WaitAsync(ct);
        try
        {
            if (_isTurnInFlight || _pending.Count > 0)
            {
                _pending.Enqueue(text);
            }
            else
            {
                await SendTurnAsync(text, ct);
            }
        }
        finally
        {
            _sendLock.Release();
        }
    }

    public async Task InterruptAsync(CancellationToken ct)
    {
        await _sendLock.WaitAsync(ct);
        try
        {
            if (!_isTurnInFlight) return;

            var requestId = Guid.NewGuid().ToString();
            var payload = JsonSerializer.Serialize(new
            {
                type = "control_request",
                request_id = requestId,
                request = new { subtype = "interrupt" }
            });

            try { await _transport.WriteLineAsync(payload, ct); }
            catch (Exception ex) { _logger.LogWarning(ex, "Failed to write interrupt control_request; degrading gracefully."); }
        }
        finally
        {
            _sendLock.Release();
        }
    }

    private async Task SendTurnAsync(string text, CancellationToken ct)
    {
        _isTurnInFlight = true;

        var payload = JsonSerializer.Serialize(new
        {
            type = "user",
            message = new
            {
                role = "user",
                content = new[]
                {
                    new { type = "text", text }
                }
            },
            parent_tool_use_id = (string?)null
        });

        await _transport.WriteLineAsync(payload, ct);
    }

    public async Task StopAsync()
    {
        _transport.Kill();
        await _transport.WaitForExitAsync();
    }

    public async ValueTask DisposeAsync()
    {
        await StopAsync();
        await _transport.DisposeAsync();
        _sendLock.Dispose();
    }
}