Preserve base AgentRunOptions properties and avoid double-wrap on user factory

Agent-Logs-Url: https://github.com/microsoft/agent-framework/sessions/3afbf18c-de22-4236-a2f2-02ca1e98ae21

Co-authored-by: rogerbarreto <19890735+rogerbarreto@users.noreply.github.com>
This commit is contained in:
copilot-swe-agent[bot]
2026-05-12 08:06:28 +00:00
committed by Roger Barreto
Unverified
parent ca732733d3
commit 921162425d
2 changed files with 96 additions and 5 deletions
@@ -229,24 +229,46 @@ public sealed class OpenTelemetryAgent : DelegatingAIAgent, IDisposable
}
string? sourceName = this._sourceName;
IChatClient WrapWithOpenTelemetry(IChatClient cc) =>
cc.AsBuilder().UseOpenTelemetry(sourceName: sourceName).Build();
static IChatClient WrapIfNeeded(IChatClient cc, string? sourceName) =>
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.
// If the user factory already returns an OpenTelemetry-instrumented client, don't double-wrap.
var clone = (ChatClientAgentRunOptions)ccOptions.Clone();
var userFactory = clone.ChatClientFactory;
clone.ChatClientFactory = cc => WrapWithOpenTelemetry(userFactory is null ? cc : userFactory(cc));
clone.ChatClientFactory = cc => WrapIfNeeded(userFactory is null ? cc : userFactory(cc), sourceName);
return clone;
}
return new ChatClientAgentRunOptions
// For a plain AgentRunOptions (or null), create a ChatClientAgentRunOptions and preserve
// any base AgentRunOptions properties from the caller so they reach the inner agent.
var newOptions = new ChatClientAgentRunOptions
{
ChatClientFactory = WrapWithOpenTelemetry,
ChatClientFactory = cc => WrapIfNeeded(cc, sourceName),
};
if (options is not null)
{
CopyBaseAgentRunOptions(options, newOptions);
}
return newOptions;
}
#pragma warning disable MEAI001 // ContinuationToken is experimental; copy it through to preserve caller-provided value.
private static void CopyBaseAgentRunOptions(AgentRunOptions source, AgentRunOptions target)
{
target.ContinuationToken = source.ContinuationToken;
target.AllowBackgroundResponses = source.AllowBackgroundResponses;
target.AdditionalProperties = source.AdditionalProperties?.Clone();
target.ResponseFormat = source.ResponseFormat;
}
#pragma warning restore MEAI001
/// <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
@@ -792,6 +792,75 @@ public class OpenTelemetryAgentTests
Assert.Contains(activities, a => string.Equals(a.GetTagItem("gen_ai.operation.name") as string, "chat", StringComparison.Ordinal));
}
[Fact]
public async Task AutoWireChatClient_PlainAgentRunOptions_PreservesBaseProperties_Async()
{
// Auto-wiring converts a plain AgentRunOptions into a ChatClientAgentRunOptions. The base
// properties (ContinuationToken, AllowBackgroundResponses, AdditionalProperties, ResponseFormat)
// must be preserved so they reach the inner agent.
AgentRunOptions? observedOptions = null;
var fakeChatClient = new AutoWireTestChatClient();
var innerChatClientAgent = new ChatClientAgent(fakeChatClient);
// Wrapping agent: surfaces the ChatClientAgent via GetService (so auto-wiring activates),
// but captures the AgentRunOptions passed to RunAsync by the OpenTelemetryAgent.
var wrapper = new TestAIAgent
{
GetServiceFunc = (type, key) =>
type == typeof(ChatClientAgent) ? innerChatClientAgent : null,
RunAsyncFunc = (messages, session, options, ct) =>
{
observedOptions = options;
return Task.FromResult(new AgentResponse(new ChatMessage(ChatRole.Assistant, "ok")));
},
};
using var agent = new OpenTelemetryAgent(wrapper);
var additionalProps = new AdditionalPropertiesDictionary { ["customKey"] = "customValue" };
var inputOptions = new AgentRunOptions
{
AllowBackgroundResponses = true,
AdditionalProperties = additionalProps,
ResponseFormat = ChatResponseFormat.Json,
};
_ = await agent.RunAsync("hi", options: inputOptions);
Assert.NotNull(observedOptions);
Assert.IsType<ChatClientAgentRunOptions>(observedOptions);
Assert.Equal(true, observedOptions!.AllowBackgroundResponses);
Assert.Same(ChatResponseFormat.Json, observedOptions.ResponseFormat);
Assert.NotNull(observedOptions.AdditionalProperties);
Assert.Equal("customValue", observedOptions.AdditionalProperties!["customKey"]);
}
[Fact]
public async Task AutoWireChatClient_UserFactoryReturnsInstrumentedClient_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();
var inner = new ChatClientAgent(fakeChatClient);
using var agent = new OpenTelemetryAgent(inner, sourceName);
// User factory wraps the chat client with OpenTelemetryChatClient itself.
var runOptions = new ChatClientAgentRunOptions
{
ChatClientFactory = cc => cc.AsBuilder().UseOpenTelemetry(sourceName: sourceName).Build(),
};
_ = await agent.RunAsync("hi", options: runOptions);
// Expect 2 activities (invoke_agent + a single chat span). If we double-wrapped, we would see 3.
Assert.Equal(2, activities.Count);
}
private sealed class AutoWireTestChatClient : IChatClient
{
public Task<ChatResponse> GetResponseAsync(IEnumerable<ChatMessage> messages, ChatOptions? options = null, CancellationToken cancellationToken = default) =>