agent-framework/dotnet/src/Microsoft.Agents.AI.Foundry/ClientHeadersAgent.cs

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

using System.Collections.Generic;
using System.Runtime.CompilerServices;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Extensions.AI;

namespace Microsoft.Agents.AI.Foundry;

/// <summary>
/// Delegating <see cref="AIAgent"/> that captures any <c>x-client-*</c> headers stored on
/// <see cref="ChatClientAgentRunOptions.ChatOptions"/> by callers of
/// <see cref="ClientHeadersExtensions.WithClientHeader(ChatOptions, string, string)"/> and pushes
/// them onto a <see cref="ClientHeadersScope"/> for the lifetime of the run. The scope is read by
/// <see cref="ClientHeadersPolicy"/> inside the SCM transport pipeline and stamped onto the
/// outbound request.
/// </summary>
/// <remarks>
/// <para>
/// The decorator snapshots the header dictionary at scope-push time so concurrent runs that share
/// the same <see cref="ChatOptions"/> reference are isolated; mutating the source dictionary after
/// <c>RunAsync</c> begins does not leak into in-flight requests.
/// </para>
/// <para>
/// Streaming uses the async-iterator pattern so the AsyncLocal scope stays alive across yields,
/// which is required because the underlying HTTP send happens during enumeration.
/// </para>
/// </remarks>
internal sealed class ClientHeadersAgent : DelegatingAIAgent
{
    public ClientHeadersAgent(AIAgent innerAgent)
        : base(innerAgent)
    {
    }

    /// <inheritdoc/>
    protected override Task<AgentResponse> RunCoreAsync(
        IEnumerable<ChatMessage> messages,
        AgentSession? session = null,
        AgentRunOptions? options = null,
        CancellationToken cancellationToken = default)
    {
        var snapshot = TrySnapshot(options);
        if (snapshot is null)
        {
            return this.InnerAgent.RunAsync(messages, session, options, cancellationToken);
        }

        return RunAsyncCoreAsync(messages, session, options, snapshot, cancellationToken);

        async Task<AgentResponse> RunAsyncCoreAsync(
            IEnumerable<ChatMessage> innerMessages,
            AgentSession? innerSession,
            AgentRunOptions? innerOptions,
            Dictionary<string, string> innerSnapshot,
            CancellationToken innerCt)
        {
            using var _ = ClientHeadersScope.Push(innerSnapshot);
            return await this.InnerAgent.RunAsync(innerMessages, innerSession, innerOptions, innerCt).ConfigureAwait(false);
        }
    }

    /// <inheritdoc/>
    protected override async IAsyncEnumerable<AgentResponseUpdate> RunCoreStreamingAsync(
        IEnumerable<ChatMessage> messages,
        AgentSession? session = null,
        AgentRunOptions? options = null,
        [EnumeratorCancellation] CancellationToken cancellationToken = default)
    {
        var snapshot = TrySnapshot(options);
        using var _ = snapshot is null ? default : ClientHeadersScope.Push(snapshot);

        await foreach (var update in this.InnerAgent.RunStreamingAsync(messages, session, options, cancellationToken).ConfigureAwait(false))
        {
            yield return update;
        }
    }

    /// <summary>Reads the header dictionary stamped by <c>WithClientHeader(s)</c> and returns an immutable snapshot, or <see langword="null"/> if none.</summary>
    private static Dictionary<string, string>? TrySnapshot(AgentRunOptions? options)
    {
        if (options is not ChatClientAgentRunOptions { ChatOptions: { } chatOptions })
        {
            return null;
        }

        var headers = chatOptions.GetClientHeaders();
        if (headers is null || headers.Count == 0)
        {
            return null;
        }

        // Copy to defeat caller mutation after RunAsync starts.
        var copy = new Dictionary<string, string>(headers.Count, System.StringComparer.OrdinalIgnoreCase);
        foreach (var kvp in headers)
        {
            copy[kvp.Key] = kvp.Value;
        }

        return copy;
    }
}