chore: import upstream snapshot with attribution

dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/WorkflowRunActivityStopTests.cs

@@ -0,0 +1,344 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.Collections.Concurrent;
+using System.Diagnostics;
+using System.Linq;
+using System.Threading.Tasks;
+using FluentAssertions;
+using Microsoft.Agents.AI.Workflows.Observability;
+
+namespace Microsoft.Agents.AI.Workflows.UnitTests;
+
+/// <summary>
+/// Regression test for https://github.com/microsoft/agent-framework/issues/4155
+/// Verifies that the workflow_invoke Activity is properly stopped/disposed so it gets exported
+/// to telemetry backends. The ActivityStopped callback must fire for the workflow_invoke span.
+/// </summary>
+[Collection("ObservabilityTests")]
+public sealed class WorkflowRunActivityStopTests : IDisposable
+{
+    private readonly ActivityListener _activityListener;
+    private readonly ConcurrentBag<Activity> _startedActivities = [];
+    private readonly ConcurrentBag<Activity> _stoppedActivities = [];
+    private bool _isDisposed;
+
+    public WorkflowRunActivityStopTests()
+    {
+        this._activityListener = new ActivityListener
+        {
+            ShouldListenTo = source => source.Name.Contains(typeof(Workflow).Namespace!),
+            Sample = (ref ActivityCreationOptions<ActivityContext> options) => ActivitySamplingResult.AllData,
+            ActivityStarted = activity => this._startedActivities.Add(activity),
+            ActivityStopped = activity => this._stoppedActivities.Add(activity),
+        };
+        ActivitySource.AddActivityListener(this._activityListener);
+    }
+
+    public void Dispose()
+    {
+        if (!this._isDisposed)
+        {
+            this._activityListener?.Dispose();
+            this._isDisposed = true;
+        }
+    }
+
+    /// <summary>
+    /// Creates a simple sequential workflow with OpenTelemetry enabled.
+    /// </summary>
+    private static Workflow CreateWorkflow()
+    {
+        Func<string, string> uppercaseFunc = s => s.ToUpperInvariant();
+        var uppercase = uppercaseFunc.BindAsExecutor("UppercaseExecutor");
+
+        Func<string, string> reverseFunc = s => new string(s.Reverse().ToArray());
+        var reverse = reverseFunc.BindAsExecutor("ReverseTextExecutor");
+
+        WorkflowBuilder builder = new(uppercase);
+        builder.AddEdge(uppercase, reverse).WithOutputFrom(reverse);
+
+        return builder.WithOpenTelemetry().Build();
+    }
+
+    /// <summary>
+    /// Verifies that the workflow_invoke Activity is stopped (and thus exportable) when
+    /// using the Lockstep execution environment.
+    /// Bug: The Activity created by LockstepRunEventStream.TakeEventStreamAsync is never
+    /// disposed because yield break in async iterators does not trigger using disposal.
+    /// </summary>
+    [Fact]
+    public async Task WorkflowRunActivity_IsStopped_LockstepAsync()
+    {
+        // Arrange
+        using var testActivity = new Activity("WorkflowRunStopTest_Lockstep").Start();
+
+        // Act
+        var workflow = CreateWorkflow();
+        Run run = await InProcessExecution.Lockstep.RunAsync(workflow, "Hello, World!");
+        await run.DisposeAsync();
+
+        // Assert - workflow.session should have been started and stopped
+        var startedSessions = this._startedActivities
+            .Where(a => a.RootId == testActivity.RootId &&
+                        a.OperationName.StartsWith(ActivityNames.WorkflowSession, StringComparison.Ordinal))
+            .ToList();
+        startedSessions.Should().HaveCount(1, "workflow.session Activity should be started");
+
+        var stoppedSessions = this._stoppedActivities
+            .Where(a => a.RootId == testActivity.RootId &&
+                        a.OperationName.StartsWith(ActivityNames.WorkflowSession, StringComparison.Ordinal))
+            .ToList();
+        stoppedSessions.Should().HaveCount(1,
+            "workflow.session Activity should be stopped/disposed so it is exported to telemetry backends");
+
+        // Assert - workflow_invoke should have been started and stopped
+        var startedWorkflowRuns = this._startedActivities
+            .Where(a => a.RootId == testActivity.RootId &&
+                        a.OperationName.StartsWith(ActivityNames.WorkflowInvoke, StringComparison.Ordinal))
+            .ToList();
+        startedWorkflowRuns.Should().HaveCount(1, "workflow_invoke Activity should be started");
+
+        var stoppedWorkflowRuns = this._stoppedActivities
+            .Where(a => a.RootId == testActivity.RootId &&
+                        a.OperationName.StartsWith(ActivityNames.WorkflowInvoke, StringComparison.Ordinal))
+            .ToList();
+        stoppedWorkflowRuns.Should().HaveCount(1,
+            "workflow_invoke Activity should be stopped/disposed so it is exported to telemetry backends (issue #4155)");
+    }
+
+    /// <summary>
+    /// Verifies that the workflow_invoke Activity is stopped when using the OffThread (Default)
+    /// execution environment (StreamingRunEventStream).
+    /// </summary>
+    [Fact]
+    public async Task WorkflowRunActivity_IsStopped_OffThreadAsync()
+    {
+        // Arrange
+        using var testActivity = new Activity("WorkflowRunStopTest_OffThread").Start();
+
+        // Act
+        var workflow = CreateWorkflow();
+        Run run = await InProcessExecution.OffThread.RunAsync(workflow, "Hello, World!");
+        await run.DisposeAsync();
+
+        // Assert - workflow.session should have been started and stopped
+        var startedSessions = this._startedActivities
+            .Where(a => a.RootId == testActivity.RootId &&
+                        a.OperationName.StartsWith(ActivityNames.WorkflowSession, StringComparison.Ordinal))
+            .ToList();
+        startedSessions.Should().HaveCount(1, "workflow.session Activity should be started");
+
+        var stoppedSessions = this._stoppedActivities
+            .Where(a => a.RootId == testActivity.RootId &&
+                        a.OperationName.StartsWith(ActivityNames.WorkflowSession, StringComparison.Ordinal))
+            .ToList();
+        stoppedSessions.Should().HaveCount(1,
+            "workflow.session Activity should be stopped/disposed so it is exported to telemetry backends");
+
+        // Assert - workflow_invoke should have been started and stopped
+        var startedWorkflowRuns = this._startedActivities
+            .Where(a => a.RootId == testActivity.RootId &&
+                        a.OperationName.StartsWith(ActivityNames.WorkflowInvoke, StringComparison.Ordinal))
+            .ToList();
+        startedWorkflowRuns.Should().HaveCount(1, "workflow_invoke Activity should be started");
+
+        var stoppedWorkflowRuns = this._stoppedActivities
+            .Where(a => a.RootId == testActivity.RootId &&
+                        a.OperationName.StartsWith(ActivityNames.WorkflowInvoke, StringComparison.Ordinal))
+            .ToList();
+        stoppedWorkflowRuns.Should().HaveCount(1,
+            "workflow_invoke Activity should be stopped/disposed so it is exported to telemetry backends (issue #4155)");
+    }
+
+    /// <summary>
+    /// Verifies that the workflow_invoke Activity is stopped when using the streaming API
+    /// (StreamingRun.WatchStreamAsync) with the OffThread execution environment.
+    /// This matches the exact usage pattern described in the issue.
+    /// </summary>
+    [Fact]
+    public async Task WorkflowRunActivity_IsStopped_Streaming_OffThreadAsync()
+    {
+        // Arrange
+        using var testActivity = new Activity("WorkflowRunStopTest_Streaming_OffThread").Start();
+
+        // Act - use streaming path (WatchStreamAsync), which is the pattern from the issue
+        var workflow = CreateWorkflow();
+        StreamingRun run = await InProcessExecution.OffThread.RunStreamingAsync(workflow, "Hello, World!");
+        await foreach (WorkflowEvent evt in run.WatchStreamAsync())
+        {
+            // Consume all events
+        }
+
+        // Dispose the run before asserting — the run Activity is disposed when the
+        // run loop exits, which happens during DisposeAsync. Without this, assertions
+        // can race against the background run loop's finally block.
+        await run.DisposeAsync();
+
+        // Assert - workflow.session should have been started
+        var startedSessions = this._startedActivities
+            .Where(a => a.RootId == testActivity.RootId &&
+                        a.OperationName.StartsWith(ActivityNames.WorkflowSession, StringComparison.Ordinal))
+            .ToList();
+        startedSessions.Should().HaveCount(1, "workflow.session Activity should be started");
+
+        // Assert - workflow_invoke should have been started
+        var startedWorkflowRuns = this._startedActivities
+            .Where(a => a.RootId == testActivity.RootId &&
+                        a.OperationName.StartsWith(ActivityNames.WorkflowInvoke, StringComparison.Ordinal))
+            .ToList();
+        startedWorkflowRuns.Should().HaveCount(1, "workflow_invoke Activity should be started");
+
+        // Assert - workflow_invoke should have been stopped
+        var stoppedWorkflowRuns = this._stoppedActivities
+            .Where(a => a.RootId == testActivity.RootId &&
+                        a.OperationName.StartsWith(ActivityNames.WorkflowInvoke, StringComparison.Ordinal))
+            .ToList();
+        stoppedWorkflowRuns.Should().HaveCount(1,
+            "workflow_invoke Activity should be stopped/disposed so it is exported to telemetry backends (issue #4155)");
+    }
+
+    /// <summary>
+    /// Verifies that a new workflow_invoke activity is started and stopped for each
+    /// streaming invocation, even when using the same workflow in a multi-turn pattern,
+    /// and that each session gets its own session activity.
+    /// </summary>
+    [Fact]
+    public async Task WorkflowRunActivity_IsStopped_Streaming_OffThread_MultiTurnAsync()
+    {
+        // Arrange
+        using var testActivity = new Activity("WorkflowRunStopTest_Streaming_OffThread_MultiTurn").Start();
+
+        var workflow = CreateWorkflow();
+
+        // Act - first streaming run
+        await using (StreamingRun run1 = await InProcessExecution.OffThread.RunStreamingAsync(workflow, "Hello, World!"))
+        {
+            await foreach (WorkflowEvent evt in run1.WatchStreamAsync())
+            {
+                // Consume all events from first turn
+            }
+        }
+
+        // Act - second streaming run (multi-turn scenario with same workflow)
+        await using (StreamingRun run2 = await InProcessExecution.OffThread.RunStreamingAsync(workflow, "Second turn!"))
+        {
+            await foreach (WorkflowEvent evt in run2.WatchStreamAsync())
+            {
+                // Consume all events from second turn
+            }
+        }
+
+        // Assert - two workflow.session activities should have been started and stopped
+        var startedSessions = this._startedActivities
+            .Where(a => a.RootId == testActivity.RootId &&
+                        a.OperationName.StartsWith(ActivityNames.WorkflowSession, StringComparison.Ordinal))
+            .ToList();
+        startedSessions.Should().HaveCount(2,
+            "each streaming invocation should start its own workflow.session Activity");
+
+        var stoppedSessions = this._stoppedActivities
+            .Where(a => a.RootId == testActivity.RootId &&
+                        a.OperationName.StartsWith(ActivityNames.WorkflowSession, StringComparison.Ordinal))
+            .ToList();
+        stoppedSessions.Should().HaveCount(2,
+            "each workflow.session Activity should be stopped/disposed so it is exported to telemetry backends");
+
+        // Assert - two workflow_invoke activities should have been started and stopped
+        var startedWorkflowRuns = this._startedActivities
+            .Where(a => a.RootId == testActivity.RootId &&
+                        a.OperationName.StartsWith(ActivityNames.WorkflowInvoke, StringComparison.Ordinal))
+            .ToList();
+        startedWorkflowRuns.Should().HaveCount(2,
+            "each streaming invocation should start its own workflow_invoke Activity");
+
+        var stoppedWorkflowRuns = this._stoppedActivities
+            .Where(a => a.RootId == testActivity.RootId &&
+                        a.OperationName.StartsWith(ActivityNames.WorkflowInvoke, StringComparison.Ordinal))
+            .ToList();
+        stoppedWorkflowRuns.Should().HaveCount(2,
+            "each workflow_invoke Activity should be stopped/disposed so it is exported to telemetry backends in multi-turn scenarios");
+    }
+
+    /// <summary>
+    /// Verifies that all started activities (not just workflow_invoke) are properly stopped.
+    /// This ensures no spans are "leaked" without being exported.
+    /// </summary>
+    [Fact]
+    public async Task AllActivities_AreStopped_AfterWorkflowCompletionAsync()
+    {
+        // Arrange
+        using var testActivity = new Activity("AllActivitiesStopTest").Start();
+
+        // Act
+        var workflow = CreateWorkflow();
+        Run run = await InProcessExecution.Lockstep.RunAsync(workflow, "Hello, World!");
+        await run.DisposeAsync();
+
+        // Assert - every started activity should also be stopped
+        var started = this._startedActivities
+            .Where(a => a.RootId == testActivity.RootId)
+            .Select(a => a.Id)
+            .ToHashSet();
+
+        var stopped = this._stoppedActivities
+            .Where(a => a.RootId == testActivity.RootId)
+            .Select(a => a.Id)
+            .ToHashSet();
+
+        var neverStopped = started.Except(stopped).ToList();
+        if (neverStopped.Count > 0)
+        {
+            var neverStoppedNames = this._startedActivities
+                .Where(a => neverStopped.Contains(a.Id))
+                .Select(a => a.OperationName)
+                .ToList();
+            neverStoppedNames.Should().BeEmpty(
+                "all started activities should be stopped so they are exported. " +
+                $"Activities started but never stopped: [{string.Join(", ", neverStoppedNames)}]");
+        }
+    }
+
+    /// <summary>
+    /// Verifies that Activity.Current is not leaked after lockstep RunAsync.
+    /// Application code creating activities after RunAsync returns should not
+    /// be parented under the workflow session span. The run activity should
+    /// still nest correctly under the session.
+    /// </summary>
+    [Fact]
+    public async Task Lockstep_SessionActivity_DoesNotLeak_IntoCaller_ActivityCurrentAsync()
+    {
+        // Arrange
+        using var testActivity = new Activity("SessionLeakTest").Start();
+        var workflow = CreateWorkflow();
+
+        // Act — run the workflow via lockstep (Start + drain happen inside RunAsync)
+        Run run = await InProcessExecution.Lockstep.RunAsync(workflow, "Hello, World!");
+
+        // Create an application activity after RunAsync returns.
+        // If the session leaked into Activity.Current, this would be parented under it.
+        using var appActivity = new Activity("AppWork").Start();
+        appActivity.Stop();
+
+        await run.DisposeAsync();
+
+        // Assert — the app activity should be parented under the test root, not the session
+        var sessionActivities = this._startedActivities
+            .Where(a => a.RootId == testActivity.RootId &&
+                        a.OperationName.StartsWith(ActivityNames.WorkflowSession, StringComparison.Ordinal))
+            .ToList();
+        sessionActivities.Should().HaveCount(1, "one session activity should exist");
+
+        appActivity.ParentId.Should().Be(testActivity.Id,
+            "application activity should be parented under the test root, not the workflow session");
+
+        // Assert — the run activity should still be parented under the session
+        var invokeActivities = this._startedActivities
+            .Where(a => a.RootId == testActivity.RootId &&
+                        a.OperationName.StartsWith(ActivityNames.WorkflowInvoke, StringComparison.Ordinal))
+            .ToList();
+        invokeActivities.Should().HaveCount(1, "one workflow_invoke activity should exist");
+        invokeActivities[0].ParentId.Should().Be(sessionActivities[0].Id,
+            "workflow_invoke activity should be nested under the session activity");
+    }
+}