.NET: Auto-wire ChatClient with OpenTelemetryChatClient in OpenTelemetryAgent

Agent-Logs-Url: https://github.com/microsoft/agent-framework/sessions/96dd033a-0c48-4d3f-9148-324bfd436b75

Co-authored-by: rogerbarreto <19890735+rogerbarreto@users.noreply.github.com>

dotnet/src/Microsoft.Agents.AI/OpenTelemetryAgent.cs

@@ -32,6 +32,13 @@ public sealed class OpenTelemetryAgent : DelegatingAIAgent, IDisposable
     private readonly OpenTelemetryChatClient _otelClient;
     /// <summary>The provider name extracted from <see cref="AIAgentMetadata"/>.</summary>
     private readonly string? _providerName;
+    /// <summary>The configured source name for telemetry. May be <see langword="null"/> to use the default.</summary>
+    private readonly string? _sourceName;
+    /// <summary>
+    /// Indicates whether the underlying <see cref="IChatClient"/> of a <see cref="ChatClientAgent"/> inner agent
+    /// should be automatically wrapped with <see cref="OpenTelemetryChatClient"/> on each invocation.
+    /// </summary>
+    private readonly bool _autoWireChatClient;
 
     /// <summary>Initializes a new instance of the <see cref="OpenTelemetryAgent"/> class.</summary>
     /// <param name="innerAgent">The underlying <see cref="AIAgent"/> to be augmented with telemetry capabilities.</param>
@@ -39,14 +46,22 @@ public sealed class OpenTelemetryAgent : DelegatingAIAgent, IDisposable
     /// An optional source name that will be used to identify telemetry data from this agent.
     /// If not provided, a default source name will be used for telemetry identification.
     /// </param>
+    /// <param name="autoWireChatClient">
+    /// When <see langword="true"/> (the default) and the inner agent is a <see cref="ChatClientAgent"/>, the underlying
+    /// <see cref="IChatClient"/> is automatically wrapped with <see cref="OpenTelemetryChatClient"/> for each invocation
+    /// so that chat-level telemetry flows alongside agent-level telemetry. If the underlying chat client is already
+    /// instrumented, no additional wrapping is applied. Set to <see langword="false"/> to opt-out of this behavior.
+    /// </param>
     /// <exception cref="ArgumentNullException"><paramref name="innerAgent"/> is <see langword="null"/>.</exception>
     /// <remarks>
     /// The constructor automatically extracts provider metadata from the inner agent and configures
     /// telemetry collection according to OpenTelemetry semantic conventions for AI systems.
     /// </remarks>
