OpenTelemetryAgent.cs

// Copyright (c) Microsoft. All rights reserved.

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Diagnostics.CodeAnalysis;
using System.Runtime.CompilerServices;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Extensions.AI;
using Microsoft.Shared.DiagnosticIds;

namespace Microsoft.Agents.AI;

/// <summary>
/// Provides a delegating <see cref="AIAgent"/> implementation that implements the OpenTelemetry Semantic Conventions for Generative AI systems.
/// </summary>
/// <remarks>
/// This class provides an implementation of the Semantic Conventions for Generative AI systems v1.37, defined at <see href="https://opentelemetry.io/docs/specs/semconv/gen-ai/" />.
/// The specification is still experimental and subject to change; as such, the telemetry output by this client is also subject to change.
/// <para>
/// <strong>Security considerations:</strong> This class emits telemetry via the standard .NET
/// <c>System.Diagnostics</c> (<see cref="Activity"/>/metrics) APIs — it does not itself contact any
/// external system. Where that telemetry is sent (a local collector, a hosted observability backend,
/// etc.) is entirely determined by the OpenTelemetry exporters and pipeline the developer configures
/// outside of Agent Framework. By default, emitted telemetry is limited to metadata (e.g. token counts,
/// operation names, durations) and does not include message content. Enabling
/// <see cref="EnableSensitiveData"/> is an explicit, separate opt-in that additionally emits raw chat
/// message content, function-call arguments, and function-call results — treat that data as sensitive
/// and ensure it is not sent to, or retained by, a telemetry backend you have not secured appropriately.
/// </para>
/// </remarks>
public sealed class OpenTelemetryAgent : DelegatingAIAgent, IDisposable
{
    // IMPLEMENTATION NOTE: The OpenTelemetryChatClient from Microsoft.Extensions.AI provides a full and up-to-date
    // implementationof the OpenTelemetry Semantic Conventions for Generative AI systems, specifically for the client
    // metrics and the chat span. But the chat span is almost identical to the invoke_agent span, just with invoke_agent
    // have a different value for the operation name and a few additional tags. To avoid needing to reimplement the
    // convention, then, and keep it up-to-date as the convention evolves, for now this implementation just delegates
    // to OpenTelemetryChatClient for the actual telemetry work. For RunAsync and RunStreamingAsync, it delegates to the
    // inner agent not directly but rather via OpenTelemetryChatClient, which wraps a ForwardingChatClient that in turn
    // calls back into the inner agent.

    /// <summary>
    /// Gets the default <see cref="ActivitySource"/> name used by <see cref="OpenTelemetryAgent"/> when no source
    /// name is supplied to the constructor.
    /// </summary>
    /// <remarks>
    /// Pass this value to the tracing pipeline (for example, <c>TracerProviderBuilder.AddSource</c>) to subscribe to
    /// the spans emitted by agents that use the default source name, instead of hardcoding the literal name. This is
    /// a property rather than a constant so that the value is read at run time: a consumer that upgrades the package
    /// picks up the current source name without recompiling.
    /// </remarks>
    [Experimental(DiagnosticIds.Experiments.AgentsAIExperiments)]
    public static string DefaultSourceName => OpenTelemetryConsts.DefaultSourceName;

    /// <summary>The <see cref="OpenTelemetryChatClient"/> providing the bulk of the telemetry.</summary>
    private readonly OpenTelemetryChatClient _otelClient;
    /// <summary>The provider name extracted from <see cref="AIAgentMetadata"/>.</summary>
    private readonly string? _providerName;
    /// <summary>The resolved source name for telemetry. Always non-empty; defaults to <see cref="OpenTelemetryConsts.DefaultSourceName"/>.</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>
    /// The auto-wired below-FICC telemetry slot, when one was activated. Cached so that updates to
    /// <see cref="EnableSensitiveData"/> made after construction can be propagated to it.
    /// </summary>
    private DeferredOpenTelemetryChatClient? _innerTelemetrySlot;

