using System;
using System.Collections.Generic;
using System.IO;
using System.IO.Pipes;
using System.Runtime.InteropServices;
using System.Threading;
using System.Threading.Tasks;
using Google.Protobuf.WellKnownTypes;
using ZB.MOM.WW.MxGateway.Contracts;
using ZB.MOM.WW.MxGateway.Contracts.Proto;
using ZB.MOM.WW.MxGateway.Worker.Ipc;
using ZB.MOM.WW.MxGateway.Worker.MxAccess;
using ZB.MOM.WW.MxGateway.Worker.Tests.TestSupport;

namespace ZB.MOM.WW.MxGateway.Worker.Tests.Ipc;

public sealed class WorkerPipeSessionTests
{
    private const string SessionId = "session-1";
    private const string Nonce = "nonce-secret";

    /// <summary>Verifies that valid gateway hello triggers worker hello and ready responses.</summary>
    [Fact]
    public async Task CompleteStartupHandshakeAsync_WithValidGatewayHello_SendsHelloThenReady()
    {
        WorkerFrameProtocolOptions options = CreateOptions();
        using MemoryStream inbound = new();
        await new WorkerFrameWriter(inbound, options).WriteAsync(CreateGatewayHelloEnvelope());
        inbound.Position = 0;
        using MemoryStream outbound = new();
        WorkerPipeSession session = CreateSession(inbound, outbound, options);
        bool initialized = false;

        await session.CompleteStartupHandshakeAsync(
            _ =>
            {
                initialized = true;
                return Task.CompletedTask;
            });

        Assert.True(initialized);
        WorkerEnvelope[] written = ReadWrittenFrames(outbound, options);
        Assert.Equal(2, written.Length);
        Assert.Equal(WorkerEnvelope.BodyOneofCase.WorkerHello, written[0].BodyCase);
        Assert.Equal(WorkerEnvelope.BodyOneofCase.WorkerReady, written[1].BodyCase);
        Assert.Equal(Nonce, written[0].WorkerHello.Nonce);
        Assert.Equal(1234, written[1].WorkerReady.WorkerProcessId);
        Assert.Equal(ZB.MOM.WW.MxGateway.Worker.MxAccess.MxAccessInteropInfo.ProgId, written[1].WorkerReady.MxaccessProgid);
        Assert.Equal(ZB.MOM.WW.MxGateway.Worker.MxAccess.MxAccessInteropInfo.Clsid, written[1].WorkerReady.MxaccessClsid);
        Assert.NotNull(written[1].WorkerReady.ReadyTimestamp);
    }

    /// <summary>Verifies that wrong nonce causes protocol violation fault before initialization.</summary>
    [Fact]
    public async Task CompleteStartupHandshakeAsync_WithWrongNonce_FaultsBeforeInitialization()
    {
        WorkerFrameProtocolOptions options = CreateOptions();
        using MemoryStream inbound = new();
        await new WorkerFrameWriter(inbound, options).WriteAsync(CreateGatewayHelloEnvelope(nonce: "wrong"));
        inbound.Position = 0;
        using MemoryStream outbound = new();
        WorkerPipeSession session = CreateSession(inbound, outbound, options);
        bool initialized = false;

        WorkerFrameProtocolException exception =
            await Assert.ThrowsAsync<WorkerFrameProtocolException>(
                async () => await session.CompleteStartupHandshakeAsync(
                    _ =>
                    {
                        initialized = true;
                        return Task.CompletedTask;
                    }));

        Assert.False(initialized);
        Assert.Equal(WorkerFrameProtocolErrorCode.NonceMismatch, exception.ErrorCode);
        WorkerEnvelope fault = Assert.Single(ReadWrittenFrames(outbound, options));
        Assert.Equal(WorkerEnvelope.BodyOneofCase.WorkerFault, fault.BodyCase);
        Assert.Equal(WorkerFaultCategory.ProtocolViolation, fault.WorkerFault.Category);
    }

    /// <summary>Verifies that unsupported protocol version causes mismatch fault before initialization.</summary>
    [Fact]
    public async Task CompleteStartupHandshakeAsync_WithWrongProtocol_FaultsBeforeInitialization()
    {
        WorkerFrameProtocolOptions options = CreateOptions();
        using MemoryStream inbound = new();
        await new WorkerFrameWriter(inbound, options).WriteAsync(CreateGatewayHelloEnvelope(supportedProtocolVersion: 999));
        inbound.Position = 0;
        using MemoryStream outbound = new();
        WorkerPipeSession session = CreateSession(inbound, outbound, options);
        bool initialized = false;

        WorkerFrameProtocolException exception =
            await Assert.ThrowsAsync<WorkerFrameProtocolException>(
                async () => await session.CompleteStartupHandshakeAsync(
                    _ =>
                    {
                        initialized = true;
                        return Task.CompletedTask;
                    }));

        Assert.False(initialized);
        Assert.Equal(WorkerFrameProtocolErrorCode.ProtocolVersionMismatch, exception.ErrorCode);
        WorkerEnvelope fault = Assert.Single(ReadWrittenFrames(outbound, options));
        Assert.Equal(WorkerFaultCategory.ProtocolMismatch, fault.WorkerFault.Category);
    }

    /// <summary>Verifies that malformed frame causes protocol violation fault.</summary>
    [Fact]
    public async Task CompleteStartupHandshakeAsync_WithMalformedFrame_WritesWorkerFault()
    {
        WorkerFrameProtocolOptions options = CreateOptions();
        using MemoryStream inbound = new(WorkerFrameTestHelpers.CreateFrame(new byte[] { 0x80 }));
        using MemoryStream outbound = new();
        WorkerPipeSession session = CreateSession(inbound, outbound, options);
        bool initialized = false;

        WorkerFrameProtocolException exception =
            await Assert.ThrowsAsync<WorkerFrameProtocolException>(
                async () => await session.CompleteStartupHandshakeAsync(
                    _ =>
                    {
                        initialized = true;
                        return Task.CompletedTask;
                    }));

        Assert.False(initialized);
        Assert.Equal(WorkerFrameProtocolErrorCode.InvalidEnvelope, exception.ErrorCode);
        WorkerEnvelope fault = Assert.Single(ReadWrittenFrames(outbound, options));
        Assert.Equal(WorkerEnvelope.BodyOneofCase.WorkerFault, fault.BodyCase);
        Assert.Equal(WorkerFaultCategory.ProtocolViolation, fault.WorkerFault.Category);
    }

    /// <summary>Verifies that MXAccess COM creation failure produces fault instead of ready.</summary>
    [Fact]
    public async Task CompleteStartupHandshakeAsync_WhenMxAccessCreationFails_WritesFaultInsteadOfReady()
    {
        const int hresult = unchecked((int)0x80040154);
        WorkerFrameProtocolOptions options = CreateOptions();
        using MemoryStream inbound = new();
        await new WorkerFrameWriter(inbound, options).WriteAsync(CreateGatewayHelloEnvelope());
        inbound.Position = 0;
        using MemoryStream outbound = new();
        WorkerPipeSession session = CreateSession(inbound, outbound, options);

        await Assert.ThrowsAsync<COMException>(
            async () => await session.CompleteStartupHandshakeAsync(
                _ => Task.FromException<WorkerReady>(new COMException("Class not registered.", hresult))));

        WorkerEnvelope[] written = ReadWrittenFrames(outbound, options);
        Assert.Equal(2, written.Length);
        Assert.Equal(WorkerEnvelope.BodyOneofCase.WorkerHello, written[0].BodyCase);
        Assert.Equal(WorkerEnvelope.BodyOneofCase.WorkerFault, written[1].BodyCase);
        Assert.Equal(WorkerFaultCategory.MxaccessCreationFailed, written[1].WorkerFault.Category);
        Assert.Equal(hresult, written[1].WorkerFault.Hresult);
        Assert.Equal(typeof(COMException).FullName, written[1].WorkerFault.ExceptionType);
        Assert.Equal(ProtocolStatusCode.WorkerUnavailable, written[1].WorkerFault.ProtocolStatus.Code);
    }

