< Summary

Information
Class: NexusLabs.Needlr.AgentFramework.Diagnostics.DiagnosticsChatClientMiddleware
Assembly: NexusLabs.Needlr.AgentFramework
File(s): /home/runner/work/needlr/needlr/src/NexusLabs.Needlr.AgentFramework/Diagnostics/DiagnosticsChatClientMiddleware.cs
Line coverage
100%
Covered lines: 242
Uncovered lines: 0
Coverable lines: 242
Total lines: 404
Line coverage: 100%
Branch coverage
83%
Covered branches: 122
Total branches: 146
Branch coverage: 83.5%
Method coverage

Feature is only available for sponsors

Upgrade to PRO version

Metrics

MethodBranch coverage Crap Score Cyclomatic complexity Line coverage
.ctor(...)100%11100%
DrainCompletions()100%22100%
HandleAsync()79.41%6868100%
HandleStreamingAsync()85.29%6868100%
StartChatActivity(...)100%88100%

File(s)

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

#LineLine coverage
 1using System.Collections.Concurrent;
 2using System.Diagnostics;
 3using System.Runtime.CompilerServices;
 4
 5using Microsoft.Extensions.AI;
 6
 7using NexusLabs.Needlr.AgentFramework.Progress;
 8
 9namespace NexusLabs.Needlr.AgentFramework.Diagnostics;
 10
 11/// <summary>
 12/// Single writer for chat completion diagnostics. Wraps each
 13/// <c>IChatClient.GetResponseAsync()</c> call to capture per-completion timing,
 14/// token usage, and full request/response payloads. Records to the AsyncLocal
 15/// <see cref="AgentRunDiagnosticsBuilder"/> and a thread-safe collection (for
 16/// workflow runs where AsyncLocal doesn't propagate). Optionally emits
 17/// <see cref="LlmCallStartedEvent"/>/<see cref="LlmCallCompletedEvent"/> to the
 18/// progress reporter and OTel metrics via <see cref="IAgentMetrics"/>.
 19/// </summary>
 20/// <remarks>
 21/// <para>
 22/// <c>IterativeAgentLoop</c> wraps its chat client with this middleware
 23/// internally, making it the sole writer for <see cref="ChatCompletionDiagnostics"/>.
 24/// No other code should call <see cref="AgentRunDiagnosticsBuilder.AddChatCompletion"/>
 25/// for calls that pass through this middleware.
 26/// </para>
 27/// <para>
 28/// <see cref="IAgentMetrics"/> and <see cref="IProgressReporterAccessor"/> are optional.
 29/// When null, recording still occurs but OTel metrics and progress events are skipped.
 30/// </para>
 31/// </remarks>
 32[DoNotAutoRegister]
 33internal sealed class DiagnosticsChatClientMiddleware : IChatCompletionCollector
 34{
 35    private readonly IAgentMetrics? _metrics;
 36    private readonly IProgressReporterAccessor? _progressAccessor;
 37    private readonly ChatCompletionActivityMode _activityMode;
 17638    private readonly ConcurrentQueue<ChatCompletionDiagnostics> _allCompletions = new();
 39    private int _sequenceCounter;
 40
 17641    internal DiagnosticsChatClientMiddleware(
 17642        IAgentMetrics? metrics = null,
 17643        IProgressReporterAccessor? progressAccessor = null,
 17644        ChatCompletionActivityMode activityMode = ChatCompletionActivityMode.Always)
 45    {
 17646        _metrics = metrics;
 17647        _progressAccessor = progressAccessor;
 17648        _activityMode = activityMode;
 17649    }
 50
 51    /// <summary>
 52    /// Drains all captured completions since the last drain. Thread-safe.
 53    /// </summary>
 54    public IReadOnlyList<ChatCompletionDiagnostics> DrainCompletions()
 55    {
 2356        var results = new List<ChatCompletionDiagnostics>();
 4457        while (_allCompletions.TryDequeue(out var completion))
 58        {
 2159            results.Add(completion);
 2160        }
 2361        return results;
 62    }
 63
 64    internal async Task<ChatResponse> HandleAsync(
 65        IEnumerable<ChatMessage> messages,
 66        ChatOptions? options,
 67        IChatClient innerChatClient,
 68        CancellationToken cancellationToken)
 69    {
 28170        var builder = AgentRunDiagnosticsBuilder.GetCurrent();
 28171        var sequence = builder?.NextChatCompletionSequence()
 28172            ?? Interlocked.Increment(ref _sequenceCounter) - 1;
 28173        var startedAt = DateTimeOffset.UtcNow;
 28174        var stopwatch = Stopwatch.StartNew();
 75
 28176        var (ownedActivity, targetActivity) = StartChatActivity("agent.chat");
 28177        using var _ = ownedActivity;
 78
 28179        if (_progressAccessor is not null)
 80        {
 3581            _progressAccessor.Current.Report(new LlmCallStartedEvent(
 3582                Timestamp: startedAt,
 3583                WorkflowId: _progressAccessor.Current.WorkflowId,
 3584                AgentId: _progressAccessor.Current.AgentId,
 3585                ParentAgentId: builder?.ParentAgentName,
 3586                Depth: _progressAccessor.Current.Depth,
 3587                SequenceNumber: _progressAccessor.Current.NextSequence(),
 3588                CallSequence: sequence));
 89        }
 90
 91        try
 92        {
 28193            var response = await innerChatClient.GetResponseAsync(messages, options, cancellationToken)
 28194                .ConfigureAwait(false);
 95
 27296            stopwatch.Stop();
 97
 27298            var model = response.ModelId ?? "unknown";
 99
 272100            targetActivity?.SetTag("gen_ai.response.model", model);
 272101            targetActivity?.SetTag("agent.chat.sequence", sequence);
 272102            targetActivity?.SetTag("status", "success");
 103
 272104            _metrics?.RecordChatCompletion(model, stopwatch.Elapsed, succeeded: true, agentName: builder?.AgentName);
 105
 272106            var usage = response.Usage;
 272107            var tokens = new TokenUsage(
 272108                InputTokens: usage?.InputTokenCount ?? 0,
 272109                OutputTokens: usage?.OutputTokenCount ?? 0,
 272110                TotalTokens: usage?.TotalTokenCount ?? 0,
 272111                CachedInputTokens: usage?.AdditionalCounts?.GetValueOrDefault("CachedInputTokens") ?? 0,
 272112                ReasoningTokens: usage?.AdditionalCounts?.GetValueOrDefault("ReasoningTokens") ?? 0);
 113
 272114            targetActivity?.SetTag("gen_ai.usage.input_tokens", tokens.InputTokens);
 272115            targetActivity?.SetTag("gen_ai.usage.output_tokens", tokens.OutputTokens);
 116
 272117            var messageList = messages as ICollection<ChatMessage> ?? messages.ToList();
 118
 272119            var diagnostics = new ChatCompletionDiagnostics(
 272120                Sequence: sequence,
 272121                Model: model,
 272122                Tokens: tokens,
 272123                InputMessageCount: messageList.Count,
 272124                Duration: stopwatch.Elapsed,
 272125                Succeeded: true,
 272126                ErrorMessage: null,
 272127                StartedAt: startedAt,
 272128                CompletedAt: DateTimeOffset.UtcNow)
 272129            {
 272130                AgentName = builder?.AgentName,
 272131                RequestMessages = messageList as IReadOnlyList<ChatMessage> ?? messageList.ToList(),
 272132                Response = response,
 272133                RequestCharCount = DiagnosticsCharCounter.ChatMessagesLength(messageList as IReadOnlyList<ChatMessage> ?
 272134                ResponseCharCount = DiagnosticsCharCounter.ChatResponseLength(response),
 272135            };
 136
 272137            builder?.AddChatCompletion(diagnostics);
 272138            _allCompletions.Enqueue(diagnostics);
 139
 272140            if (_progressAccessor is not null)
 141            {
 34142                _progressAccessor.Current.Report(new LlmCallCompletedEvent(
 34143                    Timestamp: DateTimeOffset.UtcNow,
 34144                    WorkflowId: _progressAccessor.Current.WorkflowId,
 34145                    AgentId: _progressAccessor.Current.AgentId,
 34146                    ParentAgentId: builder?.ParentAgentName,
 34147                    Depth: _progressAccessor.Current.Depth,
 34148                    SequenceNumber: _progressAccessor.Current.NextSequence(),
 34149                    CallSequence: sequence,
 34150                    Model: model,
 34151                    Duration: stopwatch.Elapsed,
 34152                    InputTokens: tokens.InputTokens,
 34153                    OutputTokens: tokens.OutputTokens,
 34154                    TotalTokens: tokens.TotalTokens));
 155            }
 156
 272157            return response;
 158        }
 9159        catch (Exception ex)
 160        {
 9161            stopwatch.Stop();
 162
 9163            targetActivity?.SetStatus(ActivityStatusCode.Error, ex.Message);
 9164            targetActivity?.SetTag("status", "failed");
 165
 9166            _metrics?.RecordChatCompletion("unknown", stopwatch.Elapsed, succeeded: false, agentName: builder?.AgentName
 167
 9168            var failedMessageList = messages as IReadOnlyList<ChatMessage> ?? messages.ToList();
 169
 9170            var diagnostics = new ChatCompletionDiagnostics(
 9171                Sequence: sequence,
 9172                Model: "unknown",
 9173                Tokens: new TokenUsage(0, 0, 0, 0, 0),
 9174                InputMessageCount: 0,
 9175                Duration: stopwatch.Elapsed,
 9176                Succeeded: false,
 9177                ErrorMessage: ex.Message,
 9178                StartedAt: startedAt,
 9179                CompletedAt: DateTimeOffset.UtcNow)
 9180            {
 9181                AgentName = builder?.AgentName,
 9182                RequestMessages = failedMessageList,
 9183                RequestCharCount = DiagnosticsCharCounter.ChatMessagesLength(failedMessageList),
 9184            };
 185
 9186            builder?.AddChatCompletion(diagnostics);
 9187            _allCompletions.Enqueue(diagnostics);
 188
 9189            if (_progressAccessor is not null)
 190            {
 1191                _progressAccessor.Current.Report(new LlmCallFailedEvent(
 1192                    Timestamp: DateTimeOffset.UtcNow,
 1193                    WorkflowId: _progressAccessor.Current.WorkflowId,
 1194                    AgentId: _progressAccessor.Current.AgentId,
 1195                    ParentAgentId: builder?.ParentAgentName,
 1196                    Depth: _progressAccessor.Current.Depth,
 1197                    SequenceNumber: _progressAccessor.Current.NextSequence(),
 1198                    CallSequence: sequence,
 1199                    ErrorMessage: ex.Message,
 1200                    Duration: stopwatch.Elapsed));
 201            }
 202
 9203            throw;
 204        }
 272205    }
 206
 207    internal async IAsyncEnumerable<ChatResponseUpdate> HandleStreamingAsync(
 208        IEnumerable<ChatMessage> messages,
 209        ChatOptions? options,
 210        IChatClient innerChatClient,
 211        [EnumeratorCancellation] CancellationToken cancellationToken)
 212    {
 25213        var builder = AgentRunDiagnosticsBuilder.GetCurrent();
 25214        var sequence = builder?.NextChatCompletionSequence()
 25215            ?? Interlocked.Increment(ref _sequenceCounter) - 1;
 25216        var startedAt = DateTimeOffset.UtcNow;
 25217        var stopwatch = Stopwatch.StartNew();
 218
 25219        var (ownedStreamActivity, targetActivity) = StartChatActivity("agent.chat.stream");
 25220        using var _s = ownedStreamActivity;
 221
 25222        if (_progressAccessor is not null)
 223        {
 23224            _progressAccessor.Current.Report(new LlmCallStartedEvent(
 23225                Timestamp: startedAt,
 23226                WorkflowId: _progressAccessor.Current.WorkflowId,
 23227                AgentId: _progressAccessor.Current.AgentId,
 23228                ParentAgentId: builder?.ParentAgentName,
 23229                Depth: _progressAccessor.Current.Depth,
 23230                SequenceNumber: _progressAccessor.Current.NextSequence(),
 23231                CallSequence: sequence));
 232        }
 233
 25234        var messageList = messages as IReadOnlyList<ChatMessage> ?? messages.ToList();
 25235        var buffered = new List<ChatResponseUpdate>();
 25236        Exception? failure = null;
 237
 25238        var enumerable = innerChatClient.GetStreamingResponseAsync(messages, options, cancellationToken);
 24239        var enumerator = enumerable.GetAsyncEnumerator(cancellationToken);
 240
 241        try
 242        {
 243            while (true)
 244            {
 245                ChatResponseUpdate update;
 246                try
 247                {
 54248                    if (!await enumerator.MoveNextAsync().ConfigureAwait(false))
 249                    {
 22250                        break;
 251                    }
 30252                    update = enumerator.Current;
 30253                }
 2254                catch (Exception ex)
 255                {
 2256                    failure = ex;
 2257                    break;
 258                }
 259
 30260                buffered.Add(update);
 30261                yield return update;
 262            }
 263        }
 264        finally
 265        {
 24266            await enumerator.DisposeAsync().ConfigureAwait(false);
 267        }
 268
 24269        stopwatch.Stop();
 270
 24271        var aggregated = buffered.ToChatResponse();
 272
 24273        if (failure is null)
 274        {
 22275            var model = aggregated.ModelId ?? "unknown";
 276
 22277            targetActivity?.SetTag("gen_ai.response.model", model);
 22278            targetActivity?.SetTag("agent.chat.sequence", sequence);
 22279            targetActivity?.SetTag("status", "success");
 280
 22281            _metrics?.RecordChatCompletion(model, stopwatch.Elapsed, succeeded: true, agentName: builder?.AgentName);
 282
 22283            var usage = aggregated.Usage;
 22284            var tokens = new TokenUsage(
 22285                InputTokens: usage?.InputTokenCount ?? 0,
 22286                OutputTokens: usage?.OutputTokenCount ?? 0,
 22287                TotalTokens: usage?.TotalTokenCount ?? 0,
 22288                CachedInputTokens: usage?.AdditionalCounts?.GetValueOrDefault("CachedInputTokens") ?? 0,
 22289                ReasoningTokens: usage?.AdditionalCounts?.GetValueOrDefault("ReasoningTokens") ?? 0);
 290
 22291            targetActivity?.SetTag("gen_ai.usage.input_tokens", tokens.InputTokens);
 22292            targetActivity?.SetTag("gen_ai.usage.output_tokens", tokens.OutputTokens);
 293
 22294            var diagnostics = new ChatCompletionDiagnostics(
 22295                Sequence: sequence,
 22296                Model: model,
 22297                Tokens: tokens,
 22298                InputMessageCount: messageList.Count,
 22299                Duration: stopwatch.Elapsed,
 22300                Succeeded: true,
 22301                ErrorMessage: null,
 22302                StartedAt: startedAt,
 22303                CompletedAt: DateTimeOffset.UtcNow)
 22304            {
 22305                AgentName = builder?.AgentName,
 22306                RequestMessages = messageList,
 22307                Response = aggregated,
 22308                RequestCharCount = DiagnosticsCharCounter.ChatMessagesLength(messageList),
 22309                ResponseCharCount = DiagnosticsCharCounter.ChatResponseLength(aggregated),
 22310            };
 311
 22312            builder?.AddChatCompletion(diagnostics);
 22313            _allCompletions.Enqueue(diagnostics);
 314
 22315            if (_progressAccessor is not null)
 316            {
 20317                _progressAccessor.Current.Report(new LlmCallCompletedEvent(
 20318                    Timestamp: DateTimeOffset.UtcNow,
 20319                    WorkflowId: _progressAccessor.Current.WorkflowId,
 20320                    AgentId: _progressAccessor.Current.AgentId,
 20321                    ParentAgentId: builder?.ParentAgentName,
 20322                    Depth: _progressAccessor.Current.Depth,
 20323                    SequenceNumber: _progressAccessor.Current.NextSequence(),
 20324                    CallSequence: sequence,
 20325                    Model: model,
 20326                    Duration: stopwatch.Elapsed,
 20327                    InputTokens: tokens.InputTokens,
 20328                    OutputTokens: tokens.OutputTokens,
 20329                    TotalTokens: tokens.TotalTokens));
 330            }
 331        }
 332        else
 333        {
 2334            targetActivity?.SetStatus(ActivityStatusCode.Error, failure.Message);
 2335            targetActivity?.SetTag("status", "failed");
 336
 2337            _metrics?.RecordChatCompletion("unknown", stopwatch.Elapsed, succeeded: false, agentName: builder?.AgentName
 338
 2339            var diagnostics = new ChatCompletionDiagnostics(
 2340                Sequence: sequence,
 2341                Model: aggregated.ModelId ?? "unknown",
 2342                Tokens: new TokenUsage(0, 0, 0, 0, 0),
 2343                InputMessageCount: 0,
 2344                Duration: stopwatch.Elapsed,
 2345                Succeeded: false,
 2346                ErrorMessage: failure.Message,
 2347                StartedAt: startedAt,
 2348                CompletedAt: DateTimeOffset.UtcNow)
 2349            {
 2350                AgentName = builder?.AgentName,
 2351                RequestMessages = messageList,
 2352                Response = aggregated,
 2353                RequestCharCount = DiagnosticsCharCounter.ChatMessagesLength(messageList),
 2354                ResponseCharCount = DiagnosticsCharCounter.ChatResponseLength(aggregated),
 2355            };
 356
 2357            builder?.AddChatCompletion(diagnostics);
 2358            _allCompletions.Enqueue(diagnostics);
 359
 2360            if (_progressAccessor is not null)
 361            {
 2362                _progressAccessor.Current.Report(new LlmCallFailedEvent(
 2363                    Timestamp: DateTimeOffset.UtcNow,
 2364                    WorkflowId: _progressAccessor.Current.WorkflowId,
 2365                    AgentId: _progressAccessor.Current.AgentId,
 2366                    ParentAgentId: builder?.ParentAgentName,
 2367                    Depth: _progressAccessor.Current.Depth,
 2368                    SequenceNumber: _progressAccessor.Current.NextSequence(),
 2369                    CallSequence: sequence,
 2370                    ErrorMessage: failure.Message,
 2371                    Duration: stopwatch.Elapsed));
 372            }
 373
 2374            throw failure;
 375        }
 22376    }
 377
 378    /// <summary>
 379    /// Creates a chat completion activity respecting <see cref="_activityMode"/>.
 380    /// When <see cref="ChatCompletionActivityMode.EnrichParent"/> is active and a
 381    /// parent <c>gen_ai.*</c> activity exists, returns <c>created = null</c> and
 382    /// <c>target = parent</c> so callers enrich the parent span without creating a
 383    /// duplicate child. The caller must only dispose <c>created</c>, never <c>target</c>.
 384    /// </summary>
 385    private (Activity? Created, Activity? Target) StartChatActivity(string operationName)
 386    {
 306387        if (_metrics is null)
 388        {
 231389            return (null, null);
 390        }
 391
 75392        if (_activityMode == ChatCompletionActivityMode.EnrichParent)
 393        {
 7394            var parent = Activity.Current;
 7395            if (parent?.OperationName.StartsWith("gen_ai.", StringComparison.Ordinal) == true)
 396            {
 5397                return (Created: null, Target: parent);
 398            }
 399        }
 400
 70401        var created = _metrics.ActivitySource.StartActivity(operationName, ActivityKind.Client);
 70402        return (Created: created, Target: created);
 403    }
 404}
Generated by: ReportGenerator 5.5.9.0
05/04/2026 - 06:46:16
GitHub | reportgenerator.io

Methods/Properties

.ctor(NexusLabs.Needlr.AgentFramework.Diagnostics.IAgentMetrics,NexusLabs.Needlr.AgentFramework.Progress.IProgressReporterAccessor,NexusLabs.Needlr.AgentFramework.Diagnostics.ChatCompletionActivityMode)
DrainCompletions()
HandleAsync()
HandleStreamingAsync()
StartChatActivity(System.String)