From 54982ba54a0b2f3bed476a621d41f93c1534a71b Mon Sep 17 00:00:00 2001 From: Roger Barreto <19890735+rogerbarreto@users.noreply.github.com> Date: Fri, 20 Feb 2026 12:47:00 +0000 Subject: [PATCH] Fix Activity.Current being set to null when FunctionInvokingChatClient runs inside invoke_agent span When CurrentActivityIsInvokeAgent is true, the orchestrate_tools activity is correctly suppressed (activity = null). However, the Activity.Current = activity workaround for dotnet/runtime#47802 after each yield return was then setting Activity.Current to null, destroying the ambient trace context. This caused all subsequent spans (tool execution, second LLM call) to become disconnected root traces instead of being children of the invoke_agent span. The fix saves the parent activity (Activity.Current) before the orchestrate_tools decision and uses activity ?? parentActivity as the restoration target. When orchestrate_tools is suppressed inside invoke_agent, Activity.Current is restored to the invoke_agent span instead of null. This preserves both: - The orchestrate_tools suppression fix from PR #7224 (issue #7223) - Correct trace context propagation across streaming + tool call cycles Fixes microsoft/agent-framework#4074 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../FunctionInvokingChatClient.cs | 15 +++-- .../FunctionInvokingChatClientTests.cs | 55 +++++++++++++++++++ 2 files changed, 64 insertions(+), 6 deletions(-) diff --git a/src/Libraries/Microsoft.Extensions.AI/ChatCompletion/FunctionInvokingChatClient.cs b/src/Libraries/Microsoft.Extensions.AI/ChatCompletion/FunctionInvokingChatClient.cs index 9fd95df0424..5a87a25ca5d 100644 --- a/src/Libraries/Microsoft.Extensions.AI/ChatCompletion/FunctionInvokingChatClient.cs +++ b/src/Libraries/Microsoft.Extensions.AI/ChatCompletion/FunctionInvokingChatClient.cs @@ -272,6 +272,7 @@ public override async Task GetResponseAsync( // A single request into this GetResponseAsync may result in multiple requests to the inner client. // Create an activity to group them together for better observability. If there's already a genai "invoke_agent" // span that's current, however, we just consider that the group and don't add a new one. + Activity? parentActivity = Activity.Current; using Activity? activity = CurrentActivityIsInvokeAgent ? null : _activitySource?.StartActivity(OpenTelemetryConsts.GenAI.OrchestrateToolsName); // Copy the original messages in order to avoid enumerating the original messages multiple times. @@ -420,7 +421,9 @@ public override async IAsyncEnumerable GetStreamingResponseA // A single request into this GetStreamingResponseAsync may result in multiple requests to the inner client. // Create an activity to group them together for better observability. If there's already a genai "invoke_agent" // span that's current, however, we just consider that the group and don't add a new one. + Activity? parentActivity = Activity.Current; using Activity? activity = CurrentActivityIsInvokeAgent ? null : _activitySource?.StartActivity(OpenTelemetryConsts.GenAI.OrchestrateToolsName); + Activity? activityToRestore = activity ?? parentActivity; UsageDetails? totalUsage = activity is { IsAllDataRequested: true } ? new() : null; // tracked usage across all turns, to be used for activity purposes // Copy the original messages in order to avoid enumerating the original messages multiple times. @@ -460,7 +463,7 @@ public override async IAsyncEnumerable GetStreamingResponseA foreach (var message in preDownstreamCallHistory) { yield return ConvertToolResultMessageToUpdate(message, options?.ConversationId, message.MessageId); - Activity.Current = activity; // workaround for https://github.com/dotnet/runtime/issues/47802 + Activity.Current = activityToRestore; // workaround for https://github.com/dotnet/runtime/issues/47802 } } @@ -474,7 +477,7 @@ public override async IAsyncEnumerable GetStreamingResponseA { message.MessageId = toolMessageId; yield return ConvertToolResultMessageToUpdate(message, options?.ConversationId, message.MessageId); - Activity.Current = activity; // workaround for https://github.com/dotnet/runtime/issues/47802 + Activity.Current = activityToRestore; // workaround for https://github.com/dotnet/runtime/issues/47802 } if (shouldTerminate) @@ -557,7 +560,7 @@ public override async IAsyncEnumerable GetStreamingResponseA // we can yield the update as-is. lastYieldedUpdateIndex++; yield return update; - Activity.Current = activity; // workaround for https://github.com/dotnet/runtime/issues/47802 + Activity.Current = activityToRestore; // workaround for https://github.com/dotnet/runtime/issues/47802 continue; } @@ -584,7 +587,7 @@ public override async IAsyncEnumerable GetStreamingResponseA } yield return updateToYield; - Activity.Current = activity; // workaround for https://github.com/dotnet/runtime/issues/47802 + Activity.Current = activityToRestore; // workaround for https://github.com/dotnet/runtime/issues/47802 } continue; @@ -601,7 +604,7 @@ public override async IAsyncEnumerable GetStreamingResponseA { var updateToYield = updates[lastYieldedUpdateIndex]; yield return updateToYield; - Activity.Current = activity; // workaround for https://github.com/dotnet/runtime/issues/47802 + Activity.Current = activityToRestore; // workaround for https://github.com/dotnet/runtime/issues/47802 } // If there's nothing more to do, break out of the loop and allow the handling at the @@ -632,7 +635,7 @@ public override async IAsyncEnumerable GetStreamingResponseA foreach (var message in modeAndMessages.MessagesAdded) { yield return ConvertToolResultMessageToUpdate(message, response.ConversationId, toolMessageId); - Activity.Current = activity; // workaround for https://github.com/dotnet/runtime/issues/47802 + Activity.Current = activityToRestore; // workaround for https://github.com/dotnet/runtime/issues/47802 } if (modeAndMessages.ShouldTerminate) diff --git a/test/Libraries/Microsoft.Extensions.AI.Tests/ChatCompletion/FunctionInvokingChatClientTests.cs b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatCompletion/FunctionInvokingChatClientTests.cs index 2ddf757d185..b5d9077d2c0 100644 --- a/test/Libraries/Microsoft.Extensions.AI.Tests/ChatCompletion/FunctionInvokingChatClientTests.cs +++ b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatCompletion/FunctionInvokingChatClientTests.cs @@ -1668,6 +1668,7 @@ public async Task ClonesChatOptionsAndResetContinuationTokenForBackgroundRespons [InlineData("invoke_agent")] [InlineData("invoke_agent my_agent")] [InlineData("invoke_agent ")] + [InlineData("invoke_agent MyAgent(agent-123)")] public async Task DoesNotCreateOrchestrateToolsSpanWhenInvokeAgentIsParent(string displayName) { string agentSourceName = Guid.NewGuid().ToString(); @@ -1713,6 +1714,60 @@ public async Task DoesNotCreateOrchestrateToolsSpanWhenInvokeAgentIsParent(strin Assert.All(childActivities, activity => Assert.Same(invokeAgent, activity.Parent)); } + [Theory] + [InlineData("invoke_agent")] + [InlineData("invoke_agent my_agent")] + [InlineData("invoke_agent MyAgent(agent-123)")] + public async Task StreamingPreservesActivityCurrentWhenInvokeAgentIsParent(string displayName) + { + string agentSourceName = Guid.NewGuid().ToString(); + string clientSourceName = Guid.NewGuid().ToString(); + + List plan = + [ + new ChatMessage(ChatRole.User, "hello"), + new ChatMessage(ChatRole.Assistant, [new FunctionCallContent("callId1", "Func1")]), + new ChatMessage(ChatRole.Tool, [new FunctionResultContent("callId1", result: "Result 1")]), + new ChatMessage(ChatRole.Assistant, "world"), + ]; + + ChatOptions options = new() + { + Tools = [AIFunctionFactory.Create(() => "Result 1", "Func1")] + }; + + Func configure = b => b.Use(c => + new FunctionInvokingChatClient(new OpenTelemetryChatClient(c, sourceName: clientSourceName))); + + var activities = new List(); + + using TracerProvider tracerProvider = OpenTelemetry.Sdk.CreateTracerProviderBuilder() + .AddSource(agentSourceName) + .AddSource(clientSourceName) + .AddInMemoryExporter(activities) + .Build(); + + using (var agentSource = new ActivitySource(agentSourceName)) + using (var invokeAgentActivity = agentSource.StartActivity(displayName)) + { + Assert.NotNull(invokeAgentActivity); + await InvokeAndAssertStreamingAsync(options, plan, configurePipeline: configure); + + // Activity.Current must still be the invoke_agent activity after streaming completes. + // This is the regression test for https://github.com/microsoft/agent-framework/issues/4074: + // Before the fix, Activity.Current was set to null after the first streaming + tool call cycle. + Assert.Same(invokeAgentActivity, Activity.Current); + } + + Assert.DoesNotContain(activities, a => a.DisplayName == "orchestrate_tools"); + Assert.Contains(activities, a => a.DisplayName == "chat"); + Assert.Contains(activities, a => a.DisplayName == "execute_tool Func1"); + + var invokeAgent = Assert.Single(activities, a => a.DisplayName == displayName); + var childActivities = activities.Where(a => a != invokeAgent).ToList(); + Assert.All(childActivities, activity => Assert.Same(invokeAgent, activity.Parent)); + } + [Theory] [InlineData("invoke_agen")] [InlineData("invoke_agent_extra")]