Fix Activity.Current nulled during streaming tool invocation

src/Libraries/Microsoft.Extensions.AI/ChatCompletion/FunctionInvokingChatClient.cs

@@ -421,6 +421,7 @@ public override async IAsyncEnumerable<ChatResponseUpdate> 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<ChatResponseUpdate> 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<ChatResponseUpdate> 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<ChatResponseUpdate> 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<ChatResponseUpdate> 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<ChatResponseUpdate> 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<ChatResponseUpdate> 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)

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<Activity>();
+
+        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")]