    /// <summary>Verifies that heartbeat payload reflects current runtime snapshot.</summary>
    [Fact]
    public async Task RunAsync_SendsHeartbeatPayloadFromRuntimeSnapshot()
    {
        using CancellationTokenSource cancellation = new(TimeSpan.FromSeconds(5));
        using PipePair pipePair = await PipePair.CreateAsync(cancellation.Token);
        FakeRuntimeSession runtime = new();
        runtime.SetSnapshot(new WorkerRuntimeHeartbeatSnapshot(
            DateTimeOffset.UtcNow,
            pendingCommandCount: 2,
            outboundEventQueueDepth: 3,
            lastEventSequence: 42,
            currentCommandCorrelationId: "current-command"));
        WorkerPipeSession session = CreatePipeSession(
            pipePair.WorkerStream,
            runtime,
            new WorkerPipeSessionOptions
            {
                HeartbeatInterval = TimeSpan.FromMilliseconds(20),
                HeartbeatGrace = TimeSpan.FromSeconds(5),
            });
        Task runTask = session.RunAsync(cancellation.Token);

        await CompleteGatewayHandshakeAsync(pipePair, cancellation.Token);

        // Deterministic race: read the first heartbeat while watching runTask.
        // A faulted RunAsync would complete the run task first; if it wins the
        // race the test fails immediately with the underlying fault instead of
        // waiting out an arbitrary fixed delay.
        Task<WorkerEnvelope> heartbeatTask = ReadUntilAsync(
            pipePair.GatewayReader,
            WorkerEnvelope.BodyOneofCase.WorkerHeartbeat,
            cancellation.Token);
        Task winner = await Task.WhenAny(runTask, heartbeatTask);
        if (winner == runTask)
        {
            // Surface the RunAsync fault (or assert it did not exit early).
            await runTask;
            Assert.Fail("RunAsync completed before the first heartbeat was received.");
        }

        WorkerEnvelope heartbeat = await heartbeatTask;

        Assert.Equal(WorkerState.ExecutingCommand, heartbeat.WorkerHeartbeat.State);
        Assert.Equal(1234, heartbeat.WorkerHeartbeat.WorkerProcessId);
        Assert.Equal(2u, heartbeat.WorkerHeartbeat.PendingCommandCount);
        Assert.Equal(3u, heartbeat.WorkerHeartbeat.OutboundEventQueueDepth);
        Assert.Equal(42UL, heartbeat.WorkerHeartbeat.LastEventSequence);
        Assert.Equal("current-command", heartbeat.WorkerHeartbeat.CurrentCommandCorrelationId);

        await SendShutdownAndWaitAsync(pipePair, runTask, cancellation.Token);
    }

    /// <summary>Verifies that heartbeat reports current command correlation during execution.</summary>
    [Fact]
    public async Task RunAsync_WhenCommandIsExecuting_HeartbeatReportsCurrentCorrelation()
    {
        using CancellationTokenSource cancellation = new(TimeSpan.FromSeconds(5));
        using PipePair pipePair = await PipePair.CreateAsync(cancellation.Token);
        FakeRuntimeSession runtime = new()
        {
            BlockDispatch = true,
        };
        WorkerPipeSession session = CreatePipeSession(
            pipePair.WorkerStream,
            runtime,
            new WorkerPipeSessionOptions
            {
                HeartbeatInterval = TimeSpan.FromMilliseconds(20),
                HeartbeatGrace = TimeSpan.FromSeconds(5),
            });
        Task runTask = session.RunAsync(cancellation.Token);
        await CompleteGatewayHandshakeAsync(pipePair, cancellation.Token);

        await pipePair.GatewayWriter.WriteAsync(
            CreateCommandEnvelope("command-1"),
            cancellation.Token);

        Assert.True(runtime.DispatchStarted.Wait(TimeSpan.FromSeconds(2)));
        WorkerEnvelope heartbeat = await ReadUntilAsync(
            pipePair.GatewayReader,
            WorkerEnvelope.BodyOneofCase.WorkerHeartbeat,
            envelope => envelope.WorkerHeartbeat.CurrentCommandCorrelationId == "command-1",
            cancellation.Token);

        Assert.Equal("command-1", heartbeat.WorkerHeartbeat.CurrentCommandCorrelationId);
        Assert.Equal(WorkerState.ExecutingCommand, heartbeat.WorkerHeartbeat.State);

        runtime.ReleaseDispatch();
        WorkerEnvelope reply = await ReadUntilAsync(
            pipePair.GatewayReader,
            WorkerEnvelope.BodyOneofCase.WorkerCommandReply,
            cancellation.Token);

        Assert.Equal("command-1", reply.CorrelationId);
        Assert.Equal(ProtocolStatusCode.Ok, reply.WorkerCommandReply.Reply.ProtocolStatus.Code);

        await SendShutdownAndWaitAsync(pipePair, runTask, cancellation.Token);
    }

    /// <summary>Verifies that worker events are written to the pipe.</summary>
    [Fact]
    public async Task RunAsync_WhenRuntimeHasEvents_WritesWorkerEventEnvelope()
    {
        using CancellationTokenSource cancellation = new(TimeSpan.FromSeconds(5));
        using PipePair pipePair = await PipePair.CreateAsync(cancellation.Token);
        FakeRuntimeSession runtime = new();
        WorkerPipeSession session = CreatePipeSession(
            pipePair.WorkerStream,
            runtime,
            new WorkerPipeSessionOptions
            {
                HeartbeatInterval = TimeSpan.FromMilliseconds(100),
                HeartbeatGrace = TimeSpan.FromSeconds(5),
            });
        Task runTask = session.RunAsync(cancellation.Token);
        await CompleteGatewayHandshakeAsync(pipePair, cancellation.Token);

        runtime.EnqueueEvent(CreateWorkerEvent(sequence: 7));

        WorkerEnvelope workerEvent = await ReadUntilAsync(
            pipePair.GatewayReader,
            WorkerEnvelope.BodyOneofCase.WorkerEvent,
            cancellation.Token);

        Assert.Equal(MxEventFamily.OnDataChange, workerEvent.WorkerEvent.Event.Family);
        Assert.Equal(7UL, workerEvent.WorkerEvent.Event.WorkerSequence);

        await SendShutdownAndWaitAsync(pipePair, runTask, cancellation.Token);
    }

    /// <summary>
    ///     Verifies that a Ping control command is answered on the worker side
    ///     (not dispatched to the STA) with an OK reply that echoes the ping
    ///     message into the reply's diagnostic field.
    /// </summary>
    [Fact]
    public async Task RunAsync_PingControlCommand_RepliesOkAndEchoesMessage()
    {
        using CancellationTokenSource cancellation = new(TimeSpan.FromSeconds(5));
        using PipePair pipePair = await PipePair.CreateAsync(cancellation.Token);
        FakeRuntimeSession runtime = new();
        WorkerPipeSession session = CreatePipeSession(pipePair.WorkerStream, runtime);
        Task runTask = session.RunAsync(cancellation.Token);
        await CompleteGatewayHandshakeAsync(pipePair, cancellation.Token);

        await pipePair.GatewayWriter
            .WriteAsync(CreatePingCommandEnvelope("ping-1", "hello-worker"), cancellation.Token);

        WorkerEnvelope replyEnvelope = await ReadUntilAsync(
            pipePair.GatewayReader,
            WorkerEnvelope.BodyOneofCase.WorkerCommandReply,
            cancellation.Token);

        MxCommandReply reply = replyEnvelope.WorkerCommandReply.Reply;
        Assert.Equal("ping-1", reply.CorrelationId);
        Assert.Equal(MxCommandKind.Ping, reply.Kind);
        Assert.Equal(ProtocolStatusCode.Ok, reply.ProtocolStatus.Code);
        Assert.Equal("hello-worker", reply.DiagnosticMessage);

        await SendShutdownAndWaitAsync(pipePair, runTask, cancellation.Token);
    }

