Skip to content

FunctionInvokingChatClient streaming sets Activity.Current = null when parent is an invoke_agent <name> span #7320

@flaviocdc

Description

@flaviocdc

Description

FunctionInvokingChatClient.GetStreamingResponseAsync destroys Activity.Current after each yield return when the current activity matches the broadened invoke_agent span detection introduced in #7224. This causes all spans after the first streaming response + tool call cycle to become disconnected root traces.

Regression

This is a regression from 10.2.0 → 10.3.0, introduced by #7224.

In 10.2.0, CurrentActivityIsInvokeAgent used exact match (== "invoke_agent"), so callers using "invoke_agent <name>" as the display name did not match — the orchestrate_tools activity was created, and the Activity.Current = activity workaround preserved trace context.

In 10.3.0, the match was broadened to prefix + space, so "invoke_agent MyAgent" now matches. When it matches, activity is set to null, and all 6 Activity.Current = activity workaround sites in the streaming method set Activity.Current = null, breaking trace propagation.

Downstream impact: microsoft/agent-framework#4074

Root Cause

In GetStreamingResponseAsync (line ~423):

using Activity? activity = CurrentActivityIsInvokeAgent ? null : _activitySource?.StartActivity("orchestrate_tools");

When CurrentActivityIsInvokeAgent is true, activity is null. Then every yield return is followed by:

Activity.Current = activity; // workaround for dotnet/runtime#47802

This sets Activity.Current = null instead of restoring the parent invoke_agent activity.

Minimal Reproduction

This test can be added to FunctionInvokingChatClientTests.cs and fails on main:

[Fact]
public async Task StreamingPreservesTraceContextWhenInvokeAgentWithNameIsParent()
{
    string agentSourceName = Guid.NewGuid().ToString();
    string clientSourceName = Guid.NewGuid().ToString();
    var activities = new List<Activity>();

    using var 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))
    {
    }

    Assert.Equal(2, callCount);

    var chatActivities = activities.Where(a => a.DisplayName.StartsWith("chat", StringComparison.Ordinal)).ToList();
    Assert.Equal(2, chatActivities.Count);

    // FAILS: execute_tool and second chat span have different trace IDs (root traces)
    var nonAgentActivities = activities.Where(a => a != invokeAgentActivity).ToList();
    Assert.All(nonAgentActivities, a =>
        Assert.Equal(invokeAgentActivity.TraceId, a.TraceId));
}

Expected Behavior

All child spans share the same TraceId as the parent invoke_agent span.

Actual Behavior

After the first streaming response + tool call, Activity.Current is null. The execute_tool span and second chat span start as disconnected root traces with different TraceIds.

Environment

  • Microsoft.Extensions.AI 10.3.0
  • .NET 9.0 / .NET 8.0 / net462 (all affected)

Metadata

Metadata

Assignees

No one assigned

    Labels

    area-aiMicrosoft.Extensions.AI libraries

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions