ClaudeDo/src/ClaudeDo.Worker/Runner/StreamAnalyzer.cs

using System.Linq;
using System.Text.Json;

namespace ClaudeDo.Worker.Runner;

public sealed class StreamResult
{
    public string? ResultMarkdown { get; set; }
    public string? StructuredOutputJson { get; set; }
    public string? SessionId { get; set; }
    public int TurnCount { get; set; }
    public int TokensIn { get; set; }
    public int TokensOut { get; set; }
    public int ApiRetryCount { get; set; }
    public IReadOnlyList<string> Blocks { get; set; } = Array.Empty<string>();
}

public sealed class StreamAnalyzer
{
    private string? _resultMarkdown;
    private string? _structuredOutputJson;
    private string? _sessionId;
    private int _turnCount;
    private int _tokensIn;
    private int _tokensOut;
    private int _apiRetryCount;
    private readonly List<string> _blocks = new();
    private const string BlockedPrefix = "CLAUDEDO_BLOCKED:";

    public void ProcessLine(string ndjsonLine)
    {
        if (string.IsNullOrWhiteSpace(ndjsonLine)) return;

        try
        {
            using var doc = JsonDocument.Parse(ndjsonLine);
            var root = doc.RootElement;

            if (!root.TryGetProperty("type", out var typeProp)) return;
            var type = typeProp.GetString();

            switch (type)
            {
                case "result":
                    if (root.TryGetProperty("result", out var resultProp))
                        _resultMarkdown = StripAndCollect(resultProp.GetString());
                    if (root.TryGetProperty("structured_output", out var structuredProp))
                        _structuredOutputJson = structuredProp.ToString();
                    if (root.TryGetProperty("session_id", out var sessionProp))
                        _sessionId = sessionProp.GetString();
                    // Authoritative token totals live on the result event.
                    if (root.TryGetProperty("usage", out var resultUsage))
                    {
                        if (resultUsage.TryGetProperty("input_tokens", out var inp))
                            _tokensIn = inp.GetInt32();
                        if (resultUsage.TryGetProperty("output_tokens", out var outp))
                            _tokensOut = outp.GetInt32();
                    }
                    break;

                case "assistant":
                    _turnCount++;
                    CollectFromAssistant(root);
                    break;

                case "system":
                    if (root.TryGetProperty("subtype", out var subtypeProp) &&
                        subtypeProp.GetString() == "api_retry")
                        _apiRetryCount++;
                    break;

                case "stream_event":
                    TryAccumulateUsage(root);
                    break;
            }
        }
        catch (JsonException) { /* Malformed JSON — skip */ }
    }

    public StreamResult GetResult() => new()
    {
        ResultMarkdown = FallbackResult(),
        StructuredOutputJson = _structuredOutputJson,
        SessionId = _sessionId,
        TurnCount = _turnCount,
        TokensIn = _tokensIn,
        TokensOut = _tokensOut,
        ApiRetryCount = _apiRetryCount,
        Blocks = _blocks.Distinct().ToList(),
    };

    private string? FallbackResult()
    {
        // A --json-schema run sometimes echoes the structured JSON as the result text; unwrap
        // that to its summary. A plain markdown result is kept verbatim (even when a separate
        // structured_output is also present). Only when there is no result text do we fall back
        // to the structured summary, then the raw JSON.
        if (SummaryFromJson(_resultMarkdown) is { } fromResult) return fromResult;
        if (!string.IsNullOrEmpty(_resultMarkdown)) return _resultMarkdown;
        if (SummaryFromJson(_structuredOutputJson) is { } fromStructured) return fromStructured;
        return _structuredOutputJson;
    }

    /// <summary>Extracts a non-empty string <c>summary</c> from a JSON object, or null when
    /// the input isn't such an object (a plain markdown result parses to null and is kept).</summary>
    private static string? SummaryFromJson(string? json)
    {
        if (string.IsNullOrWhiteSpace(json)) return null;
        try
        {
            using var doc = JsonDocument.Parse(json);
            if (doc.RootElement.ValueKind == JsonValueKind.Object &&
                doc.RootElement.TryGetProperty("summary", out var s) &&
                s.ValueKind == JsonValueKind.String)
            {
                var text = s.GetString();
                return string.IsNullOrWhiteSpace(text) ? null : text;
            }
        }
        catch { }
        return null;
    }

    private void CollectFromAssistant(JsonElement root)
    {
        if (!root.TryGetProperty("message", out var msg)) return;
        if (msg.ValueKind != JsonValueKind.Object) return;
        if (!msg.TryGetProperty("content", out var content) || content.ValueKind != JsonValueKind.Array) return;
        foreach (var block in content.EnumerateArray())
            if (block.TryGetProperty("type", out var t) && t.GetString() == "text"
                && block.TryGetProperty("text", out var txt))
                ScanForBlocks(txt.GetString());
    }

    private void ScanForBlocks(string? text)
    {
        if (string.IsNullOrEmpty(text)) return;
        foreach (var line in text.Split('\n'))
        {
            var trimmed = line.Trim();
            if (trimmed.StartsWith(BlockedPrefix, StringComparison.Ordinal))
                _blocks.Add(trimmed[BlockedPrefix.Length..].Trim());
        }
    }

    private string? StripAndCollect(string? text)
    {
        if (string.IsNullOrEmpty(text)) return text;
        ScanForBlocks(text);
        var kept = text.Split('\n')
            .Where(l => !l.Trim().StartsWith(BlockedPrefix, StringComparison.Ordinal));
        return string.Join('\n', kept).Trim();
    }

    private void TryAccumulateUsage(JsonElement root)
    {
        if (!root.TryGetProperty("event", out var eventProp)) return;
        if (eventProp.TryGetProperty("message", out var msgProp) &&
            msgProp.TryGetProperty("usage", out var usageProp))
        {
            if (usageProp.TryGetProperty("input_tokens", out var inp))
                _tokensIn += inp.GetInt32();
            if (usageProp.TryGetProperty("output_tokens", out var outp))
                _tokensOut += outp.GetInt32();
        }
    }
}