    /// <summary>
    ///     Verifies that GetSessionState reports the worker's lifecycle as the
    ///     proto SessionState — READY while the message loop is serving.
    /// </summary>
    [Fact]
    public async Task RunAsync_GetSessionStateControlCommand_RepliesReady()
    {
        using CancellationTokenSource cancellation = new(TimeSpan.FromSeconds(5));
        using PipePair pipePair = await PipePair.CreateAsync(cancellation.Token);
        FakeRuntimeSession runtime = new();
        WorkerPipeSession session = CreatePipeSession(pipePair.WorkerStream, runtime);
        Task runTask = session.RunAsync(cancellation.Token);
        await CompleteGatewayHandshakeAsync(pipePair, cancellation.Token);

        await pipePair.GatewayWriter
            .WriteAsync(
                CreateControlCommandEnvelope(
                    "state-1",
                    MxCommandKind.GetSessionState,
                    command => command.GetSessionState = new GetSessionStateCommand()),
                cancellation.Token);

        WorkerEnvelope replyEnvelope = await ReadUntilAsync(
            pipePair.GatewayReader,
            WorkerEnvelope.BodyOneofCase.WorkerCommandReply,
            cancellation.Token);

        MxCommandReply reply = replyEnvelope.WorkerCommandReply.Reply;
        Assert.Equal(MxCommandKind.GetSessionState, reply.Kind);
        Assert.Equal(ProtocolStatusCode.Ok, reply.ProtocolStatus.Code);
        Assert.Equal(SessionState.Ready, reply.SessionState.State);

        await SendShutdownAndWaitAsync(pipePair, runTask, cancellation.Token);
    }

    /// <summary>
    ///     Verifies that GetWorkerInfo populates the worker process id, version,
    ///     and MXAccess ProgID/CLSID from the worker's own metadata.
    /// </summary>
    [Fact]
    public async Task RunAsync_GetWorkerInfoControlCommand_PopulatesWorkerInfoFields()
    {
        using CancellationTokenSource cancellation = new(TimeSpan.FromSeconds(5));
        using PipePair pipePair = await PipePair.CreateAsync(cancellation.Token);
        FakeRuntimeSession runtime = new();
        WorkerPipeSession session = CreatePipeSession(pipePair.WorkerStream, runtime);
        Task runTask = session.RunAsync(cancellation.Token);
        await CompleteGatewayHandshakeAsync(pipePair, cancellation.Token);

        await pipePair.GatewayWriter
            .WriteAsync(
                CreateControlCommandEnvelope(
                    "info-1",
                    MxCommandKind.GetWorkerInfo,
                    command => command.GetWorkerInfo = new GetWorkerInfoCommand()),
                cancellation.Token);

        WorkerEnvelope replyEnvelope = await ReadUntilAsync(
            pipePair.GatewayReader,
            WorkerEnvelope.BodyOneofCase.WorkerCommandReply,
            cancellation.Token);

        MxCommandReply reply = replyEnvelope.WorkerCommandReply.Reply;
        Assert.Equal(MxCommandKind.GetWorkerInfo, reply.Kind);
        Assert.Equal(ProtocolStatusCode.Ok, reply.ProtocolStatus.Code);
        WorkerInfoReply info = reply.WorkerInfo;
        Assert.Equal(1234, info.WorkerProcessId);
        Assert.False(string.IsNullOrEmpty(info.WorkerVersion));
        Assert.Equal(MxAccessInteropInfo.ProgId, info.MxaccessProgid);
        Assert.Equal(MxAccessInteropInfo.Clsid, info.MxaccessClsid);

        await SendShutdownAndWaitAsync(pipePair, runTask, cancellation.Token);
    }

    /// <summary>
    ///     Verifies that DrainEvents drains the runtime session's queued events
    ///     into the reply rather than streaming them as WorkerEvent envelopes.
    /// </summary>
    [Fact]
    public async Task RunAsync_DrainEventsControlCommand_ReturnsQueuedEvents()
    {
        using CancellationTokenSource cancellation = new(TimeSpan.FromSeconds(5));
        using PipePair pipePair = await PipePair.CreateAsync(cancellation.Token);
        // Suppress the background drain loop's fixed-batch drains so the
        // queued events survive for the explicit DrainEvents command (which
        // drains all via max_events == 0). 128 mirrors
        // WorkerPipeSession.EventDrainBatchSize.
        FakeRuntimeSession runtime = new() { SuppressDrainForBatchSize = 128 };
        WorkerPipeSession session = CreatePipeSession(pipePair.WorkerStream, runtime);
        runtime.EnqueueEvent(CreateWorkerEvent(sequence: 11));
        runtime.EnqueueEvent(CreateWorkerEvent(sequence: 12));
        Task runTask = session.RunAsync(cancellation.Token);
        await CompleteGatewayHandshakeAsync(pipePair, cancellation.Token);

        await pipePair.GatewayWriter
            .WriteAsync(
                CreateControlCommandEnvelope(
                    "drain-1",
                    MxCommandKind.DrainEvents,
                    command => command.DrainEvents = new DrainEventsCommand { MaxEvents = 0 }),
                cancellation.Token);

        WorkerEnvelope replyEnvelope = await ReadUntilAsync(
            pipePair.GatewayReader,
            WorkerEnvelope.BodyOneofCase.WorkerCommandReply,
            cancellation.Token);

        MxCommandReply reply = replyEnvelope.WorkerCommandReply.Reply;
        Assert.Equal(MxCommandKind.DrainEvents, reply.Kind);
        Assert.Equal(ProtocolStatusCode.Ok, reply.ProtocolStatus.Code);
        Assert.Equal(2, reply.DrainEvents.Events.Count);
        Assert.Contains(reply.DrainEvents.Events, e => e.WorkerSequence == 11UL);
        Assert.Contains(reply.DrainEvents.Events, e => e.WorkerSequence == 12UL);

        await SendShutdownAndWaitAsync(pipePair, runTask, cancellation.Token);
    }

    /// <summary>
    ///     Verifies that ShutdownWorker returns its OK reply BEFORE the graceful
    ///     shutdown runs and disposes the runtime session, and that the message
    ///     loop then stops.
    /// </summary>
    [Fact]
    public async Task RunAsync_ShutdownWorkerControlCommand_RepliesOkThenShutsDown()
    {
        using CancellationTokenSource cancellation = new(TimeSpan.FromSeconds(5));
        using PipePair pipePair = await PipePair.CreateAsync(cancellation.Token);
        FakeRuntimeSession runtime = new();
        WorkerPipeSession session = CreatePipeSession(pipePair.WorkerStream, runtime);
        Task runTask = session.RunAsync(cancellation.Token);
        await CompleteGatewayHandshakeAsync(pipePair, cancellation.Token);

        await pipePair.GatewayWriter
            .WriteAsync(
                CreateControlCommandEnvelope(
                    "shutdown-1",
                    MxCommandKind.ShutdownWorker,
                    command => command.ShutdownWorker = new ShutdownWorkerCommand
                    {
                        GracePeriod = Duration.FromTimeSpan(TimeSpan.FromSeconds(1)),
                    }),
                cancellation.Token);

        WorkerEnvelope replyEnvelope = await ReadUntilAsync(
            pipePair.GatewayReader,
            WorkerEnvelope.BodyOneofCase.WorkerCommandReply,
            cancellation.Token);

        MxCommandReply reply = replyEnvelope.WorkerCommandReply.Reply;
        Assert.Equal("shutdown-1", reply.CorrelationId);
        Assert.Equal(MxCommandKind.ShutdownWorker, reply.Kind);
        Assert.Equal(ProtocolStatusCode.Ok, reply.ProtocolStatus.Code);

        // The OK reply is followed by a shutdown ack, then the loop stops and
        // the runtime session is disposed.
        WorkerEnvelope ack = await ReadUntilAsync(
            pipePair.GatewayReader,
            WorkerEnvelope.BodyOneofCase.WorkerShutdownAck,
            cancellation.Token);
        Assert.Equal(ProtocolStatusCode.Ok, ack.WorkerShutdownAck.Status.Code);

        Task completedTask = await Task
            .WhenAny(runTask, Task.Delay(TimeSpan.FromSeconds(5), cancellation.Token));
        Assert.Same(runTask, completedTask);
        await runTask;
        Assert.True(runtime.Disposed, "ShutdownWorker must dispose the runtime session.");
    }


