< Summary

Information
Class: NexusLabs.Needlr.AgentFramework.Workflows.Diagnostics.DiagnosticsAgentRunMiddleware
Assembly: NexusLabs.Needlr.AgentFramework.Workflows
File(s): /home/runner/work/needlr/needlr/src/NexusLabs.Needlr.AgentFramework.Workflows/Diagnostics/DiagnosticsAgentRunMiddleware.cs
Line coverage
96%
Covered lines: 96
Uncovered lines: 4
Coverable lines: 100
Total lines: 224
Line coverage: 96%
Branch coverage
64%
Covered branches: 48
Total branches: 74
Branch coverage: 64.8%
Method coverage

Feature is only available for sponsors

Upgrade to PRO version

Metrics

MethodBranch coverage Crap Score Cyclomatic complexity Line coverage
.ctor(...)100%11100%
HandleStreamingAsync()62.5%3232100%
HandleAsync()41.66%262485.18%
SynthesizeResponse(...)100%1818100%

File(s)

/home/runner/work/needlr/needlr/src/NexusLabs.Needlr.AgentFramework.Workflows/Diagnostics/DiagnosticsAgentRunMiddleware.cs

#LineLine coverage
 1using System.Diagnostics;
 2using System.Runtime.CompilerServices;
 3
 4using Microsoft.Agents.AI;
 5using Microsoft.Extensions.AI;
 6
 7using NexusLabs.Needlr.AgentFramework.Diagnostics;
 8
 9namespace NexusLabs.Needlr.AgentFramework.Workflows.Diagnostics;
 10
 11/// <summary>
 12/// Outermost middleware layer: wraps <c>agent.RunAsync()</c> and
 13/// <c>agent.RunStreamingAsync()</c> to capture per-run diagnostics including
 14/// total duration, message counts, and success/failure state. Emits
 15/// <see cref="IAgentMetrics"/> counters on start and completion.
 16/// </summary>
 17/// <remarks>
 18/// Both the non-streaming and streaming paths produce equivalent
 19/// <see cref="IAgentRunDiagnostics"/> via <see cref="IAgentDiagnosticsWriter.Set"/>.
 20/// </remarks>
 21internal sealed class DiagnosticsAgentRunMiddleware
 22{
 23    private readonly string _agentName;
 24    private readonly IAgentDiagnosticsWriter _writer;
 25    private readonly IAgentMetrics _metrics;
 26
 4627    internal DiagnosticsAgentRunMiddleware(
 4628        string agentName,
 4629        IAgentDiagnosticsWriter writer,
 4630        IAgentMetrics metrics)
 31    {
 4632        _agentName = agentName;
 4633        _writer = writer;
 4634        _metrics = metrics;
 4635    }
 36
 37    internal async IAsyncEnumerable<AgentResponseUpdate> HandleStreamingAsync(
 38        IEnumerable<ChatMessage> messages,
 39        AgentSession? session,
 40        AgentRunOptions? options,
 41        AIAgent innerAgent,
 42        [EnumeratorCancellation] CancellationToken cancellationToken)
 43    {
 2544        var resolvedName = !string.IsNullOrEmpty(innerAgent.Name) ? innerAgent.Name : _agentName;
 45
 2546        _metrics.RecordRunStarted(resolvedName);
 2547        using var activity = _metrics.ActivitySource.StartActivity($"agent.run {resolvedName}", ActivityKind.Internal);
 2548        activity?.SetTag("gen_ai.agent.name", resolvedName);
 2549        activity?.SetTag("gen_ai.agent.streaming", true);
 50
 2551        using var builder = AgentRunDiagnosticsBuilder.StartNew(resolvedName);
 52
 2553        var messageList = messages as IReadOnlyList<ChatMessage> ?? messages.ToList();
 2554        builder.RecordInputMessageCount(messageList.Count);
 2555        builder.RecordInputMessages(messageList);
 56
 2557        var messageIds = new HashSet<string>(StringComparer.Ordinal);
 2558        var accumulated = new List<AgentResponseUpdate>();
 2559        Exception? failure = null;
 60
 2561        var enumerator = innerAgent
 2562            .RunStreamingAsync(messageList, session, options, cancellationToken)
 2563            .GetAsyncEnumerator(cancellationToken);
 64        try
 65        {
 66            while (true)
 67            {
 68                AgentResponseUpdate update;
 69                try
 70                {
 5871                    if (!await enumerator.MoveNextAsync().ConfigureAwait(false))
 72                    {
 2273                        break;
 74                    }
 3375                    update = enumerator.Current;
 3376                }
 377                catch (Exception ex)
 78                {
 379                    failure = ex;
 380                    break;
 81                }
 82
 3383                if (!string.IsNullOrEmpty(update.MessageId))
 84                {
 1285                    messageIds.Add(update.MessageId);
 86                }
 3387                accumulated.Add(update);
 88
 3389                yield return update;
 90            }
 91        }
 92        finally
 93        {
 2594            await enumerator.DisposeAsync().ConfigureAwait(false);
 95        }
 96
 2597        builder.RecordOutputMessageCount(messageIds.Count);
 2598        builder.RecordOutputResponse(SynthesizeResponse(accumulated));
 99
 25100        if (failure is not null)
 101        {
 3102            builder.RecordFailure(failure.Message);
 3103            activity?.SetStatus(ActivityStatusCode.Error, failure.Message);
 104        }
 105
 25106        var diagnostics = builder.Build();
 25107        _writer.Set(diagnostics);
 25108        _metrics.RecordRunCompleted(diagnostics);
 109
 25110        activity?.SetTag("status", diagnostics.Succeeded ? "success" : "failed");
 25111        activity?.SetTag("gen_ai.usage.input_tokens", diagnostics.AggregateTokenUsage.InputTokens);
 25112        activity?.SetTag("gen_ai.usage.output_tokens", diagnostics.AggregateTokenUsage.OutputTokens);
 25113        activity?.SetTag("gen_ai.usage.total_tokens", diagnostics.AggregateTokenUsage.TotalTokens);
 25114        activity?.SetTag("gen_ai.usage.cached_input_tokens", diagnostics.AggregateTokenUsage.CachedInputTokens);
 25115        activity?.SetTag("gen_ai.usage.reasoning_tokens", diagnostics.AggregateTokenUsage.ReasoningTokens);
 116
 25117        if (failure is not null)
 118        {
 3119            throw failure;
 120        }
 22121    }
 122
 123    internal async Task<AgentResponse> HandleAsync(
 124        IEnumerable<ChatMessage> messages,
 125        AgentSession? session,
 126        AgentRunOptions? options,
 127        AIAgent innerAgent,
 128        CancellationToken cancellationToken)
 129    {
 130        // Resolve the agent name at runtime from the inner agent. The plugin creates
 131        // this middleware before the agent is fully built, so the name passed at
 132        // construction time is a fallback.
 21133        var resolvedName = !string.IsNullOrEmpty(innerAgent.Name) ? innerAgent.Name : _agentName;
 134
 21135        _metrics.RecordRunStarted(resolvedName);
 21136        using var activity = _metrics.ActivitySource.StartActivity($"agent.run {resolvedName}", ActivityKind.Internal);
 21137        activity?.SetTag("gen_ai.agent.name", resolvedName);
 138
 21139        using var builder = AgentRunDiagnosticsBuilder.StartNew(resolvedName);
 140
 141        try
 142        {
 21143            var messageList = messages as IReadOnlyList<ChatMessage> ?? messages.ToList();
 21144            builder.RecordInputMessageCount(messageList.Count);
 21145            builder.RecordInputMessages(messageList);
 146
 21147            var response = await innerAgent.RunAsync(messageList, session, options, cancellationToken)
 21148                .ConfigureAwait(false);
 149
 21150            builder.RecordOutputMessageCount(response.Messages?.Count ?? 0);
 21151            builder.RecordOutputResponse(response);
 152
 21153            return response;
 154        }
 0155        catch (Exception ex)
 156        {
 0157            builder.RecordFailure(ex.Message);
 0158            activity?.SetStatus(ActivityStatusCode.Error, ex.Message);
 0159            throw;
 160        }
 161        finally
 162        {
 21163            var diagnostics = builder.Build();
 21164            _writer.Set(diagnostics);
 21165            _metrics.RecordRunCompleted(diagnostics);
 166
 21167            activity?.SetTag("status", diagnostics.Succeeded ? "success" : "failed");
 21168            activity?.SetTag("gen_ai.usage.input_tokens", diagnostics.AggregateTokenUsage.InputTokens);
 21169            activity?.SetTag("gen_ai.usage.output_tokens", diagnostics.AggregateTokenUsage.OutputTokens);
 21170            activity?.SetTag("gen_ai.usage.total_tokens", diagnostics.AggregateTokenUsage.TotalTokens);
 21171            activity?.SetTag("gen_ai.usage.cached_input_tokens", diagnostics.AggregateTokenUsage.CachedInputTokens);
 21172            activity?.SetTag("gen_ai.usage.reasoning_tokens", diagnostics.AggregateTokenUsage.ReasoningTokens);
 173        }
 21174    }
 175
 176    /// <summary>
 177    /// Synthesizes an <see cref="AgentResponse"/> from the raw stream of
 178    /// <see cref="AgentResponseUpdate"/> items observed during a streaming run.
 179    /// Groups contents by <c>MessageId</c> so each logical message becomes one
 180    /// <see cref="ChatMessage"/>. Updates with no <c>MessageId</c> are grouped
 181    /// positionally so partial streams (mid-failure) still capture what was
 182    /// observed. Returns <see langword="null"/> when no updates were received.
 183    /// </summary>
 184    private static AgentResponse? SynthesizeResponse(List<AgentResponseUpdate> updates)
 185    {
 25186        if (updates.Count == 0)
 187        {
 1188            return null;
 189        }
 190
 24191        var order = new List<string>();
 24192        var groups = new Dictionary<string, (ChatRole Role, List<AIContent> Contents, string? AuthorName)>(StringCompare
 193
 114194        for (var i = 0; i < updates.Count; i++)
 195        {
 33196            var u = updates[i];
 33197            var key = !string.IsNullOrEmpty(u.MessageId) ? u.MessageId : $"__ordinal_{i}";
 33198            if (!groups.TryGetValue(key, out var entry))
 199            {
 29200                entry = (u.Role ?? ChatRole.Assistant, new List<AIContent>(), u.AuthorName);
 29201                groups[key] = entry;
 29202                order.Add(key);
 203            }
 33204            if (u.Contents is { Count: > 0 })
 205            {
 33206                entry.Contents.AddRange(u.Contents);
 207            }
 208        }
 209
 24210        var messages = new List<ChatMessage>(order.Count);
 106211        foreach (var k in order)
 212        {
 29213            var (role, contents, authorName) = groups[k];
 29214            var msg = new ChatMessage(role, contents);
 29215            if (!string.IsNullOrEmpty(authorName))
 216            {
 29217                msg.AuthorName = authorName;
 218            }
 29219            messages.Add(msg);
 220        }
 221
 24222        return new AgentResponse(messages);
 223    }
 224}