    /// <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>
    /// <param name="sourceName">
    /// An optional source name used to identify telemetry data from this agent.
    /// When specified, register the same value with <c>TracerProviderBuilder.AddSource</c> so the tracing pipeline
    /// subscribes to these spans. When omitted, <see cref="DefaultSourceName"/> is used.
    /// </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)
#pragma warning disable MAAI001 // Auto-wiring is the new default; the experimental opt-out lives on the 3-arg overload.
        : this(innerAgent, sourceName, autoWireChatClient: true)
#pragma warning restore MAAI001
    {
    }

    /// <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>
    /// <param name="sourceName">
    /// An optional source name used to identify telemetry data from this agent.
    /// When specified, register the same value with <c>TracerProviderBuilder.AddSource</c> so the tracing pipeline
    /// subscribes to these spans. When omitted, <see cref="DefaultSourceName"/> is used.
    /// </param>
    /// <param name="autoWireChatClient">
    /// When <see langword="true"/> 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>
    [Experimental(DiagnosticIds.Experiments.AgentsAIExperiments)]
    public OpenTelemetryAgent(AIAgent innerAgent, string? sourceName, bool autoWireChatClient) : base(innerAgent)
    {
        this._providerName = innerAgent.GetService<AIAgentMetadata>()?.ProviderName;

        // Resolve once so the outer OpenTelemetryChatClient and the auto-wired inner
        // OpenTelemetryChatClient always emit spans under the same ActivitySource, even when
        // the caller passes "" or whitespace (which neither client should treat as a real source).
        this._sourceName = string.IsNullOrWhiteSpace(sourceName) ? OpenTelemetryConsts.DefaultSourceName : sourceName!;
        this._autoWireChatClient = autoWireChatClient;

        this._otelClient = new OpenTelemetryChatClient(
            new ForwardingChatClient(this),
            sourceName: this._sourceName);

        this.TryActivateInnerChatClientTelemetry();
    }

    /// <inheritdoc/>
    public void Dispose() => this._otelClient.Dispose();

    /// <summary>
    /// Gets or sets a value indicating whether potentially sensitive information should be included in telemetry.
    /// </summary>
    /// <value>
    /// <see langword="true"/> if potentially sensitive information should be included in telemetry;
    /// <see langword="false"/> if telemetry shouldn't include raw inputs and outputs.
    /// The default value is <see langword="false"/>, unless the <c>OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT</c>
    /// environment variable is set to "true" (case-insensitive).
    /// </value>
    /// <remarks>
    /// By default, telemetry includes metadata, such as token counts, but not raw inputs
    /// and outputs, such as message content, function call arguments, and function call results.
    /// The default value can be overridden by setting the <c>OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT</c>
    /// environment variable to "true". Explicitly setting this property will override the environment variable.
    /// <para>
    /// <strong>Security consideration:</strong> When sensitive data capture is enabled, the full text of chat messages —
    /// including user inputs, LLM responses, function call arguments, and function results — is emitted as telemetry.
    /// This data may contain PII or other sensitive information. Ensure that your telemetry pipeline is configured
    /// with appropriate access controls and data retention policies.
    /// </para>
    /// </remarks>
    public bool EnableSensitiveData
    {
        get => this._otelClient.EnableSensitiveData;
        set
        {
            this._otelClient.EnableSensitiveData = value;

            // Keep the auto-wired below-FICC slot in sync so its chat span captures message content too.
            if (this._innerTelemetrySlot is { } slot)
            {
                slot.EnableSensitiveData = value;
            }
        }
    }

    /// <inheritdoc/>
    protected override async Task<AgentResponse> RunCoreAsync(
        IEnumerable<ChatMessage> messages, AgentSession? session = null, AgentRunOptions? options = null, CancellationToken cancellationToken = default)
    {
        ChatOptions co = new ForwardedOptions(options, session, Activity.Current);

        var response = await this._otelClient.GetResponseAsync(messages, co, cancellationToken).ConfigureAwait(false);

        return response.RawRepresentation as AgentResponse ?? new AgentResponse(response);
    }