    /// <summary>
    ///     Verifies that stale STA activity with no command in flight triggers
    ///     the watchdog StaHung fault. Worker-017 changed the watchdog to skip
    ///     the fault while a command is in flight (the worker is busy
    ///     executing it, not hung), so this test deliberately leaves the
    ///     current-command correlation id empty to assert the genuine-hung
    ///     path still fires.
    /// </summary>
    [Fact]
    public async Task RunAsync_WhenStaActivityIsStale_WritesWatchdogFault()
    {
        using CancellationTokenSource cancellation = new(TimeSpan.FromSeconds(5));
        using PipePair pipePair = await PipePair.CreateAsync(cancellation.Token);
        FakeRuntimeSession runtime = new();
        runtime.SetSnapshot(new WorkerRuntimeHeartbeatSnapshot(
            DateTimeOffset.UtcNow - TimeSpan.FromSeconds(5),
            pendingCommandCount: 0,
            outboundEventQueueDepth: 0,
            lastEventSequence: 0,
            currentCommandCorrelationId: string.Empty));
        WorkerPipeSession session = CreatePipeSession(
            pipePair.WorkerStream,
            runtime,
            new WorkerPipeSessionOptions
            {
                HeartbeatInterval = TimeSpan.FromMilliseconds(20),
                HeartbeatGrace = TimeSpan.FromMilliseconds(50),
            });
        Task runTask = session.RunAsync(cancellation.Token);
        await CompleteGatewayHandshakeAsync(pipePair, cancellation.Token);

        WorkerEnvelope fault = await ReadUntilAsync(
            pipePair.GatewayReader,
            WorkerEnvelope.BodyOneofCase.WorkerFault,
            cancellation.Token);

        Assert.Equal(WorkerFaultCategory.StaHung, fault.WorkerFault.Category);
        Assert.Contains("STA activity is stale", fault.WorkerFault.DiagnosticMessage);

        await SendShutdownAndWaitAsync(pipePair, runTask, cancellation.Token);
    }

    /// <summary>
    ///     Worker-017 regression: while a command is in flight (snapshot's
    ///     current command correlation id is non-empty), stale STA activity
    ///     must NOT trigger the watchdog StaHung fault. The STA is busy
    ///     executing the command, not hung; <c>StaRuntime.ProcessQueuedCommands</c>
    ///     only calls <c>MarkActivity()</c> before and after each work item,
    ///     so a synchronously long-running command (e.g. <c>ReadBulk</c>
    ///     waiting <c>timeout_ms</c> for OnDataChange) legitimately freezes
    ///     <c>LastActivityUtc</c>. The heartbeat already advertises the
    ///     in-flight correlation id so the gateway can apply its own per-command
    ///     timeout.
    /// </summary>
    [Fact]
    public async Task RunAsync_WhenStaActivityIsStaleWithCommandInFlight_DoesNotWriteWatchdogFault()
    {
        using CancellationTokenSource cancellation = new(TimeSpan.FromSeconds(10));
        using PipePair pipePair = await PipePair.CreateAsync(cancellation.Token);
        FakeRuntimeSession runtime = new();
        runtime.SetSnapshot(new WorkerRuntimeHeartbeatSnapshot(
            DateTimeOffset.UtcNow - TimeSpan.FromSeconds(5),
            pendingCommandCount: 0,
            outboundEventQueueDepth: 0,
            lastEventSequence: 0,
            currentCommandCorrelationId: "slow-bulk-read"));
        WorkerPipeSession session = CreatePipeSession(
            pipePair.WorkerStream,
            runtime,
            new WorkerPipeSessionOptions
            {
                HeartbeatInterval = TimeSpan.FromMilliseconds(20),
                HeartbeatGrace = TimeSpan.FromMilliseconds(50),
            });
        Task runTask = session.RunAsync(cancellation.Token);
        await CompleteGatewayHandshakeAsync(pipePair, cancellation.Token);

        // Read several frames over a window much larger than HeartbeatGrace.
        // None must be a WorkerFault; multiple heartbeats must all carry the
        // in-flight correlation id. Reading a bounded count of frames keeps
        // the pipe frame-aligned for the subsequent shutdown handshake.
        const int framesToInspect = 6;
        int heartbeatsObserved = 0;
        for (int index = 0; index < framesToInspect; index++)
        {
            WorkerEnvelope envelope = await pipePair.GatewayReader
                .ReadAsync(cancellation.Token);
            Assert.NotEqual(
                WorkerEnvelope.BodyOneofCase.WorkerFault,
                envelope.BodyCase);
            if (envelope.BodyCase == WorkerEnvelope.BodyOneofCase.WorkerHeartbeat)
            {
                Assert.Equal(
                    "slow-bulk-read",
                    envelope.WorkerHeartbeat.CurrentCommandCorrelationId);
                heartbeatsObserved++;
            }
        }

        Assert.True(
            heartbeatsObserved >= 2,
            $"Expected multiple heartbeats during in-flight command window; observed {heartbeatsObserved}.");

        await SendShutdownAndWaitAsync(pipePair, runTask, cancellation.Token);
    }

    /// <summary>
    ///     Worker-004 regression: once the watchdog reports an StaHung fault,
    ///     subsequent heartbeats must report <see cref="WorkerState.Faulted"/>
    ///     rather than a non-faulted state that contradicts the fault. The
    ///     snapshot uses an empty current-command correlation id so the
    ///     heartbeat State is derived from the session state, not forced to
    ///     ExecutingCommand.
    /// </summary>
    [Fact]
    public async Task RunAsync_AfterWatchdogFault_HeartbeatReportsFaultedState()
    {
        using CancellationTokenSource cancellation = new(TimeSpan.FromSeconds(10));
        using PipePair pipePair = await PipePair.CreateAsync(cancellation.Token);
        FakeRuntimeSession runtime = new();
        runtime.SetSnapshot(new WorkerRuntimeHeartbeatSnapshot(
            DateTimeOffset.UtcNow - TimeSpan.FromSeconds(5),
            pendingCommandCount: 0,
            outboundEventQueueDepth: 0,
            lastEventSequence: 0,
            currentCommandCorrelationId: string.Empty));
        WorkerPipeSession session = CreatePipeSession(
            pipePair.WorkerStream,
            runtime,
            new WorkerPipeSessionOptions
            {
                HeartbeatInterval = TimeSpan.FromMilliseconds(20),
                HeartbeatGrace = TimeSpan.FromMilliseconds(50),
            });
        Task runTask = session.RunAsync(cancellation.Token);
        await CompleteGatewayHandshakeAsync(pipePair, cancellation.Token);

        WorkerEnvelope fault = await ReadUntilAsync(
            pipePair.GatewayReader,
            WorkerEnvelope.BodyOneofCase.WorkerFault,
            cancellation.Token);
        Assert.Equal(WorkerFaultCategory.StaHung, fault.WorkerFault.Category);

        // The next heartbeat after the fault must agree with it.
        WorkerEnvelope heartbeat = await ReadUntilAsync(
            pipePair.GatewayReader,
            WorkerEnvelope.BodyOneofCase.WorkerHeartbeat,
            cancellation.Token);
        Assert.Equal(WorkerState.Faulted, heartbeat.WorkerHeartbeat.State);

        await SendShutdownAndWaitAsync(pipePair, runTask, cancellation.Token);
    }