-    public OpenTelemetryAgent(AIAgent innerAgent, string? sourceName = null) : base(innerAgent)
+    public OpenTelemetryAgent(AIAgent innerAgent, string? sourceName = null, bool autoWireChatClient = true) : base(innerAgent)
     {
         this._providerName = innerAgent.GetService<AIAgentMetadata>()?.ProviderName;
+        this._sourceName = sourceName;
+        this._autoWireChatClient = autoWireChatClient;
 
         this._otelClient = new OpenTelemetryChatClient(
             new ForwardingChatClient(this),
@@ -163,6 +178,53 @@ public sealed class OpenTelemetryAgent : DelegatingAIAgent, IDisposable
         public Activity? CurrentActivity { get; }
     }
 
+    /// <summary>
+    /// If auto-wiring is enabled and the inner agent is a <see cref="ChatClientAgent"/> whose underlying
+    /// <see cref="IChatClient"/> is not already instrumented with <see cref="OpenTelemetryChatClient"/>, returns a
+    /// new <see cref="AgentRunOptions"/> with a <see cref="ChatClientAgentRunOptions.ChatClientFactory"/> that
+    /// wraps the chat client with <see cref="OpenTelemetryChatClient"/>. Otherwise, returns <paramref name="options"/> unchanged.
+    /// </summary>
+    private AgentRunOptions? GetRunOptionsWithChatClientWiring(AgentRunOptions? options)
+    {
+        if (!this._autoWireChatClient)
+        {
+            return options;
+        }
+
+        // The auto-wiring only applies when the inner agent is a ChatClientAgent. Otherwise, no-op.
+        if (this.InnerAgent.GetService<ChatClientAgent>() is null)
+        {
+            return options;
+        }
+
+        // Capture the underlying IChatClient and check whether it is already instrumented.
+        var chatClient = this.InnerAgent.GetService<IChatClient>();
+        if (chatClient is null || chatClient.GetService(typeof(OpenTelemetryChatClient)) is not null)
+        {
+            return options;
+        }
+
+        string? sourceName = this._sourceName;
+        IChatClient WrapWithOpenTelemetry(IChatClient cc) =>
+            cc.GetService(typeof(OpenTelemetryChatClient)) is not null
+                ? cc
+                : cc.AsBuilder().UseOpenTelemetry(sourceName: sourceName).Build();
+
+        if (options is ChatClientAgentRunOptions ccOptions)
+        {
+            // Don't mutate the caller's options; clone and chain any caller-provided factory.
+            var clone = (ChatClientAgentRunOptions)ccOptions.Clone();
+            var userFactory = clone.ChatClientFactory;
+            clone.ChatClientFactory = cc => WrapWithOpenTelemetry(userFactory is null ? cc : userFactory(cc));
+            return clone;
+        }
+
+        return new ChatClientAgentRunOptions
+        {
+            ChatClientFactory = WrapWithOpenTelemetry,
+        };
+    }
+
     /// <summary>The stub <see cref="IChatClient"/> used to delegate from the <see cref="OpenTelemetryChatClient"/> into the inner <see cref="AIAgent"/>.</summary>
     /// <param name="parentAgent"></param>
     private sealed class ForwardingChatClient(OpenTelemetryAgent parentAgent) : IChatClient
@@ -175,8 +237,11 @@ public sealed class OpenTelemetryAgent : DelegatingAIAgent, IDisposable
             // Update the current activity to reflect the agent invocation.
             parentAgent.UpdateCurrentActivity(fo?.CurrentActivity);
 
+            // If enabled, wire the underlying chat client with OpenTelemetryChatClient via ChatClientFactory.
+            var runOptions = parentAgent.GetRunOptionsWithChatClientWiring(fo?.Options);
+
             // Invoke the inner agent.
-            var response = await parentAgent.InnerAgent.RunAsync(messages, fo?.Session, fo?.Options, cancellationToken).ConfigureAwait(false);
+            var response = await parentAgent.InnerAgent.RunAsync(messages, fo?.Session, runOptions, cancellationToken).ConfigureAwait(false);
 
             // Wrap the response in a ChatResponse so we can pass it back through OpenTelemetryChatClient.
             return response.AsChatResponse();
@@ -190,8 +255,11 @@ public sealed class OpenTelemetryAgent : DelegatingAIAgent, IDisposable
             // Update the current activity to reflect the agent invocation.
             parentAgent.UpdateCurrentActivity(fo?.CurrentActivity);
 
+            // If enabled, wire the underlying chat client with OpenTelemetryChatClient via ChatClientFactory.
+            var runOptions = parentAgent.GetRunOptionsWithChatClientWiring(fo?.Options);
+
             // Invoke the inner agent.
-            await foreach (var update in parentAgent.InnerAgent.RunStreamingAsync(messages, fo?.Session, fo?.Options, cancellationToken).ConfigureAwait(false))
+            await foreach (var update in parentAgent.InnerAgent.RunStreamingAsync(messages, fo?.Session, runOptions, cancellationToken).ConfigureAwait(false))
             {
                 // Wrap the response updates in ChatResponseUpdates so we can pass them back through OpenTelemetryChatClient.
                 yield return update.AsChatResponseUpdate();

dotnet/src/Microsoft.Agents.AI/OpenTelemetryAgentBuilderExtensions.cs

@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft. All rights reserved.
 
 using System;
+using Microsoft.Extensions.AI;
 using Microsoft.Shared.Diagnostics;
 
 namespace Microsoft.Agents.AI;
@@ -48,9 +49,39 @@ public static class OpenTelemetryAgentBuilderExtensions
         this AIAgentBuilder builder,
         string? sourceName = null,
         Action<OpenTelemetryAgent>? configure = null) =>
+        builder.UseOpenTelemetry(autoWireChatClient: true, sourceName: sourceName, configure: configure);
+
+    /// <summary>
+    /// Adds OpenTelemetry instrumentation to the agent pipeline, with explicit control over whether the underlying
+    /// <see cref="IChatClient"/> of a <see cref="ChatClientAgent"/> inner agent is automatically wrapped with
+    /// <see cref="Microsoft.Extensions.AI.OpenTelemetryChatClient"/>.
+    /// </summary>
+    /// <param name="builder">The <see cref="AIAgentBuilder"/> to which OpenTelemetry support will be added.</param>
+    /// <param name="autoWireChatClient">
+    /// When <see langword="true"/>, the underlying <see cref="IChatClient"/> of a <see cref="ChatClientAgent"/> inner agent
+    /// is automatically wrapped with <see cref="Microsoft.Extensions.AI.OpenTelemetryChatClient"/> on each invocation so
+    /// that chat-level telemetry flows alongside agent-level telemetry. If the underlying chat client is already
+    /// instrumented, no additional wrapping is applied. If the inner agent is not a <see cref="ChatClientAgent"/>, no
+    /// auto-wiring is performed. Set to <see langword="false"/> to opt out of this behavior.
+    /// </param>
+    /// <param name="sourceName">
+    /// An optional source name that will be used to identify telemetry data from this agent.
+    /// If not specified, a default source name will be used.
+    /// </param>
+    /// <param name="configure">
+    /// An optional callback that provides additional configuration of the <see cref="OpenTelemetryAgent"/> instance.
+    /// This allows for fine-tuning telemetry behavior such as enabling sensitive data collection.
+    /// </param>
+    /// <returns>The <see cref="AIAgentBuilder"/> with OpenTelemetry instrumentation added, enabling method chaining.</returns>
+    /// <exception cref="ArgumentNullException"><paramref name="builder"/> is <see langword="null"/>.</exception>
+    public static AIAgentBuilder UseOpenTelemetry(
+        this AIAgentBuilder builder,
+        bool autoWireChatClient,
+        string? sourceName = null,
+        Action<OpenTelemetryAgent>? configure = null) =>
         Throw.IfNull(builder).Use((innerAgent, services) =>
         {
-            var agent = new OpenTelemetryAgent(innerAgent, sourceName);
+            var agent = new OpenTelemetryAgent(innerAgent, sourceName, autoWireChatClient);
             configure?.Invoke(agent);
 
             return agent;

dotnet/tests/Microsoft.Agents.AI.UnitTests/OpenTelemetryAgentBuilderExtensionsTests.cs

@@ -123,4 +123,59 @@ public class OpenTelemetryAgentBuilderExtensionsTests
         Assert.True(configureWasCalled);
         Assert.IsType<OpenTelemetryAgent>(result);
     }
+
+    /// <summary>
+    /// Verify that UseOpenTelemetry with autoWireChatClient parameter works correctly.
+    /// </summary>
+    [Theory]
+    [InlineData(true)]
+    [InlineData(false)]
+    public void UseOpenTelemetry_WithAutoWireChatClientFlag_ReturnsOpenTelemetryAgent(bool autoWireChatClient)
+    {
+        // Arrange
+        var mockAgent = new Mock<AIAgent>();
+        var builder = new AIAgentBuilder(mockAgent.Object);
+
+        // Act
+        var result = builder.UseOpenTelemetry(autoWireChatClient: autoWireChatClient).Build();
+
+        // Assert
+        Assert.IsType<OpenTelemetryAgent>(result);
+    }
+
+    /// <summary>
+    /// Verify that UseOpenTelemetry with autoWireChatClient and all parameters works correctly.
+    /// </summary>
+    [Fact]
+    public void UseOpenTelemetry_WithAutoWireChatClientAndAllParameters_CallsConfigureAction()
+    {
+        // Arrange
+        var mockAgent = new Mock<AIAgent>();
+        var builder = new AIAgentBuilder(mockAgent.Object);
+        var configureWasCalled = false;
+
+        // Act
+        var result = builder.UseOpenTelemetry(
+            autoWireChatClient: false,
+            sourceName: "TestSource",
+            configure: agent =>
+            {
+                configureWasCalled = true;
+                Assert.NotNull(agent);
+            }).Build();
+
+        // Assert
+        Assert.True(configureWasCalled);
+        Assert.IsType<OpenTelemetryAgent>(result);
+    }
+
+    /// <summary>
+    /// Verify that UseOpenTelemetry with autoWireChatClient throws ArgumentNullException when builder is null.
+    /// </summary>
+    [Fact]
+    public void UseOpenTelemetry_WithAutoWireChatClient_WithNullBuilder_ThrowsArgumentNullException()
+    {
+        // Act & Assert
+        Assert.Throws<ArgumentNullException>("builder", () => ((AIAgentBuilder)null!).UseOpenTelemetry(autoWireChatClient: true));
+    }
 }

dotnet/tests/Microsoft.Agents.AI.UnitTests/OpenTelemetryAgentTests.cs

@@ -627,4 +627,166 @@ public class OpenTelemetryAgentTests
     }
 
     private static string ReplaceWhitespace(string? input) => Regex.Replace(input ?? "", @"\s+", "").Trim();
