test: WIP

2026-06-16 21:04:09 +08:00 · 2025-12-05 16:55:19 -05:00
parent 5b5960a546
commit 4011c97b75
5 changed files with 481 additions and 9 deletions
@@ -53,6 +53,9 @@ internal sealed class LockstepRunEventStream : IRunEventStream
        using Activity? activity = this._stepRunner.TelemetryContext.StartWorkflowRunActivity();
        activity?.SetTag(Tags.WorkflowId, this._stepRunner.StartExecutorId).SetTag(Tags.RunId, this._stepRunner.RunId);

+        bool hadException = false;
+        bool hadCancellation = false;
+
        try
        {
            this.RunStatus = RunStatus.Running;
@@ -74,14 +77,21 @@ internal sealed class LockstepRunEventStream : IRunEventStream
                    }
                    catch (OperationCanceledException)
                    {
+                        hadCancellation = true;
                    }
-                    catch (Exception ex) when (activity is not null)
+                    catch (Exception ex)
                    {
-                        activity.AddEvent(new ActivityEvent(EventNames.WorkflowError, tags: new() {
-                             { Tags.ErrorType, ex.GetType().FullName },
-                             { Tags.BuildErrorMessage, ex.Message },
-                        }));
-                        activity.CaptureException(ex);
+                        hadException = true;
+
+                        if (activity != null)
+                        {
+                            activity.AddEvent(new ActivityEvent(EventNames.WorkflowError, tags: new() {
+                                 { Tags.ErrorType, ex.GetType().FullName },
+                                 { Tags.BuildErrorMessage, ex.Message },
+                            }));
+                            activity.CaptureException(ex);
+                        }
+
                        throw;
                    }

@@ -133,7 +143,19 @@ internal sealed class LockstepRunEventStream : IRunEventStream
        }
        finally
        {
-            this.RunStatus = this._stepRunner.HasUnservicedRequests ? RunStatus.PendingRequests : RunStatus.Idle;
+            if (hadException || hadCancellation || linkedSource.Token.IsCancellationRequested)
+            {
+                this.RunStatus = RunStatus.Ended;
+            }
+            else if (this._stepRunner.HasUnservicedRequests)
+            {
+                this.RunStatus = RunStatus.PendingRequests;
+            }
+            else
+            {
+                this.RunStatus = RunStatus.Idle;
+            }
+
            this._stepRunner.OutgoingEvents.EventRaised -= OnWorkflowEventAsync;
        }

@@ -118,11 +118,12 @@ internal sealed class StreamingRunEventStream : IRunEventStream
        }
        finally
        {
+            // Mark as ended when run loop exits
+            this._runStatus = RunStatus.Ended;
+
            this._stepRunner.OutgoingEvents.EventRaised -= OnEventRaisedAsync;
            this._eventChannel.Writer.Complete();

-            // Mark as ended when run loop exits
-            this._runStatus = RunStatus.Ended;
            activity?.AddEvent(new ActivityEvent(EventNames.WorkflowCompleted));
        }

