Address PR review: pre-wire AsAIAgent path and dedup TryApplyUserAgent

* FoundryAgent: extract WireClientHeaders helper and call it from the
  internal (AIProjectClient, ChatClientAgent) constructor used by
  AzureAIProjectChatClientExtensions.AsAIAgent so those Foundry-built
  agents also pre-wire the x-client header pipeline.
* Foundry.Hosting TryApplyUserAgent: dedup HostedAgentUserAgentPolicy
  registration per OpenAIRequestPolicies instance via
  ConditionalWeakTable so per-request resolution does not grow the
  policy list unboundedly on singleton agents.
This commit is contained in:
Roger Barreto
2026-05-05 14:18:40 +01:00
Unverified
parent 9f31bb6795
commit a4c8f911b7
2 changed files with 39 additions and 23 deletions
@@ -3,6 +3,7 @@
using System;
using System.ClientModel.Primitives;
using System.Diagnostics.CodeAnalysis;
using System.Runtime.CompilerServices;
using Azure.AI.AgentServer.Responses;
using Azure.Core;
using Azure.Identity;
@@ -230,12 +231,21 @@ public static class FoundryHostingExtensions
var chatClient = agent.GetService<IChatClient>();
if (chatClient?.GetService<OpenAIRequestPolicies>() is { } policies)
{
// The HostedAgentUserAgentPolicy is idempotent on the wire (it skips when the
// supplement is already present in the User-Agent header), so we can register it
// unconditionally without dedup. MEAI's lock-free CAS-append on _entries is safe.
policies.AddPolicy(HostedAgentUserAgentPolicy.Instance, PipelinePosition.PerCall);
// Hosted agents are typically singletons resolved per request, so AddPolicy must be
// called at most once per OpenAIRequestPolicies instance to avoid unbounded growth of
// the policy list (each entry adds per-request CPU work even though the User-Agent
// value stays stable). Track which instances we have already wired with a
// ConditionalWeakTable keyed on the OpenAIRequestPolicies reference; the table holds
// weak references so it does not extend the lifetime of the chat client.
if (s_userAgentRegistrations.TryAdd(policies, s_boxedTrue))
{
policies.AddPolicy(HostedAgentUserAgentPolicy.Instance, PipelinePosition.PerCall);
}
}
return agent;
}
private static readonly object s_boxedTrue = new();
private static readonly ConditionalWeakTable<OpenAIRequestPolicies, object> s_userAgentRegistrations = new();
}
@@ -102,7 +102,7 @@ public sealed class FoundryAgent : DelegatingAIAgent
/// Internal constructor used by <c>AsAIAgent</c> extension methods that already have an <see cref="AIProjectClient"/> and a configured <see cref="ChatClientAgent"/>.
/// </summary>
internal FoundryAgent(AIProjectClient aiProjectClient, ChatClientAgent innerAgent)
: base(Throw.IfNull(innerAgent))
: base(WireClientHeaders(Throw.IfNull(innerAgent)))
{
this._aiProjectClient = Throw.IfNull(aiProjectClient);
}
@@ -166,7 +166,7 @@ public sealed class FoundryAgent : DelegatingAIAgent
#region Private helpers
private static ClientHeadersAgent CreateInnerAgent(
private static AIAgent CreateInnerAgent(
AIProjectClient aiProjectClient,
string model, string instructions,
string? name, string? description,
@@ -196,7 +196,7 @@ public sealed class FoundryAgent : DelegatingAIAgent
return CreateResponsesChatClientAgent(aiProjectClient, options, clientFactory, loggerFactory, services);
}
private static ClientHeadersAgent CreateResponsesChatClientAgent(
private static AIAgent CreateResponsesChatClientAgent(
AIProjectClient aiProjectClient,
ChatClientAgentOptions agentOptions,
Func<IChatClient, IChatClient>? clientFactory,
@@ -215,9 +215,25 @@ public sealed class FoundryAgent : DelegatingAIAgent
chatClient = clientFactory(chatClient);
}
// Register the ClientHeadersPolicy on the chat client's OpenAIRequestPolicies, if available.
// Silent no-op when the chat client is not OpenAI-backed.
if (chatClient.GetService<OpenAIRequestPolicies>() is { } policies)
return WireClientHeaders(new ChatClientAgent(chatClient, agentOptions, loggerFactory, services));
}
/// <summary>
/// Registers <see cref="ClientHeadersPolicy"/> on the agent's underlying chat client (if it
/// exposes <see cref="OpenAIRequestPolicies"/>) and wraps the agent in a
/// <see cref="ClientHeadersAgent"/> so per-call <c>x-client-*</c> headers stamped via
/// <see cref="ClientHeadersExtensions.WithClientHeader(ChatOptions, string, string)"/> reach
/// the wire. Idempotent: if the chain already contains a <see cref="ClientHeadersAgent"/>,
/// the original instance is returned unchanged.
/// </summary>
private static AIAgent WireClientHeaders(ChatClientAgent innerAgent)
{
if (innerAgent.GetService<ClientHeadersAgent>() is not null)
{
return innerAgent;
}
if (innerAgent.ChatClient.GetService<OpenAIRequestPolicies>() is { } policies)
{
OpenAIRequestPoliciesReflection.AddPolicyIfMissing(
policies,
@@ -225,11 +241,10 @@ public sealed class FoundryAgent : DelegatingAIAgent
System.ClientModel.Primitives.PipelinePosition.PerCall);
}
var inner = new ChatClientAgent(chatClient, agentOptions, loggerFactory, services);
return new ClientHeadersAgent(inner);
return new ClientHeadersAgent(innerAgent);
}
private static ClientHeadersAgent CreateInnerAgentFromEndpoint(
private static AIAgent CreateInnerAgentFromEndpoint(
AIProjectClient aiProjectClient,
Uri agentEndpoint,
IList<AITool>? tools,
@@ -254,16 +269,7 @@ public sealed class FoundryAgent : DelegatingAIAgent
chatClient = clientFactory(chatClient);
}
if (chatClient.GetService<OpenAIRequestPolicies>() is { } policies)
{
OpenAIRequestPoliciesReflection.AddPolicyIfMissing(
policies,
ClientHeadersPolicy.Instance,
System.ClientModel.Primitives.PipelinePosition.PerCall);
}
var inner = new ChatClientAgent(chatClient, agentOptions, services: services);
return new ClientHeadersAgent(inner);
return WireClientHeaders(new ChatClientAgent(chatClient, agentOptions, services: services));
}
private static AIProjectClient CreateProjectClient(Uri endpoint, AuthenticationTokenProvider credential, AIProjectClientOptions? clientOptions = null)