+
+    #region AutoWireChatClient
+
+    [Fact]
+    public async Task AutoWireChatClient_DefaultsToEnabled_EmitsChatSpan_Async()
+    {
+        var sourceName = Guid.NewGuid().ToString();
+        var activities = new List<Activity>();
+        using var tracerProvider = OpenTelemetry.Sdk.CreateTracerProviderBuilder()
+            .AddSource(sourceName)
+            .AddInMemoryExporter(activities)
+            .Build();
+
+        var fakeChatClient = new AutoWireTestChatClient();
+        var inner = new ChatClientAgent(fakeChatClient);
+        using var agent = new OpenTelemetryAgent(inner, sourceName);
+
+        _ = await agent.RunAsync("hi");
+
+        // Expect 2 activities: the inner chat span (from auto-wired OpenTelemetryChatClient) and the invoke_agent span.
+        Assert.Equal(2, activities.Count);
+        Assert.Contains(activities, a => a.DisplayName.StartsWith("invoke_agent", StringComparison.Ordinal));
+        Assert.Contains(activities, a => string.Equals(a.GetTagItem("gen_ai.operation.name") as string, "chat", StringComparison.Ordinal));
+    }
+
+    [Fact]
+    public async Task AutoWireChatClient_Streaming_EmitsChatSpan_Async()
+    {
+        var sourceName = Guid.NewGuid().ToString();
+        var activities = new List<Activity>();
+        using var tracerProvider = OpenTelemetry.Sdk.CreateTracerProviderBuilder()
+            .AddSource(sourceName)
+            .AddInMemoryExporter(activities)
+            .Build();
+
+        var fakeChatClient = new AutoWireTestChatClient();
+        var inner = new ChatClientAgent(fakeChatClient);
+        using var agent = new OpenTelemetryAgent(inner, sourceName);
+
+        await foreach (var _ in agent.RunStreamingAsync("hi"))
+        {
+        }
+
+        Assert.Equal(2, activities.Count);
+        Assert.Contains(activities, a => a.DisplayName.StartsWith("invoke_agent", StringComparison.Ordinal));
+        Assert.Contains(activities, a => string.Equals(a.GetTagItem("gen_ai.operation.name") as string, "chat", StringComparison.Ordinal));
+    }
+
+    [Fact]
+    public async Task AutoWireChatClient_Disabled_DoesNotEmitChatSpan_Async()
+    {
+        var sourceName = Guid.NewGuid().ToString();
+        var activities = new List<Activity>();
+        using var tracerProvider = OpenTelemetry.Sdk.CreateTracerProviderBuilder()
+            .AddSource(sourceName)
+            .AddInMemoryExporter(activities)
+            .Build();
+
+        var fakeChatClient = new AutoWireTestChatClient();
+        var inner = new ChatClientAgent(fakeChatClient);
+        using var agent = new OpenTelemetryAgent(inner, sourceName, autoWireChatClient: false);
+
+        _ = await agent.RunAsync("hi");
+
+        // Only the invoke_agent activity should be emitted; no chat span.
+        var activity = Assert.Single(activities);
+        Assert.StartsWith("invoke_agent", activity.DisplayName);
+    }
+
+    [Fact]
+    public async Task AutoWireChatClient_NonChatClientAgent_NoOp_Async()
+    {
+        // Inner is not a ChatClientAgent — auto-wiring must be a no-op and options must remain null.
+        AgentRunOptions? observedOptions = null;
+        var inner = new TestAIAgent
+        {
+            RunAsyncFunc = (messages, session, options, ct) =>
+            {
+                observedOptions = options;
+                return Task.FromResult(new AgentResponse(new ChatMessage(ChatRole.Assistant, "ok")));
+            },
+        };
+
+        using var agent = new OpenTelemetryAgent(inner);
+
+        _ = await agent.RunAsync("hi");
+
+        Assert.Null(observedOptions);
+    }
+
+    [Fact]
+    public async Task AutoWireChatClient_AlreadyInstrumented_DoesNotDoubleWrap_Async()
+    {
+        var sourceName = Guid.NewGuid().ToString();
+        var activities = new List<Activity>();
+        using var tracerProvider = OpenTelemetry.Sdk.CreateTracerProviderBuilder()
+            .AddSource(sourceName)
+            .AddInMemoryExporter(activities)
+            .Build();
+
+        var fakeChatClient = new AutoWireTestChatClient();
+        // Pre-wrap with OpenTelemetryChatClient on the same source so spans flow through the tracer.
+        IChatClient preWrapped = fakeChatClient.AsBuilder().UseOpenTelemetry(sourceName: sourceName).Build();
+        var inner = new ChatClientAgent(preWrapped);
+        using var agent = new OpenTelemetryAgent(inner, sourceName);
+
+        _ = await agent.RunAsync("hi");
+
+        // Expect exactly 2 activities (one invoke_agent + one chat from the pre-existing wrapper). If we had double-wrapped, we would see 3.
+        Assert.Equal(2, activities.Count);
+    }
+
+    [Fact]
+    public async Task AutoWireChatClient_PreservesUserChatClientFactory_Async()
+    {
+        var sourceName = Guid.NewGuid().ToString();
+        var activities = new List<Activity>();
+        using var tracerProvider = OpenTelemetry.Sdk.CreateTracerProviderBuilder()
+            .AddSource(sourceName)
+            .AddInMemoryExporter(activities)
+            .Build();
+
+        bool userFactoryCalled = false;
+        var fakeChatClient = new AutoWireTestChatClient();
+        var inner = new ChatClientAgent(fakeChatClient);
+        using var agent = new OpenTelemetryAgent(inner, sourceName);
+
+        var runOptions = new ChatClientAgentRunOptions
+        {
+            ChatClientFactory = cc =>
+            {
+                userFactoryCalled = true;
+                return cc;
+            },
+        };
+
+        _ = await agent.RunAsync("hi", options: runOptions);
+
+        Assert.True(userFactoryCalled);
+        // Auto-wiring should still produce a chat span on top of the user's factory.
+        Assert.Equal(2, activities.Count);
+        Assert.Contains(activities, a => string.Equals(a.GetTagItem("gen_ai.operation.name") as string, "chat", StringComparison.Ordinal));
+    }
+
+    private sealed class AutoWireTestChatClient : IChatClient
+    {
+        public Task<ChatResponse> GetResponseAsync(IEnumerable<ChatMessage> messages, ChatOptions? options = null, CancellationToken cancellationToken = default) =>
+            Task.FromResult(new ChatResponse(new ChatMessage(ChatRole.Assistant, "ok")));
+
+        public async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseAsync(IEnumerable<ChatMessage> messages, ChatOptions? options = null, [EnumeratorCancellation] CancellationToken cancellationToken = default)
+        {
+            await Task.Yield();
+            yield return new ChatResponseUpdate(ChatRole.Assistant, "ok");
+        }
+
+        public object? GetService(Type serviceType, object? serviceKey = null) =>
+            serviceType?.IsInstanceOfType(this) == true ? this : null;
+
+        public void Dispose() { }
+    }
+
+    #endregion
 }