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")]