Skip to content

Fix Activity.Current being set to null when FunctionInvokingChatClient #7324

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
rogerbarreto wants to merge 1 commit into from
Closed

Fix Activity.Current being set to null when FunctionInvokingChatClient #7324

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 9 additions & 6 deletions src/Libraries/Microsoft.Extensions.AI/ChatCompletion/FunctionInvokingChatClient.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -272,6 +272,7 @@ public override async Task<ChatResponse> 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;

Copilot AI Feb 20, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The parentActivity variable is captured but never used in the non-streaming GetResponseAsync method. Unlike the streaming method which uses it to restore Activity.Current after each yield, this method doesn't need activity restoration since it doesn't use yield return. Consider removing this line.

Suggested change
Activity? parentActivity = Activity.Current;

Copilot uses AI. Check for mistakes.
using Activity? activity = CurrentActivityIsInvokeAgent ? null : _activitySource?.StartActivity(OpenTelemetryConsts.GenAI.OrchestrateToolsName);

// Copy the original messages in order to avoid enumerating the original messages multiple times.
Expand Down Expand Up @@ -420,7 +421,9 @@ public override async IAsyncEnumerable<ChatResponseUpdate> 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.
Expand Down Expand Up @@ -460,7 +463,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
}
}

Expand All @@ -474,7 +477,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)
Expand Down Expand Up @@ -557,7 +560,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;
}
Expand All @@ -584,7 +587,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;
Expand All @@ -601,7 +604,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
Expand Down Expand Up @@ -632,7 +635,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)
Expand Down
55 changes: 55 additions & 0 deletions ...Libraries/Microsoft.Extensions.AI.Tests/ChatCompletion/FunctionInvokingChatClientTests.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down Expand Up @@ -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<ChatMessage> 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<ChatClientBuilder, ChatClientBuilder> configure = b => b.Use(c =>
new FunctionInvokingChatClient(new OpenTelemetryChatClient(c, sourceName: clientSourceName)));

var activities = new List<Activity>();

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")]
Expand Down
Loading