From bf2685b95c8d7bc9e593454cade3c6a9fc320763 Mon Sep 17 00:00:00 2001 From: Flavio da Costa Date: Thu, 19 Feb 2026 09:12:21 -0600 Subject: [PATCH 1/3] Fix Activity.Current nulled during streaming Fixes #7320 When an invoke_agent span with a name suffix is the current activity (e.g. "invoke_agent MyAgent(id)"), CurrentActivityIsInvokeAgent returns true and the orchestrate_tools activity is not created. The local activity variable is null, and all Activity.Current = activity workaround sites for dotnet/runtime#47802 set Activity.Current to null after each yield return. This disconnects subsequent spans from the trace. Capture the current activity before it can be lost and restore that value instead of null after each yield. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../FunctionInvokingChatClient.cs | 13 ++-- .../FunctionInvokingChatClientTests.cs | 61 +++++++++++++++++++ 2 files changed, 68 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..835c3aeab1d 100644 --- a/src/Libraries/Microsoft.Extensions.AI/ChatCompletion/FunctionInvokingChatClient.cs +++ b/src/Libraries/Microsoft.Extensions.AI/ChatCompletion/FunctionInvokingChatClient.cs @@ -421,6 +421,7 @@ public override async IAsyncEnumerable GetStreamingResponseA // 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. using Activity? activity = CurrentActivityIsInvokeAgent ? null : _activitySource?.StartActivity(OpenTelemetryConsts.GenAI.OrchestrateToolsName); + Activity? activityToRestore = activity ?? Activity.Current; // when activity is null (under invoke_agent), restore the parent instead 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 +461,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 +475,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 +558,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 +585,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 +602,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 +633,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..7d3600d7141 100644 --- a/test/Libraries/Microsoft.Extensions.AI.Tests/ChatCompletion/FunctionInvokingChatClientTests.cs +++ b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatCompletion/FunctionInvokingChatClientTests.cs @@ -1713,6 +1713,67 @@ public async Task DoesNotCreateOrchestrateToolsSpanWhenInvokeAgentIsParent(strin Assert.All(childActivities, activity => Assert.Same(invokeAgent, activity.Parent)); } + [Fact] + public async Task StreamingPreservesTraceContextWhenInvokeAgentWithNameIsParent() + { + // Reproduces: Activity.Current lost after streaming + tool call when + // parent is "invoke_agent AgentName(id)" (agent-framework naming convention). + // The broadened CurrentActivityIsInvokeAgent match (PR #7224) causes + // activity to be null, and Activity.Current = activity wipes the context. + string agentSourceName = Guid.NewGuid().ToString(); + string clientSourceName = Guid.NewGuid().ToString(); + var activities = new List(); + + using TracerProvider tracerProvider = OpenTelemetry.Sdk.CreateTracerProviderBuilder() + .AddSource(agentSourceName) + .AddSource(clientSourceName) + .AddInMemoryExporter(activities) + .Build(); + + int callCount = 0; + using var innerClient = new TestChatClient + { + GetStreamingResponseAsyncCallback = (messages, options, ct) => + { + callCount++; + ChatMessage message = callCount == 1 + ? new(ChatRole.Assistant, [new FunctionCallContent("call1", "Func1")]) + : new(ChatRole.Assistant, "Done"); + return YieldAsync(new ChatResponse(message).ToChatResponseUpdates()); + } + }; + + var client = innerClient.AsBuilder() + .Use(c => new FunctionInvokingChatClient( + new OpenTelemetryChatClient(c, sourceName: clientSourceName))) + .Build(); + + var options = new ChatOptions + { + Tools = [AIFunctionFactory.Create(() => "Result 1", "Func1")] + }; + + using var agentSource = new ActivitySource(agentSourceName); + using var invokeAgentActivity = agentSource.StartActivity("invoke_agent MyAgent(agent-123)"); + Assert.NotNull(invokeAgentActivity); + + await foreach (var update in client.GetStreamingResponseAsync( + [new ChatMessage(ChatRole.User, "hello")], options)) + { + // consume all updates + } + + Assert.Equal(2, callCount); + + var chatActivities = activities.Where(a => a.DisplayName.StartsWith("chat", StringComparison.Ordinal)).ToList(); + Assert.Equal(2, chatActivities.Count); + + // All child activities must share the same trace as invoke_agent + var nonAgentActivities = activities.Where(a => a != invokeAgentActivity).ToList(); + Assert.All(nonAgentActivities, a => + Assert.Equal(invokeAgentActivity.TraceId, a.TraceId)); + } + [Theory] [InlineData("invoke_agen")] [InlineData("invoke_agent_extra")] From 6b0b491459d0b08dca7edd2c9d6f92c60cef96bb Mon Sep 17 00:00:00 2001 From: Flavio da Costa Date: Thu, 19 Feb 2026 09:35:19 -0600 Subject: [PATCH 2/3] Address review: guard with null check instead Use a null check around Activity.Current assignment instead of capturing a separate activityToRestore variable, per reviewer feedback. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../FunctionInvokingChatClient.cs | 32 ++++++++++++++----- 1 file changed, 24 insertions(+), 8 deletions(-) diff --git a/src/Libraries/Microsoft.Extensions.AI/ChatCompletion/FunctionInvokingChatClient.cs b/src/Libraries/Microsoft.Extensions.AI/ChatCompletion/FunctionInvokingChatClient.cs index 835c3aeab1d..f1858a295cf 100644 --- a/src/Libraries/Microsoft.Extensions.AI/ChatCompletion/FunctionInvokingChatClient.cs +++ b/src/Libraries/Microsoft.Extensions.AI/ChatCompletion/FunctionInvokingChatClient.cs @@ -421,7 +421,6 @@ public override async IAsyncEnumerable GetStreamingResponseA // 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. using Activity? activity = CurrentActivityIsInvokeAgent ? null : _activitySource?.StartActivity(OpenTelemetryConsts.GenAI.OrchestrateToolsName); - Activity? activityToRestore = activity ?? Activity.Current; // when activity is null (under invoke_agent), restore the parent instead 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. @@ -461,7 +460,10 @@ public override async IAsyncEnumerable GetStreamingResponseA foreach (var message in preDownstreamCallHistory) { yield return ConvertToolResultMessageToUpdate(message, options?.ConversationId, message.MessageId); - Activity.Current = activityToRestore; // workaround for https://github.com/dotnet/runtime/issues/47802 + if (activity is not null) + { + Activity.Current = activity; // workaround for https://github.com/dotnet/runtime/issues/47802 + } } } @@ -475,7 +477,10 @@ public override async IAsyncEnumerable GetStreamingResponseA { message.MessageId = toolMessageId; yield return ConvertToolResultMessageToUpdate(message, options?.ConversationId, message.MessageId); - Activity.Current = activityToRestore; // workaround for https://github.com/dotnet/runtime/issues/47802 + if (activity is not null) + { + Activity.Current = activity; // workaround for https://github.com/dotnet/runtime/issues/47802 + } } if (shouldTerminate) @@ -558,8 +563,10 @@ public override async IAsyncEnumerable GetStreamingResponseA // we can yield the update as-is. lastYieldedUpdateIndex++; yield return update; - Activity.Current = activityToRestore; // workaround for https://github.com/dotnet/runtime/issues/47802 - + if (activity is not null) + { + Activity.Current = activity; // workaround for https://github.com/dotnet/runtime/issues/47802 + } continue; } @@ -585,7 +592,10 @@ public override async IAsyncEnumerable GetStreamingResponseA } yield return updateToYield; - Activity.Current = activityToRestore; // workaround for https://github.com/dotnet/runtime/issues/47802 + if (activity is not null) + { + Activity.Current = activity; // workaround for https://github.com/dotnet/runtime/issues/47802 + } } continue; @@ -602,7 +612,10 @@ public override async IAsyncEnumerable GetStreamingResponseA { var updateToYield = updates[lastYieldedUpdateIndex]; yield return updateToYield; - Activity.Current = activityToRestore; // workaround for https://github.com/dotnet/runtime/issues/47802 + if (activity is not null) + { + Activity.Current = activity; // 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 @@ -633,7 +646,10 @@ public override async IAsyncEnumerable GetStreamingResponseA foreach (var message in modeAndMessages.MessagesAdded) { yield return ConvertToolResultMessageToUpdate(message, response.ConversationId, toolMessageId); - Activity.Current = activityToRestore; // workaround for https://github.com/dotnet/runtime/issues/47802 + if (activity is not null) + { + Activity.Current = activity; // workaround for https://github.com/dotnet/runtime/issues/47802 + } } if (modeAndMessages.ShouldTerminate) From f50d00b16f5c798d6e6c20b3b3c7b0266399c00b Mon Sep 17 00:00:00 2001 From: Stephen Toub Date: Thu, 19 Feb 2026 17:08:32 -0500 Subject: [PATCH 3/3] Update test/Libraries/Microsoft.Extensions.AI.Tests/ChatCompletion/FunctionInvokingChatClientTests.cs --- .../ChatCompletion/FunctionInvokingChatClientTests.cs | 4 ---- 1 file changed, 4 deletions(-) diff --git a/test/Libraries/Microsoft.Extensions.AI.Tests/ChatCompletion/FunctionInvokingChatClientTests.cs b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatCompletion/FunctionInvokingChatClientTests.cs index 7d3600d7141..dd2a139b4f6 100644 --- a/test/Libraries/Microsoft.Extensions.AI.Tests/ChatCompletion/FunctionInvokingChatClientTests.cs +++ b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatCompletion/FunctionInvokingChatClientTests.cs @@ -1716,10 +1716,6 @@ public async Task DoesNotCreateOrchestrateToolsSpanWhenInvokeAgentIsParent(strin [Fact] public async Task StreamingPreservesTraceContextWhenInvokeAgentWithNameIsParent() { - // Reproduces: Activity.Current lost after streaming + tool call when - // parent is "invoke_agent AgentName(id)" (agent-framework naming convention). - // The broadened CurrentActivityIsInvokeAgent match (PR #7224) causes - // activity to be null, and Activity.Current = activity wipes the context. string agentSourceName = Guid.NewGuid().ToString(); string clientSourceName = Guid.NewGuid().ToString(); var activities = new List();