    /// <inheritdoc/>
    protected override async IAsyncEnumerable<AgentResponseUpdate> RunCoreStreamingAsync(
        IEnumerable<ChatMessage> messages, AgentSession? session = null, AgentRunOptions? options = null, [EnumeratorCancellation] CancellationToken cancellationToken = default)
    {
        ChatOptions co = new ForwardedOptions(options, session, Activity.Current);

        await foreach (var update in this._otelClient.GetStreamingResponseAsync(messages, co, cancellationToken).ConfigureAwait(false))
        {
            yield return update.RawRepresentation as AgentResponseUpdate ?? new AgentResponseUpdate(update);
        }
    }

    /// <summary>Augments the current activity created by the <see cref="OpenTelemetryChatClient"/> with agent-specific information.</summary>
    /// <param name="previousActivity">The <see cref="Activity"/> that was current prior to the <see cref="OpenTelemetryChatClient"/>'s invocation.</param>
    private void UpdateCurrentActivity(Activity? previousActivity)
    {
        // If there isn't a current activity to augment, or it's the same one that was current when the agent was invoked (meaning
        // the OpenTelemetryChatClient didn't create one), then there's nothing to do.
        if (Activity.Current is not { } activity ||
            ReferenceEquals(activity, previousActivity))
        {
            return;
        }

        // Override information set by OpenTelemetryChatClient to make it specific to invoke_agent.

        activity.DisplayName = string.IsNullOrWhiteSpace(this.Name)
            ? $"{OpenTelemetryConsts.GenAI.InvokeAgent} {this.Id}"
            : $"{OpenTelemetryConsts.GenAI.InvokeAgent} {this.Name}({this.Id})";
        activity.SetTag(OpenTelemetryConsts.GenAI.Operation.Name, OpenTelemetryConsts.GenAI.InvokeAgent);

        if (!string.IsNullOrWhiteSpace(this._providerName))
        {
            _ = activity.SetTag(OpenTelemetryConsts.GenAI.Provider.Name, this._providerName);
        }

        // Further augment the activity with agent-specific tags.

        _ = activity.SetTag(OpenTelemetryConsts.GenAI.Agent.Id, this.Id);

        if (this.Name is { } name && !string.IsNullOrWhiteSpace(name))
        {
            _ = activity.SetTag(OpenTelemetryConsts.GenAI.Agent.Name, this.Name);
        }

        if (this.Description is { } description && !string.IsNullOrWhiteSpace(description))
        {
            _ = activity.SetTag(OpenTelemetryConsts.GenAI.Agent.Description, description);
        }
    }

    /// <summary>State passed from this instance into the inner agent, circumventing the intermediate <see cref="OpenTelemetryChatClient"/>.</summary>
    private sealed class ForwardedOptions : ChatOptions
    {
        public ForwardedOptions(AgentRunOptions? options, AgentSession? session, Activity? currentActivity) :
            base((options as ChatClientAgentRunOptions)?.ChatOptions)
        {
            this.Options = options;
            this.Session = session;
            this.CurrentActivity = currentActivity;
        }

        public AgentRunOptions? Options { get; }

        public AgentSession? Session { get; }

        public Activity? CurrentActivity { get; }
    }