    /// <summary>
    ///     Worker-023 regression: the in-flight-command suppression on the
    ///     <c>StaHung</c> watchdog (Worker-017) is bounded by
    ///     <c>WorkerPipeSessionOptions.HeartbeatStuckCeiling</c>. A truly
    ///     stuck synchronous STA command (e.g. a dead MXAccess provider) would
    ///     otherwise keep <c>CurrentCommandCorrelationId</c> non-empty forever
    ///     and permanently defeat the watchdog. Once <c>LastStaActivityUtc</c>
    ///     has been stale for longer than <c>HeartbeatStuckCeiling</c> the
    ///     watchdog DOES fire <c>StaHung</c> even with a command in flight.
    /// </summary>
    [Fact]
    public async Task RunAsync_WhenStaActivityIsStaleBeyondCeilingWithCommandInFlight_WritesWatchdogFault()
    {
        using CancellationTokenSource cancellation = new(TimeSpan.FromSeconds(5));
        using PipePair pipePair = await PipePair.CreateAsync(cancellation.Token);
        FakeRuntimeSession runtime = new();
        // Stale by 5s, which exceeds the configured 200 ms ceiling — the
        // watchdog must fire even with a command in flight.
        runtime.SetSnapshot(new WorkerRuntimeHeartbeatSnapshot(
            DateTimeOffset.UtcNow - TimeSpan.FromSeconds(5),
            pendingCommandCount: 0,
            outboundEventQueueDepth: 0,
            lastEventSequence: 0,
            currentCommandCorrelationId: "stuck-command"));
        WorkerPipeSession session = CreatePipeSession(
            pipePair.WorkerStream,
            runtime,
            new WorkerPipeSessionOptions
            {
                HeartbeatInterval = TimeSpan.FromMilliseconds(20),
                HeartbeatGrace = TimeSpan.FromMilliseconds(50),
                HeartbeatStuckCeiling = TimeSpan.FromMilliseconds(200),
            });
        Task runTask = session.RunAsync(cancellation.Token);
        await CompleteGatewayHandshakeAsync(pipePair, cancellation.Token);

        WorkerEnvelope fault = await ReadUntilAsync(
            pipePair.GatewayReader,
            WorkerEnvelope.BodyOneofCase.WorkerFault,
            cancellation.Token);

        Assert.Equal(WorkerFaultCategory.StaHung, fault.WorkerFault.Category);
        Assert.Contains("STA activity is stale", fault.WorkerFault.DiagnosticMessage);

        await SendShutdownAndWaitAsync(pipePair, runTask, cancellation.Token);
    }

    /// <summary>
    ///     Worker-025 regression: <c>RunAsync</c> must throw a diagnostic
    ///     exception if the runtime-session factory returns null, rather than
    ///     deferring the failure to an NRE on the next dereference.
    /// </summary>
    [Fact]
    public async Task RunAsync_WhenRuntimeSessionFactoryReturnsNull_ThrowsDiagnosticException()
    {
        using CancellationTokenSource cancellation = new(TimeSpan.FromSeconds(5));
        using PipePair pipePair = await PipePair.CreateAsync(cancellation.Token);
        WorkerFrameProtocolOptions options = CreateOptions();
        WorkerPipeSession session = new(
            new WorkerFrameReader(pipePair.WorkerStream, options),
            new WorkerFrameWriter(pipePair.WorkerStream, options),
            options,
            () => 1234,
            new WorkerPipeSessionOptions(),
            () => null!);

        InvalidOperationException exception = await Assert.ThrowsAsync<InvalidOperationException>(
            () => session.RunAsync(cancellation.Token));

        Assert.Contains("factory returned null", exception.Message);
    }

    /// <summary>
    ///     Worker-006 regression: when graceful shutdown times out, RunAsync
    ///     must still dispose the runtime session in its finally block.
    ///     Skipping disposal on the timed-out path leaked the STA thread and
    ///     the MXAccess COM object.
    /// </summary>
    [Fact]
    public async Task RunAsync_WhenShutdownTimesOut_StillDisposesRuntimeSession()
    {
        using CancellationTokenSource cancellation = new(TimeSpan.FromSeconds(10));
        using PipePair pipePair = await PipePair.CreateAsync(cancellation.Token);
        FakeRuntimeSession runtime = new()
        {
            ThrowTimeoutOnShutdown = true,
        };
        WorkerPipeSession session = CreatePipeSession(
            pipePair.WorkerStream,
            runtime,
            new WorkerPipeSessionOptions
            {
                HeartbeatInterval = TimeSpan.FromSeconds(1),
                HeartbeatGrace = TimeSpan.FromSeconds(30),
            });
        Task runTask = session.RunAsync(cancellation.Token);
        await CompleteGatewayHandshakeAsync(pipePair, cancellation.Token);

        await pipePair.GatewayWriter
            .WriteAsync(CreateShutdownEnvelope(), cancellation.Token);

        // Drain the gateway-side pipe (heartbeats + the shutdown-timeout
        // fault) so the worker's writes never block on a full pipe buffer.
        Task drainTask = DrainReaderUntilFaultedAsync(pipePair.GatewayReader, cancellation.Token);

        // RunAsync must rethrow the TimeoutException and still reach its
        // finally block, which disposes the runtime session.
        await Assert.ThrowsAsync<TimeoutException>(async () => await runTask);
        Assert.True(
            runtime.Disposed,
            "RunAsync must dispose the runtime session even when shutdown times out.");

        await drainTask;
    }

    private static async Task DrainReaderUntilFaultedAsync(
        WorkerFrameReader reader,
        CancellationToken cancellationToken)
    {
        try
        {
            while (!cancellationToken.IsCancellationRequested)
            {
                WorkerEnvelope envelope = await reader.ReadAsync(cancellationToken).ConfigureAwait(false);
                if (envelope.BodyCase == WorkerEnvelope.BodyOneofCase.WorkerFault
                    && envelope.WorkerFault.Category == WorkerFaultCategory.ShutdownTimeout)
                {
                    return;
                }
            }
        }
        catch (Exception exception) when (
            exception is OperationCanceledException
            || exception is IOException
            || exception is WorkerFrameProtocolException)
        {
            // The worker pipe closed once RunAsync completed — expected.
        }
    }

    /// <summary>Verifies that shutdown drops late replies and sends shutdown ack.</summary>
    [Fact]
    public async Task RunAsync_WhenShutdownArrivesDuringCommand_DropsLateReplyAndWritesShutdownAck()
    {
        using CancellationTokenSource cancellation = new(TimeSpan.FromSeconds(5));
        using PipePair pipePair = await PipePair.CreateAsync(cancellation.Token);
        FakeRuntimeSession runtime = new()
        {
            BlockDispatch = true,
        };
        WorkerPipeSession session = CreatePipeSession(
            pipePair.WorkerStream,
            runtime,
            new WorkerPipeSessionOptions
            {
                HeartbeatInterval = TimeSpan.FromSeconds(1),
                HeartbeatGrace = TimeSpan.FromSeconds(5),
            });
        Task runTask = session.RunAsync(cancellation.Token);
        await CompleteGatewayHandshakeAsync(pipePair, cancellation.Token);

        await pipePair.GatewayWriter.WriteAsync(
            CreateCommandEnvelope("command-during-shutdown"),
            cancellation.Token);
        Assert.True(runtime.DispatchStarted.Wait(TimeSpan.FromSeconds(2)));

        await pipePair.GatewayWriter
            .WriteAsync(CreateShutdownEnvelope(), cancellation.Token);

        WorkerEnvelope shutdownAck = await ReadUntilAsync(
            pipePair.GatewayReader,
            WorkerEnvelope.BodyOneofCase.WorkerShutdownAck,
            cancellation.Token);

        Assert.Equal(ProtocolStatusCode.Ok, shutdownAck.WorkerShutdownAck.Status.Code);
        Task completedTask = await Task.WhenAny(runTask, Task.Delay(TimeSpan.FromSeconds(2), cancellation.Token));
        Assert.Same(runTask, completedTask);
        await runTask;
    }