@@ -0,0 +1,29 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.Threading.Tasks;
+using System.Threading.Tasks.Sources;
+
+namespace Microsoft.Agents.AI.Workflows.UnitTests;
+
+internal sealed class DelayValueTaskSource<T> : IValueTaskSource<T>
+{
+    private readonly TestValueTaskSource<T> _innerSource = new();
+    private readonly T _value;
+
+    public DelayValueTaskSource(T value)
+    {
+        this._value = value;
+    }
+
+    public ValueTask ReleaseSucceededAsync() => this._innerSource.SetSucceededAsync(this._value);
+    public ValueTask ReleaseFaultedAsync(Exception exception) => this._innerSource.SetFaultedAsync(exception);
+    public ValueTask ReleaseCanceledAsync() => this._innerSource.SetCanceledAsync();
+
+    public T GetResult(short token) => this._innerSource.GetResult(token);
+
+    public ValueTaskSourceStatus GetStatus(short token) => this._innerSource.GetStatus(token);
+
+    public void OnCompleted(Action<object?> continuation, object? state, short token, ValueTaskSourceOnCompletedFlags flags)
+        => this._innerSource.OnCompleted(continuation, state, token, flags);
+}
@@ -0,0 +1,290 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.Diagnostics;
+using System.Runtime.CompilerServices;
+using System.Threading;
+using System.Threading.Tasks;
+using FluentAssertions;
+using Microsoft.Agents.AI.Workflows.Execution;
+
+namespace Microsoft.Agents.AI.Workflows.UnitTests;
+
+public static class RunStatusTests
+{
+    internal sealed class TestStepRunner : ISuperStepRunner
+    {
+        public TestStepRunner([CallerMemberName] string? name = null)
+        {
+            Console.WriteLine($"Starting test {name}");
+        }
+
+        public string RunId { get; } = Guid.NewGuid().ToString("N");
+
+        public string StartExecutorId { get; } = "start";
+
+        public bool HasUnservicedRequests { get; set; }
+        public bool HasUnprocessedMessages { get; set; }
+
+        public ConcurrentEventSink OutgoingEvents { get; } = new();
+
+        public ValueTask<bool> EnqueueMessageAsync<T>(T message, CancellationToken cancellationToken = default)
+        {
+            this.HasUnprocessedMessages = true;
+            return new(true);
+        }
+
+        ValueTask<bool> ISuperStepRunner.EnqueueMessageUntypedAsync(object message, Type declaredType, CancellationToken cancellationToken)
+        {
+            this.HasUnprocessedMessages = true;
+            return new(true);
+        }
+
+        public async ValueTask EnqueueResponseAsync(ExternalResponse response, CancellationToken cancellationToken = default)
+        {
+            this.HasUnservicedRequests = false;
+            await this.EnqueueMessageAsync(response, cancellationToken);
+        }
+
+        public ValueTask<bool> IsValidInputTypeAsync<T>(CancellationToken cancellationToken = default) => new(true);
+
+        public ValueTask RequestEndRunAsync()
+        {
+            if (this._currentStepSource != null)
+            {
+                return this.CancelStepAsync();
+            }
+
+            return new();
+        }
+
+        private DelayValueTaskSource<bool>? _currentStepSource;
+        private CancellationTokenRegistration? _registration;
+
+        ValueTask<bool> ISuperStepRunner.RunSuperStepAsync(CancellationToken cancellationToken)
+        {
+            Debug.Assert(Interlocked.CompareExchange(ref this._currentStepSource,
+                                                     value: new DelayValueTaskSource<bool>(this.HasUnprocessedMessages),
+                                                     null) is null);
+
+            this._registration = cancellationToken.Register(() => _ = this._currentStepSource == null
+                                                                    ? Task.CompletedTask
+                                                                    : this._currentStepSource.ReleaseCanceledAsync().AsTask());
+            this.HasUnprocessedMessages = false;
+
+            return new(this._currentStepSource, 0);
+        }
+
+        private DelayValueTaskSource<bool> TakeCurrentStepSource()
+        {
+            DelayValueTaskSource<bool>? currentStepSource = Interlocked.Exchange(ref this._currentStepSource, null);
+            Debug.Assert(currentStepSource is not null);
+            this._registration?.Dispose();
+            this._registration = null;
+
+            return currentStepSource;
+        }
+
+        public ValueTask CompleteStepAsync() => this.TakeCurrentStepSource().ReleaseSucceededAsync();
+
+        public ValueTask CompleteStepWithPendingAsync()
+        {
+            this.HasUnservicedRequests = true;
+            return this.CompleteStepAsync();
+        }
+
+        public ValueTask CancelStepAsync() => this.TakeCurrentStepSource().ReleaseCanceledAsync();
+
+        public ValueTask FailStepAsync(Exception exception) => this.TakeCurrentStepSource().ReleaseFaultedAsync(exception);
+    }
+
+    public enum EventStreamKind
+    {
+        OffThread,
+        Lockstep
+    }
+
+    private static IRunEventStream GetRunStreamForKind(EventStreamKind kind, ISuperStepRunner stepRunner)
+    {
+        IRunEventStream result;
+        switch (kind)
+        {
+            case EventStreamKind.OffThread:
+                result = new StreamingRunEventStream(stepRunner);
+                break;
+            case EventStreamKind.Lockstep:
+                result = new LockstepRunEventStream(stepRunner);
+                break;
+            default:
+                throw new NotSupportedException($"Unsupported RunStream kind: {kind}");
+        }
+
+        result.Start();
+        return result;
+    }
+
+    [Theory]
+    [InlineData(EventStreamKind.OffThread)]
+    [InlineData(EventStreamKind.Lockstep)]
+    public static async Task Test_RunStatus_NotStartedWhenStartingAsync(EventStreamKind mode)
+    {
+        TestStepRunner runner = new();
+        IRunEventStream eventStream = GetRunStreamForKind(mode, runner);
+
+        RunStatus status = await eventStream.GetStatusAsync();
+        status.Should().Be(RunStatus.NotStarted);
+    }
+
+    [Theory]
+    [InlineData(EventStreamKind.OffThread)]
+    [InlineData(EventStreamKind.Lockstep)]
+    public static async Task Test_RunStatus_RunningWhenInSuperstepAsync(EventStreamKind mode)
+    {
+        TestStepRunner runner = new();
+        IRunEventStream eventStream = GetRunStreamForKind(mode, runner);
+
+        await runner.EnqueueMessageAsync(new object());
+        eventStream.SignalInput();
+
+        _ = WatchStreamAsync();
+
+        RunStatus status = await eventStream.GetStatusAsync();
+        status.Should().Be(RunStatus.Running);
+
+        await eventStream.DisposeAsync();
+
+        async Task WatchStreamAsync()
+        {
+            await foreach (var _ in eventStream.TakeEventStreamAsync(false)) { }
+        }
+    }
+
+    [Theory]
+    [InlineData(EventStreamKind.OffThread)]
+    [InlineData(EventStreamKind.Lockstep)]
+    public static async Task Test_RunStatus_IdleWhenFinishedSuperstepsAsync(EventStreamKind mode)
+    {
+        TestStepRunner runner = new();
+        IRunEventStream eventStream = GetRunStreamForKind(mode, runner);
+
+        await runner.EnqueueMessageAsync(new object());
+        eventStream.SignalInput();
+
+        Task watchTask = WatchStreamAsync();
+        await runner.CompleteStepAsync();
+        await watchTask;
+
+        RunStatus status = await eventStream.GetStatusAsync();
+        status.Should().Be(RunStatus.Idle);
+
+        await eventStream.DisposeAsync();
+
+        async Task WatchStreamAsync()
+        {
+            await foreach (var _ in eventStream.TakeEventStreamAsync(false)) { }
+        }
+    }
+
+    [Theory]
+    [InlineData(EventStreamKind.OffThread)]
+    [InlineData(EventStreamKind.Lockstep)]
+    public static async Task Test_RunStatus_EndedWhenCancelledAsync(EventStreamKind mode)
+    {
+        TestStepRunner runner = new();
+        IRunEventStream eventStream = GetRunStreamForKind(mode, runner);
+
+        await runner.EnqueueMessageAsync(new object());
+        eventStream.SignalInput();
+
+        Task watchTask = WatchStreamAsync();
+        await runner.CancelStepAsync();
+        await watchTask;
+
+        RunStatus status = await eventStream.GetStatusAsync();
+        status.Should().Be(RunStatus.Ended);
+
+        await eventStream.DisposeAsync();
+
+        async Task WatchStreamAsync()
+        {
+            await foreach (var _ in eventStream.TakeEventStreamAsync(false)) { }
+        }
+    }
+
+    [Theory]
+    [InlineData(EventStreamKind.OffThread)]
+    [InlineData(EventStreamKind.Lockstep)]
+    public static async Task Test_RunStatus_ExceptionWhenFaultedAsync(EventStreamKind mode)
+    {
+        TestStepRunner runner = new();
+        IRunEventStream eventStream = GetRunStreamForKind(mode, runner);
+
+        await runner.EnqueueMessageAsync(new object());
+        eventStream.SignalInput();
+
+        Task watchTask = WatchStreamAsync();
+        await runner.FailStepAsync(new InvalidOperationException());
+        await watchTask;
+
+        RunStatus status = await eventStream.GetStatusAsync();
+        status.Should().Be(RunStatus.Ended);
+
+        await eventStream.DisposeAsync();
+
+        async Task WatchStreamAsync()
+        {
+            await foreach (var _ in eventStream.TakeEventStreamAsync(false)) { }
+        }
+    }
+
+    //[Theory]
+    //[InlineData(EventStreamKind.OffThread)]
+    //[InlineData(EventStreamKind.Lockstep)]
+    internal static async Task Test_RunStatus_PendingRequestsAsync(EventStreamKind mode)
+    {
+        TestStepRunner runner = new();
+        IRunEventStream eventStream = GetRunStreamForKind(mode, runner);
+
+        // Act 1: Send the input object, and run the step to PendingRequest
+        await runner.EnqueueMessageAsync(new object());
+        eventStream.SignalInput();
+
+        Task watchTask = WatchStreamAsync();
+        await runner.CompleteStepWithPendingAsync();
+        await watchTask;
+
+        // Assert 1
+        RunStatus status = await eventStream.GetStatusAsync();
+        status.Should().Be(RunStatus.PendingRequests);
+
+        // Act 2: Send the response, check running state
+        await runner.EnqueueResponseAsync(
+            new ExternalResponse(
+                new Checkpointing.RequestPortInfo(new(typeof(object)), new(typeof(object)), "_"),
+                Guid.NewGuid().ToString("N"),
+                new(new())));
+        eventStream.SignalInput();
+
+        watchTask = WatchStreamAsync();
+
+        // Assert 2
+        status = await eventStream.GetStatusAsync();
+        status.Should().Be(RunStatus.Running);
+
+        // Act 3: Process the response, check state is idle
+        await runner.CompleteStepAsync();
+        await watchTask; status = await eventStream.GetStatusAsync();
+        status.Should().Be(RunStatus.Running);
+
+        // Assert 3
+        status = await eventStream.GetStatusAsync();
+        status.Should().Be(RunStatus.Idle);
+
+        await eventStream.DisposeAsync();
+
+        async Task WatchStreamAsync()
+        {
+            await foreach (var _ in eventStream.TakeEventStreamAsync(false)) { }
+        }
+    }
+}
@@ -0,0 +1,130 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.Diagnostics;
+using System.Threading;
+using System.Threading.Tasks;
+using System.Threading.Tasks.Sources;
+
+namespace Microsoft.Agents.AI.Workflows.UnitTests;
+
+internal sealed class TestValueTaskSource<T> : IValueTaskSource<T>
+{
+    private int _status = (int)ValueTaskSourceStatus.Pending;
+    private T? _value;
+    private Exception? _exception;
+    private int _continuationScheduled;
+    private readonly object _continuationMutex = new();
+
+    private bool _ranContinuation;
+    private Action _continuationClosure = () => { };
+
+    public TestValueTaskSource()
+    {
+    }
+
+    private bool TrySetCompletionStatus(ValueTaskSourceStatus status)
+        => Interlocked.CompareExchange(ref this._status,
+                                       value: (int)status,
+                                       comparand: (int)ValueTaskSourceStatus.Pending)
+                                               == (int)ValueTaskSourceStatus.Pending;
+
+    private void RunScheduledContinuation()
+    {
+        Console.WriteLine("Running scheduled continuation");
+        lock (this._continuationMutex)
+        {
+            this._ranContinuation = true;
+            this._continuationClosure();
+        }
+    }
+
+    public ValueTask SetSucceededAsync(T value)
+    {
+        Console.WriteLine($"Setting succeeded {value}");
+
+        if (this.TrySetCompletionStatus(ValueTaskSourceStatus.Succeeded))
+        {
+            // If the status was Pending, we can set it
+            this._value = value;
+        }
+
+        this.RunScheduledContinuation();
+        return new();
+    }
+
+    public ValueTask SetFaultedAsync(Exception exception)
+    {
+        Console.WriteLine($"Setting faulted {exception}");
+
+        if (this.TrySetCompletionStatus(ValueTaskSourceStatus.Faulted))
+        {
+            // If the status was Pending, we can set it
+            this._exception = exception;
+        }
+
+        this.RunScheduledContinuation();
+        return new();
+    }
+
+    public ValueTask SetCanceledAsync()
+    {
+        Console.WriteLine("Setting canceled");
+
+        this.TrySetCompletionStatus(ValueTaskSourceStatus.Canceled);
+        this.RunScheduledContinuation();
+        return new();
+    }
+
+    public T GetResult(short token)
+    {
+        Debug.Assert(token == 0);
+
+        switch (this.GetStatus(0))
+        {
+            case ValueTaskSourceStatus.Succeeded:
+                return this._value!;
+            case ValueTaskSourceStatus.Faulted:
+                throw this._exception!;
+            case ValueTaskSourceStatus.Canceled:
+                throw new TaskCanceledException();
+            case ValueTaskSourceStatus.Pending:
+                throw new InvalidOperationException("The operation is not yet completed.");
+            default:
+                throw new NotSupportedException();
+        }
+    }
+
+    public ValueTaskSourceStatus GetStatus(short token)
+    {
+        Debug.Assert(token == 0);
+
+        return (ValueTaskSourceStatus)Volatile.Read(ref this._status);
+    }
+
+    public void OnCompleted(Action<object?> continuation, object? state, short token, ValueTaskSourceOnCompletedFlags flags)
+    {
+        Debug.Assert(token == 0);
+
+        if (Interlocked.Exchange(ref this._continuationScheduled, 1) == 1)
+        {
+            throw new InvalidOperationException("Cannot schedule more than one continuation on ValueTaskSource");
+        }
+
+        lock (this._continuationMutex)
+        {
+            if (this._ranContinuation)
+            {
+                // The default no-op was run, since we have not yet scheduled a continuation
+                // Run this continuation immediately
+                Console.WriteLine("Running continuation");
+                continuation(state);
+            }
+            else
+            {
+                Console.WriteLine("Scheduling continuation");
+                this._continuationClosure = () => continuation(state);
+            }
+        }
+    }
+}