    /// <summary>
    /// When auto-wiring is enabled and the inner agent is a <see cref="ChatClientAgent"/> whose underlying
    /// <see cref="IChatClient"/> is not already instrumented, activates the in-place
    /// <see cref="DeferredOpenTelemetryChatClient"/> slot so that chat spans are emitted below the
    /// <see cref="FunctionInvokingChatClient"/> under this agent's source name. Positioning OpenTelemetry below FICC
    /// is what allows FICC to emit execute_tool spans on the agent source. Respects
    /// <see cref="ChatClientAgentOptions.UseProvidedChatClientAsIs"/> and is a no-op when no slot is reachable.
    /// </summary>
    private void TryActivateInnerChatClientTelemetry()
    {
        if (!this._autoWireChatClient)
        {
            return;
        }

        // Auto-wiring only applies when a ChatClientAgent is reachable from the inner agent. Otherwise, no-op.
        // Use GetService rather than a type check so wrapping agents that expose a nested ChatClientAgent are supported.
        var chatClientAgent = this.InnerAgent.GetService<ChatClientAgent>();
        if (chatClientAgent is null)
        {
            return;
        }

        // Respect ChatClientAgentOptions.UseProvidedChatClientAsIs: don't decorate the chat client when the user opted out.
        if (chatClientAgent.GetService<ChatClientAgentOptions>()?.UseProvidedChatClientAsIs is true)
        {
            return;
        }

        // Don't activate when the chat client is already instrumented (e.g. the caller added their own
        // OpenTelemetryChatClient), to avoid emitting duplicate chat spans.
        var chatClient = chatClientAgent.GetService<IChatClient>();
        if (chatClient is null || chatClient.GetService(typeof(OpenTelemetryChatClient)) is not null)
        {
            return;
        }

        // Activate the pre-placed slot in-place (below FICC) rather than wrapping a new OpenTelemetryChatClient
        // around the whole pipeline on each run. Cache it and seed its EnableSensitiveData from the current
        // value so a later change to this agent's EnableSensitiveData can be propagated to the inner chat span.
        if (chatClient.GetService(typeof(DeferredOpenTelemetryChatClient)) is DeferredOpenTelemetryChatClient slot)
        {
            slot.Activate(this._sourceName);
            slot.EnableSensitiveData = this.EnableSensitiveData;
            this._innerTelemetrySlot = slot;
        }
    }

    /// <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
    {
        public async Task<ChatResponse> GetResponseAsync(
            IEnumerable<ChatMessage> messages, ChatOptions? options = null, CancellationToken cancellationToken = default)
        {
            ForwardedOptions? fo = options as ForwardedOptions;

            // Update the current activity to reflect the agent invocation.
            parentAgent.UpdateCurrentActivity(fo?.CurrentActivity);

            // Invoke the inner agent. Chat-level telemetry is emitted by the in-place DeferredOpenTelemetryChatClient
            // slot (below FICC), activated once at construction; no per-run chat-client wiring is needed here.
            var response = await parentAgent.InnerAgent.RunAsync(messages, fo?.Session, fo?.Options, cancellationToken).ConfigureAwait(false);

            // Wrap the response in a ChatResponse so we can pass it back through OpenTelemetryChatClient.
            return response.AsChatResponse();
        }

        public async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseAsync(
            IEnumerable<ChatMessage> messages, ChatOptions? options = null, [EnumeratorCancellation] CancellationToken cancellationToken = default)
        {
            ForwardedOptions? fo = options as ForwardedOptions;

            // Update the current activity to reflect the agent invocation.
            parentAgent.UpdateCurrentActivity(fo?.CurrentActivity);

            // Invoke the inner agent. Chat-level telemetry is emitted by the in-place DeferredOpenTelemetryChatClient
            // slot (below FICC), activated once at construction; no per-run chat-client wiring is needed here.
            await foreach (var update in parentAgent.InnerAgent.RunStreamingAsync(messages, fo?.Session, fo?.Options, cancellationToken).ConfigureAwait(false))
            {
                // Wrap the response updates in ChatResponseUpdates so we can pass them back through OpenTelemetryChatClient.
                yield return update.AsChatResponseUpdate();
            }
        }

        public object? GetService(Type serviceType, object? serviceKey = null) =>
            // Delegate any inquiries made by the OpenTelemetryChatClient back to the parent agent.
            parentAgent.GetService(serviceType, serviceKey);

        public void Dispose() { }
    }
}