    /// <summary>Verifies that command exceptions after shutdown are dropped before ack.</summary>
    [Fact]
    public async Task RunAsync_WhenCommandThrowsAfterShutdown_DropsLateFaultAndWritesShutdownAck()
    {
        using CancellationTokenSource cancellation = new(TimeSpan.FromSeconds(5));
        using PipePair pipePair = await PipePair.CreateAsync(cancellation.Token);
        FakeRuntimeSession runtime = new()
        {
            BlockDispatch = true,
            ThrowAfterDispatchReleased = true,
        };
        WorkerPipeSession session = CreatePipeSession(
            pipePair.WorkerStream,
            runtime,
            new WorkerPipeSessionOptions
            {
                HeartbeatInterval = TimeSpan.FromSeconds(1),
                HeartbeatGrace = TimeSpan.FromSeconds(5),
            });
        Task runTask = session.RunAsync(cancellation.Token);
        await CompleteGatewayHandshakeAsync(pipePair, cancellation.Token);

        await pipePair.GatewayWriter.WriteAsync(
            CreateCommandEnvelope("command-fails-during-shutdown"),
            cancellation.Token);
        Assert.True(runtime.DispatchStarted.Wait(TimeSpan.FromSeconds(2)));

        await pipePair.GatewayWriter
            .WriteAsync(CreateShutdownEnvelope(), cancellation.Token);

        // The first heartbeat is emitted immediately on entering the loop
        // (Worker-002), so skip any interleaved heartbeats; the late fault
        // must still be dropped — no WorkerFault may precede the ack.
        WorkerEnvelope envelopeAfterShutdown;
        do
        {
            envelopeAfterShutdown = await pipePair.GatewayReader.ReadAsync(cancellation.Token);
            Assert.NotEqual(
                WorkerEnvelope.BodyOneofCase.WorkerFault,
                envelopeAfterShutdown.BodyCase);
        }
        while (envelopeAfterShutdown.BodyCase == WorkerEnvelope.BodyOneofCase.WorkerHeartbeat);

        Assert.Equal(WorkerEnvelope.BodyOneofCase.WorkerShutdownAck, envelopeAfterShutdown.BodyCase);
        Assert.Equal(ProtocolStatusCode.Ok, envelopeAfterShutdown.WorkerShutdownAck.Status.Code);
        Task completedTask = await Task.WhenAny(runTask, Task.Delay(TimeSpan.FromSeconds(2), cancellation.Token));
        Assert.Same(runTask, completedTask);
        await runTask;
    }

    /// <summary>
    ///     Worker.Tests-017 regression: the <c>WorkerCancel</c> branch of
    ///     <see cref="WorkerPipeSession.DispatchGatewayEnvelopeAsync"/> must
    ///     forward the envelope's correlation id to the runtime session via
    ///     <see cref="IWorkerRuntimeSession.CancelCommand"/> and keep the
    ///     message loop running (no fault, no exit). The handler dispatch
    ///     returns <c>true</c> (keep reading), so a subsequent
    ///     <c>WorkerShutdown</c> still produces the normal shutdown ack.
    /// </summary>
    [Fact]
    public async Task RunAsync_WhenGatewaySendsWorkerCancel_ForwardsCorrelationIdToRuntimeSession()
    {
        using CancellationTokenSource cancellation = new(TimeSpan.FromSeconds(5));
        using PipePair pipePair = await PipePair.CreateAsync(cancellation.Token);
        FakeRuntimeSession runtime = new();
        WorkerPipeSession session = CreatePipeSession(
            pipePair.WorkerStream,
            runtime,
            new WorkerPipeSessionOptions
            {
                HeartbeatInterval = TimeSpan.FromSeconds(1),
                HeartbeatGrace = TimeSpan.FromSeconds(5),
            });
        Task runTask = session.RunAsync(cancellation.Token);
        await CompleteGatewayHandshakeAsync(pipePair, cancellation.Token);

        await pipePair.GatewayWriter
            .WriteAsync(CreateCancelEnvelope("cancel-correlation-1"), cancellation.Token);

        // The session must remain in its message loop: send a follow-up
        // shutdown and observe the normal ack. If WorkerCancel had faulted
        // the pipe or exited the loop, the ack would never arrive.
        await SendShutdownAndWaitAsync(pipePair, runTask, cancellation.Token);

        Assert.Contains("cancel-correlation-1", runtime.CancelledCorrelationIds);
    }

    /// <summary>
    ///     Worker.Tests-017 regression: the <c>default:</c> arm of
    ///     <see cref="WorkerPipeSession.DispatchGatewayEnvelopeAsync"/> must
    ///     throw <see cref="WorkerFrameProtocolException"/> with
    ///     <see cref="WorkerFrameProtocolErrorCode.UnexpectedEnvelopeBody"/>
    ///     when the gateway sends an envelope body that is invalid
    ///     post-handshake (here a second <c>GatewayHello</c>) and must exit
    ///     the message loop — <see cref="WorkerPipeSession.RunAsync"/>
    ///     surfaces the exception to the caller. The message loop does not
    ///     emit a fault frame on this path (the handshake catch in
    ///     <c>CompleteStartupHandshakeAsync</c> is what writes faults for
    ///     pre-handshake protocol violations); the contract this test pins
    ///     is the exception type/error-code and message-loop exit.
    /// </summary>
    [Fact]
    public async Task RunAsync_WhenGatewaySendsUnexpectedEnvelopeBodyAfterHandshake_ThrowsAndExitsMessageLoop()
    {
        using CancellationTokenSource cancellation = new(TimeSpan.FromSeconds(10));
        using PipePair pipePair = await PipePair.CreateAsync(cancellation.Token);
        FakeRuntimeSession runtime = new();
        // Use a long heartbeat interval so no heartbeat frame fires during
        // the test window. With no heartbeats and no fault frame written on
        // the unexpected-body path, the gateway pipe receives nothing after
        // the handshake — no drain task is needed.
        WorkerPipeSession session = CreatePipeSession(
            pipePair.WorkerStream,
            runtime,
            new WorkerPipeSessionOptions
            {
                HeartbeatInterval = TimeSpan.FromSeconds(30),
                HeartbeatGrace = TimeSpan.FromSeconds(60),
            });
        Task runTask = session.RunAsync(cancellation.Token);
        await CompleteGatewayHandshakeAsync(pipePair, cancellation.Token);

        // Send a second GatewayHello — valid envelope, invalid for the
        // post-handshake state, so DispatchGatewayEnvelopeAsync falls to
        // the default arm.
        await pipePair.GatewayWriter
            .WriteAsync(CreateGatewayHelloEnvelope(), cancellation.Token);

        WorkerFrameProtocolException exception =
            await Assert.ThrowsAsync<WorkerFrameProtocolException>(async () => await runTask);
        Assert.Equal(WorkerFrameProtocolErrorCode.UnexpectedEnvelopeBody, exception.ErrorCode);
    }

    /// <summary>
    ///     Worker-002 regression: the first heartbeat must be emitted
    ///     immediately on entering the heartbeat loop, not after a full
    ///     HeartbeatInterval. A long interval is configured so a delay-first
    ///     loop would fail to deliver a heartbeat inside the assertion window.
    /// </summary>
    [Fact]
    public async Task RunAsync_SendsFirstHeartbeatImmediatelyOnEnteringLoop()
    {
        using CancellationTokenSource cancellation = new(TimeSpan.FromSeconds(10));
        using PipePair pipePair = await PipePair.CreateAsync(cancellation.Token);
        FakeRuntimeSession runtime = new();
        WorkerPipeSession session = CreatePipeSession(
            pipePair.WorkerStream,
            runtime,
            new WorkerPipeSessionOptions
            {
                // A deliberately long interval: a delay-before-first-beat
                // loop would not produce a heartbeat for 30s.
                HeartbeatInterval = TimeSpan.FromSeconds(30),
                HeartbeatGrace = TimeSpan.FromSeconds(60),
            });
        Task runTask = session.RunAsync(cancellation.Token);
        await CompleteGatewayHandshakeAsync(pipePair, cancellation.Token);

        DateTimeOffset start = DateTimeOffset.UtcNow;
        using CancellationTokenSource heartbeatWait = CancellationTokenSource
            .CreateLinkedTokenSource(cancellation.Token);
        heartbeatWait.CancelAfter(TimeSpan.FromSeconds(5));
        WorkerEnvelope heartbeat = await ReadUntilAsync(
            pipePair.GatewayReader,
            WorkerEnvelope.BodyOneofCase.WorkerHeartbeat,
            heartbeatWait.Token);
        TimeSpan elapsed = DateTimeOffset.UtcNow - start;

        Assert.Equal(WorkerEnvelope.BodyOneofCase.WorkerHeartbeat, heartbeat.BodyCase);
        Assert.True(
            elapsed < TimeSpan.FromSeconds(5),
            $"First heartbeat took {elapsed}, expected well under the 30s interval.");

        await SendShutdownAndWaitAsync(pipePair, runTask, cancellation.Token);
    }

