diff --git a/dotnet/src/Microsoft.Agents.AI/OpenTelemetryAgent.cs b/dotnet/src/Microsoft.Agents.AI/OpenTelemetryAgent.cs index 7ec8a53161a..d9df4793d61 100644 --- a/dotnet/src/Microsoft.Agents.AI/OpenTelemetryAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI/OpenTelemetryAgent.cs @@ -181,14 +181,25 @@ public async IAsyncEnumerable GetStreamingResponseAsync( { ForwardedOptions? fo = options as ForwardedOptions; - // Update the current activity to reflect the agent invocation. - parentAgent.UpdateCurrentActivity(fo?.CurrentActivity); + // Capture the current activity so we can restore it after each yield (workaround for + // dotnet/runtime#47802) and defer UpdateCurrentActivity until streaming completes to + // avoid interfering with FunctionInvokingChatClient's activity management. + Activity? activity = Activity.Current; - // Invoke the inner agent. - await foreach (var update in parentAgent.InnerAgent.RunStreamingAsync(messages, fo?.Session, fo?.Options, cancellationToken).ConfigureAwait(false)) + try + { + await foreach (var update in parentAgent.InnerAgent.RunStreamingAsync(messages, fo?.Session, fo?.Options, cancellationToken).ConfigureAwait(false)) + { + yield return update.AsChatResponseUpdate(); + + // Restore Activity.Current after yielding. + Activity.Current = activity; + } + } + finally { - // Wrap the response updates in ChatResponseUpdates so we can pass them back through OpenTelemetryChatClient. - yield return update.AsChatResponseUpdate(); + Activity.Current = activity; + parentAgent.UpdateCurrentActivity(fo?.CurrentActivity); } } diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/OpenTelemetryAgentTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/OpenTelemetryAgentTests.cs index 4cda58875ee..cd823f6b5ba 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/OpenTelemetryAgentTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/OpenTelemetryAgentTests.cs @@ -607,5 +607,245 @@ async static IAsyncEnumerable CallbackAsync( } } + /// + /// Verifies that Activity.Current is preserved throughout streaming responses that + /// include tool calls, ensuring all spans remain within the same trace. + /// + [Fact] + public async Task StreamingWithToolCalls_PreservesActivityCurrentAsync() + { + var sourceName = Guid.NewGuid().ToString(); + var activities = new List(); + using var tracerProvider = OpenTelemetry.Sdk.CreateTracerProviderBuilder() + .AddSource(sourceName) + .AddSource("Microsoft.Extensions.AI") + .AddInMemoryExporter(activities) + .Build(); + + Activity? activityDuringFirstCall = null; + Activity? activityDuringSecondCall = null; + bool toolWasCalled = false; + int callCount = 0; + + var getWeatherTool = AIFunctionFactory.Create( + (string location) => + { + toolWasCalled = true; + return $"Sunny and 72°F in {location}"; + }, + "GetCurrentWeather", + "Gets the current weather for a location."); + + var mockInnerClient = new CallbackChatClient + { + GetStreamingResponseAsyncCallback = (messages, options, ct) => + { + int currentCall = Interlocked.Increment(ref callCount); + if (currentCall == 1) + { + activityDuringFirstCall = Activity.Current; + return FirstCallStreamingAsync(ct); + } + + activityDuringSecondCall = Activity.Current; + return SecondCallStreamingAsync(ct); + }, + }; + + IChatClient chatClient = new ChatClientBuilder(mockInnerClient) + .UseOpenTelemetry(sourceName: sourceName) + .Build(); + + var agentOptions = new ChatClientAgentOptions + { + ChatOptions = new ChatOptions + { + Tools = [getWeatherTool], + }, + }; + var innerAgent = new ChatClientAgent(chatClient, agentOptions); + + using var otelAgent = new OpenTelemetryAgent(innerAgent, sourceName); + + using var parentSource = new ActivitySource(sourceName); + using var parentActivity = parentSource.StartActivity("HTTP POST /api/messages", ActivityKind.Server); + Assert.NotNull(parentActivity); + + var parentTraceId = parentActivity.TraceId; + + List messages = [new(ChatRole.User, "What's the weather in Seattle?")]; + var updates = new List(); + + await foreach (var update in otelAgent.RunStreamingAsync(messages)) + { + updates.Add(update); + } + + Assert.Equal(2, callCount); + Assert.True(toolWasCalled); + + // Activity.Current should still be the parent activity after streaming completes. + Assert.NotNull(Activity.Current); + Assert.Same(parentActivity, Activity.Current); + + // All activities should share the same TraceId. + foreach (var activity in activities) + { + Assert.Equal(parentTraceId, activity.TraceId); + } + + // Both LLM calls should have an active Activity within the same trace. + Assert.NotNull(activityDuringFirstCall); + Assert.NotNull(activityDuringSecondCall); + Assert.Equal(parentTraceId, activityDuringFirstCall!.TraceId); + Assert.Equal(parentTraceId, activityDuringSecondCall!.TraceId); + + static async IAsyncEnumerable FirstCallStreamingAsync( + [EnumeratorCancellation] CancellationToken ct) + { + await Task.Yield(); + yield return new ChatResponseUpdate + { + Role = ChatRole.Assistant, + Contents = [new FunctionCallContent("call_001", "GetCurrentWeather", + new Dictionary { ["location"] = "Seattle" })], + ResponseId = "resp_1", + }; + } + + static async IAsyncEnumerable SecondCallStreamingAsync( + [EnumeratorCancellation] CancellationToken ct) + { + await Task.Yield(); + yield return new ChatResponseUpdate + { + Role = ChatRole.Assistant, + Contents = [new TextContent("It's sunny and 72°F in Seattle!")], + ResponseId = "resp_2", + }; + } + } + + /// + /// Verifies Activity.Current preservation at the IChatClient pipeline level + /// (FunctionInvokingChatClient + OpenTelemetryChatClient) independent of the agent layer. + /// + [Fact] + public async Task StreamingWithToolCalls_ChatClientPipeline_PreservesActivityCurrentAsync() + { + var sourceName = Guid.NewGuid().ToString(); + var activities = new List(); + using var tracerProvider = OpenTelemetry.Sdk.CreateTracerProviderBuilder() + .AddSource(sourceName) + .AddSource("Microsoft.Extensions.AI") + .AddInMemoryExporter(activities) + .Build(); + + Activity? activityDuringSecondCall = null; + bool toolWasCalled = false; + int callCount = 0; + + var getWeatherTool = AIFunctionFactory.Create( + (string location) => + { + toolWasCalled = true; + return $"Sunny and 72°F in {location}"; + }, + "GetCurrentWeather", + "Gets the current weather for a location."); + + var mockInnerClient = new CallbackChatClient + { + GetStreamingResponseAsyncCallback = (messages, options, ct) => + { + int currentCall = Interlocked.Increment(ref callCount); + if (currentCall == 1) + { + return FirstCallStreamingAsync(ct); + } + + activityDuringSecondCall = Activity.Current; + return SecondCallStreamingAsync(ct); + }, + }; + + IChatClient pipeline = new ChatClientBuilder(mockInnerClient) + .UseFunctionInvocation() + .UseOpenTelemetry(sourceName: sourceName) + .Build(); + + using var parentSource = new ActivitySource(sourceName); + using var parentActivity = parentSource.StartActivity("parent-operation", ActivityKind.Server); + Assert.NotNull(parentActivity); + var parentTraceId = parentActivity.TraceId; + + var chatOptions = new ChatOptions { Tools = [getWeatherTool] }; + var updates = new List(); + await foreach (var update in pipeline.GetStreamingResponseAsync( + [new ChatMessage(ChatRole.User, "What's the weather?")], chatOptions)) + { + updates.Add(update); + } + + Assert.Equal(2, callCount); + Assert.True(toolWasCalled); + + Assert.NotNull(Activity.Current); + Assert.Same(parentActivity, Activity.Current); + + Assert.NotNull(activityDuringSecondCall); + Assert.Equal(parentTraceId, activityDuringSecondCall!.TraceId); + + foreach (var activity in activities) + { + Assert.Equal(parentTraceId, activity.TraceId); + } + + static async IAsyncEnumerable FirstCallStreamingAsync( + [EnumeratorCancellation] CancellationToken ct) + { + await Task.Yield(); + yield return new ChatResponseUpdate + { + Role = ChatRole.Assistant, + Contents = [new FunctionCallContent("call_001", "GetCurrentWeather", + new Dictionary { ["location"] = "Seattle" })], + ResponseId = "resp_1", + }; + } + + static async IAsyncEnumerable SecondCallStreamingAsync( + [EnumeratorCancellation] CancellationToken ct) + { + await Task.Yield(); + yield return new ChatResponseUpdate + { + Role = ChatRole.Assistant, + Contents = [new TextContent("Sunny in Seattle!")], + ResponseId = "resp_2", + }; + } + } + + /// Simple callback-based IChatClient for testing. + private sealed class CallbackChatClient : IChatClient + { + public Func, ChatOptions?, CancellationToken, IAsyncEnumerable>? GetStreamingResponseAsyncCallback { get; set; } + public Func, ChatOptions?, CancellationToken, Task>? GetResponseAsyncCallback { get; set; } + + public IAsyncEnumerable GetStreamingResponseAsync( + IEnumerable messages, ChatOptions? options = null, CancellationToken cancellationToken = default) + => this.GetStreamingResponseAsyncCallback?.Invoke(messages, options, cancellationToken) + ?? throw new NotSupportedException(); + + public Task GetResponseAsync( + IEnumerable messages, ChatOptions? options = null, CancellationToken cancellationToken = default) + => this.GetResponseAsyncCallback?.Invoke(messages, options, cancellationToken) + ?? throw new NotSupportedException(); + + public object? GetService(Type serviceType, object? serviceKey = null) => null; + public void Dispose() { } + } + private static string ReplaceWhitespace(string? input) => Regex.Replace(input ?? "", @"\s+", "").Trim(); }