    /// <summary>
    ///     Worker-003 regression: when a command completes after the worker
    ///     has transitioned out of a command-serving state, the dropped
    ///     reply must be logged with a diagnostic rather than discarded
    ///     silently, so a stuck gateway correlation wait can be traced.
    /// </summary>
    [Fact]
    public async Task RunAsync_WhenReplyIsDroppedAfterShutdown_LogsDiagnostic()
    {
        using CancellationTokenSource cancellation = new(TimeSpan.FromSeconds(10));
        using PipePair pipePair = await PipePair.CreateAsync(cancellation.Token);
        FakeRuntimeSession runtime = new()
        {
            BlockDispatch = true,
        };
        RecordingWorkerLogger logger = new();
        WorkerFrameProtocolOptions options = CreateOptions();
        WorkerPipeSession session = new(
            new WorkerFrameReader(pipePair.WorkerStream, options),
            new WorkerFrameWriter(pipePair.WorkerStream, options),
            options,
            () => 1234,
            new WorkerPipeSessionOptions
            {
                HeartbeatInterval = TimeSpan.FromSeconds(1),
                HeartbeatGrace = TimeSpan.FromSeconds(5),
            },
            () => runtime,
            logger);
        Task runTask = session.RunAsync(cancellation.Token);
        await CompleteGatewayHandshakeAsync(pipePair, cancellation.Token);

        await pipePair.GatewayWriter.WriteAsync(
            CreateCommandEnvelope("command-dropped-after-shutdown"),
            cancellation.Token);
        Assert.True(runtime.DispatchStarted.Wait(TimeSpan.FromSeconds(2)));

        await pipePair.GatewayWriter
            .WriteAsync(CreateShutdownEnvelope(), cancellation.Token);

        WorkerEnvelope shutdownAck = await ReadUntilAsync(
            pipePair.GatewayReader,
            WorkerEnvelope.BodyOneofCase.WorkerShutdownAck,
            cancellation.Token);
        Assert.Equal(ProtocolStatusCode.Ok, shutdownAck.WorkerShutdownAck.Status.Code);

        Task completedTask = await Task.WhenAny(runTask, Task.Delay(TimeSpan.FromSeconds(3), cancellation.Token));
        Assert.Same(runTask, completedTask);
        await runTask;

        Assert.Contains(
            logger.Events,
            entry => entry.EventName == "WorkerCommandResultDropped"
                && entry.Fields.TryGetValue("correlation_id", out object? correlationId)
                && (string?)correlationId == "command-dropped-after-shutdown");
    }

    private static WorkerPipeSession CreateSession(
        Stream inbound,
        Stream outbound,
        WorkerFrameProtocolOptions options)
    {
        return new WorkerPipeSession(
            new WorkerFrameReader(inbound, options),
            new WorkerFrameWriter(outbound, options),
            options,
            () => 1234);
    }

    private static WorkerPipeSession CreatePipeSession(
        Stream stream,
        FakeRuntimeSession runtime)
    {
        return CreatePipeSession(
            stream,
            runtime,
            new WorkerPipeSessionOptions
            {
                HeartbeatInterval = TimeSpan.FromMilliseconds(100),
                HeartbeatGrace = TimeSpan.FromSeconds(5),
            });
    }

    private static WorkerPipeSession CreatePipeSession(
        Stream stream,
        FakeRuntimeSession runtime,
        WorkerPipeSessionOptions sessionOptions)
    {
        WorkerFrameProtocolOptions options = CreateOptions();
        return new WorkerPipeSession(
            new WorkerFrameReader(stream, options),
            new WorkerFrameWriter(stream, options),
            options,
            () => 1234,
            sessionOptions,
            () => runtime);
    }

    private static WorkerFrameProtocolOptions CreateOptions()
    {
        return new WorkerFrameProtocolOptions(
            SessionId,
            GatewayContractInfo.WorkerProtocolVersion,
            Nonce);
    }

    // Inbound-envelope sequence numbers below are documentation-only: the
    // worker has no inbound monotonicity check, so the literal values do
    // not affect dispatch. Each helper exposes a sequence parameter
    // (default = position in the typical Hello/Command/Cancel/Shutdown
    // ordering) so a multi-frame test that interleaves the helpers can
    // assign monotonically increasing values and produce a wire trace
    // that reads in ascending order — see Worker.Tests-030.
    private static WorkerEnvelope CreateGatewayHelloEnvelope(
        string nonce = Nonce,
        uint supportedProtocolVersion = GatewayContractInfo.WorkerProtocolVersion,
        ulong sequence = 1)
    {
        return new WorkerEnvelope
        {
            ProtocolVersion = GatewayContractInfo.WorkerProtocolVersion,
            SessionId = SessionId,
            Sequence = sequence,
            GatewayHello = new GatewayHello
            {
                SupportedProtocolVersion = supportedProtocolVersion,
                Nonce = nonce,
                GatewayVersion = "test-gateway",
            },
        };
    }

    // A generic STA-dispatched command used by the dispatch/heartbeat/
    // shutdown-race tests. Register is a real MXAccess command kind (not a
    // worker control command), so it flows through IWorkerRuntimeSession
    // .DispatchAsync — unlike Ping/GetSessionState/etc., which are answered on
    // the message-loop thread without touching the STA.
    private static WorkerEnvelope CreateCommandEnvelope(string correlationId, ulong sequence = 2)
    {
        return new WorkerEnvelope
        {
            ProtocolVersion = GatewayContractInfo.WorkerProtocolVersion,
            SessionId = SessionId,
            Sequence = sequence,
            CorrelationId = correlationId,
            WorkerCommand = new WorkerCommand
            {
                Command = new MxCommand
                {
                    Kind = MxCommandKind.Register,
                    Register = new RegisterCommand
                    {
                        ClientName = "test-client",
                    },
                },
                EnqueueTimestamp = Timestamp.FromDateTimeOffset(DateTimeOffset.UtcNow),
            },
        };
    }

    private static WorkerEnvelope CreatePingCommandEnvelope(
        string correlationId,
        string message,
        ulong sequence = 2)
    {
        return CreateControlCommandEnvelope(
            correlationId,
            MxCommandKind.Ping,
            command => command.Ping = new PingCommand { Message = message },
            sequence);
    }

    private static WorkerEnvelope CreateControlCommandEnvelope(
        string correlationId,
        MxCommandKind kind,
        Action<MxCommand> configurePayload,
        ulong sequence = 2)
    {
        MxCommand command = new() { Kind = kind };
        configurePayload(command);
        return new WorkerEnvelope
        {
            ProtocolVersion = GatewayContractInfo.WorkerProtocolVersion,
            SessionId = SessionId,
            Sequence = sequence,
            CorrelationId = correlationId,
            WorkerCommand = new WorkerCommand
            {
                Command = command,
                EnqueueTimestamp = Timestamp.FromDateTimeOffset(DateTimeOffset.UtcNow),
            },
        };
    }

    private static WorkerEnvelope CreateCancelEnvelope(string correlationId, ulong sequence = 2)
    {
        return new WorkerEnvelope
        {
            ProtocolVersion = GatewayContractInfo.WorkerProtocolVersion,
            SessionId = SessionId,
            Sequence = sequence,
            CorrelationId = correlationId,
            WorkerCancel = new WorkerCancel
            {
                Reason = "test-cancel",
            },
        };
    }

    private static WorkerEnvelope CreateShutdownEnvelope(ulong sequence = 3)
    {
        return new WorkerEnvelope
        {
            ProtocolVersion = GatewayContractInfo.WorkerProtocolVersion,
            SessionId = SessionId,
            Sequence = sequence,
            WorkerShutdown = new WorkerShutdown
            {
                GracePeriod = Duration.FromTimeSpan(TimeSpan.FromSeconds(1)),
                Reason = "test-complete",
            },
        };
    }

    private static WorkerEvent CreateWorkerEvent(ulong sequence)
    {
        return new WorkerEvent
        {
            Event = new MxEvent
            {
                SessionId = SessionId,
                Family = MxEventFamily.OnDataChange,
                WorkerSequence = sequence,
                OnDataChange = new OnDataChangeEvent(),
            },
        };
    }

    private static async Task CompleteGatewayHandshakeAsync(
        PipePair pipePair,
        CancellationToken cancellationToken)
    {
        await pipePair.GatewayWriter
            .WriteAsync(CreateGatewayHelloEnvelope(), cancellationToken)
            .ConfigureAwait(false);

        WorkerEnvelope hello = await pipePair.GatewayReader.ReadAsync(cancellationToken).ConfigureAwait(false);
        WorkerEnvelope ready = await pipePair.GatewayReader.ReadAsync(cancellationToken).ConfigureAwait(false);

        Assert.Equal(WorkerEnvelope.BodyOneofCase.WorkerHello, hello.BodyCase);
        Assert.Equal(WorkerEnvelope.BodyOneofCase.WorkerReady, ready.BodyCase);
    }

    private static async Task SendShutdownAndWaitAsync(
        PipePair pipePair,
        Task runTask,
        CancellationToken cancellationToken)
    {
        await pipePair.GatewayWriter
            .WriteAsync(CreateShutdownEnvelope(), cancellationToken)
            .ConfigureAwait(false);

        WorkerEnvelope shutdownAck = await ReadUntilAsync(
            pipePair.GatewayReader,
            WorkerEnvelope.BodyOneofCase.WorkerShutdownAck,
            cancellationToken);

        Assert.Equal(ProtocolStatusCode.Ok, shutdownAck.WorkerShutdownAck.Status.Code);
        Task completedTask = await Task
            .WhenAny(runTask, Task.Delay(TimeSpan.FromSeconds(5), cancellationToken))
            .ConfigureAwait(false);

        Assert.Same(runTask, completedTask);
        await runTask.ConfigureAwait(false);
    }

    /// <summary>Reads frames until one matching the expected body type is found.</summary>
    /// <param name="reader">Frame reader.</param>
    /// <param name="expectedBody">Expected body case.</param>
    /// <param name="cancellationToken">Token to cancel the asynchronous operation.</param>
    /// <returns>The matching envelope.</returns>
    private static Task<WorkerEnvelope> ReadUntilAsync(
        WorkerFrameReader reader,
        WorkerEnvelope.BodyOneofCase expectedBody,
        CancellationToken cancellationToken)
    {
        return ReadUntilAsync(
            reader,
            expectedBody,
            _ => true,
            cancellationToken);
    }

    /// <summary>Reads frames until one matches the expected body type and predicate.</summary>
    /// <param name="reader">Frame reader.</param>
    /// <param name="expectedBody">Expected body case.</param>
    /// <param name="predicate">Predicate to match against envelope.</param>
    /// <param name="cancellationToken">Token to cancel the asynchronous operation.</param>
    /// <returns>The matching envelope.</returns>
    private static async Task<WorkerEnvelope> ReadUntilAsync(
        WorkerFrameReader reader,
        WorkerEnvelope.BodyOneofCase expectedBody,
        Func<WorkerEnvelope, bool> predicate,
        CancellationToken cancellationToken)
    {
        while (true)
        {
            WorkerEnvelope envelope = await reader.ReadAsync(cancellationToken).ConfigureAwait(false);
            if (envelope.BodyCase == expectedBody && predicate(envelope))
            {
                return envelope;
            }
        }
    }

    private static WorkerEnvelope[] ReadWrittenFrames(
        MemoryStream stream,
        WorkerFrameProtocolOptions options)
    {
        stream.Position = 0;
        WorkerFrameReader reader = new(stream, options);
        List<WorkerEnvelope> envelopes = new();

        while (stream.Position < stream.Length)
        {
            envelopes.Add(reader.ReadAsync(CancellationToken.None).GetAwaiter().GetResult());
        }

        return envelopes.ToArray();
    }

    private sealed class RecordingWorkerLogger : ZB.MOM.WW.MxGateway.Worker.Bootstrap.IWorkerLogger
    {
        private readonly object gate = new();
        private readonly List<LogEntry> events = new();

        /// <summary>Gets a snapshot of the recorded log entries.</summary>
        public IReadOnlyList<LogEntry> Events
        {
            get
            {
                lock (gate)
                {
                    return new List<LogEntry>(events);
                }
            }
        }

        /// <summary>Records an informational log event.</summary>
        /// <param name="eventName">The event name.</param>
        /// <param name="fields">The event fields.</param>
        public void Information(string eventName, IReadOnlyDictionary<string, object?> fields)
        {
            Record(eventName, fields);
        }

        /// <summary>Records an error log event.</summary>
        /// <param name="eventName">The event name.</param>
        /// <param name="fields">The event fields.</param>
        public void Error(string eventName, IReadOnlyDictionary<string, object?> fields)
        {
            Record(eventName, fields);
        }

        private void Record(string eventName, IReadOnlyDictionary<string, object?> fields)
        {
            Dictionary<string, object?> copy = new();
            foreach (KeyValuePair<string, object?> field in fields)
            {
                copy[field.Key] = field.Value;
            }

            lock (gate)
            {
                events.Add(new LogEntry(eventName, copy));
            }
        }

        /// <summary>A single recorded log entry.</summary>
        public sealed class LogEntry
        {
            /// <summary>Initializes a recorded log entry.</summary>
            /// <param name="eventName">The log event name.</param>
            /// <param name="fields">The log event fields.</param>
            public LogEntry(string eventName, IReadOnlyDictionary<string, object?> fields)
            {
                EventName = eventName;
                Fields = fields;
            }

            /// <summary>Gets the log event name.</summary>
            public string EventName { get; }

            /// <summary>Gets the log event fields.</summary>
            public IReadOnlyDictionary<string, object?> Fields { get; }
        }
    }

    private sealed class PipePair : IDisposable
    {
        private readonly NamedPipeServerStream gatewayStream;

        private PipePair(
            NamedPipeServerStream gatewayStream,
            NamedPipeClientStream workerStream)
        {
            this.gatewayStream = gatewayStream;
            WorkerStream = workerStream;
            WorkerFrameProtocolOptions options = CreateOptions();
            GatewayReader = new WorkerFrameReader(gatewayStream, options);
            GatewayWriter = new WorkerFrameWriter(gatewayStream, options);
        }

        /// <summary>Gets the worker side of the named pipe stream.</summary>
        public Stream WorkerStream { get; }

        /// <summary>Gets the gateway frame reader.</summary>
        public WorkerFrameReader GatewayReader { get; }

        /// <summary>Gets the gateway frame writer.</summary>
        public WorkerFrameWriter GatewayWriter { get; }

        /// <summary>Creates a connected pair of named pipes for testing.</summary>
        /// <param name="cancellationToken">Cancellation token.</param>
        /// <returns>Connected pipe pair.</returns>
        public static async Task<PipePair> CreateAsync(CancellationToken cancellationToken)
        {
            string pipeName = $"mxaccessgw-worker-session-tests-{Guid.NewGuid():N}";
            NamedPipeServerStream gatewayStream = new(
                pipeName,
                PipeDirection.InOut,
                maxNumberOfServerInstances: 1,
                PipeTransmissionMode.Byte,
                PipeOptions.Asynchronous);
            NamedPipeClientStream workerStream = new(
                ".",
                pipeName,
                PipeDirection.InOut,
                PipeOptions.Asynchronous);

            Task waitForConnectionTask = gatewayStream.WaitForConnectionAsync();
            await Task
                .Run(() => workerStream.Connect(5000), cancellationToken)
                .ConfigureAwait(false);
            await waitForConnectionTask.ConfigureAwait(false);

            return new PipePair(gatewayStream, workerStream);
        }

        /// <summary>Disposes pipe resources.</summary>
        public void Dispose()
        {
            WorkerStream.Dispose();
            gatewayStream.Dispose();
        }
    }
}
