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";

    // Byte-heavy drain fixture (WRK-21). 10,000 events at ~1.7 KiB each is ~17 MB of queue — far
    // more than one frame — so DrainEvents must truncate.
    //
    // Two limits below are harness accommodations, not properties of the fix. PipePair runs both
    // ends of a duplex pipe inside one process, with no continuous read pump and with blocking
    // FlushFileBuffers under every frame write, so it tolerates neither multi-megabyte frames nor
    // hundreds of large round trips before both ends wedge waiting on each other. Hence a small
    // negotiated frame maximum, and a smaller queue for the drain-to-empty walk. The byte cap
    // behaves identically at any frame size; exhaustive no-loss over the full 10,000 events is
    // covered without a pipe by MxAccessEventQueueTests.
    private const int ByteHeavyEventCount = 10_000;
    private const int RepeatedDrainEventCount = 1_000;
    private const int ByteHeavyEventPayloadBytes = 1_800;
    private const uint NegotiatedMaxFrameBytes = 128 * 1024;

    /// <summary>Verifies that valid gateway hello triggers worker hello and ready responses.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [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>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [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>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [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>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [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>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [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>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [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>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [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>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [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>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [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>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [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>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [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>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [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 a DrainEvents control command with <c>max_events = 0</c> is bounded by the
    ///     worker rather than draining the entire queue into one reply frame: the session passes a
    ///     capped, non-zero maximum to the runtime session.
    /// </summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task RunAsync_DrainEventsWithZeroMaxEvents_BoundsTheDrain()
    {
        using CancellationTokenSource cancellation = new(TimeSpan.FromSeconds(5));
        using PipePair pipePair = await PipePair.CreateAsync(cancellation.Token);
        FakeRuntimeSession runtime = new() { SuppressDrainForBatchSize = 128 };
        WorkerPipeSession session = CreatePipeSession(pipePair.WorkerStream, runtime);
        runtime.EnqueueEvent(CreateWorkerEvent(sequence: 11));
        Task runTask = session.RunAsync(cancellation.Token);
        await CompleteGatewayHandshakeAsync(pipePair, cancellation.Token);

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

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

        // The client asked for "all" (0) but the worker must pass a bounded, non-zero cap so the reply
        // frame cannot grow without limit.
        Assert.NotNull(runtime.LastDrainMaxEvents);
        Assert.NotEqual(0u, runtime.LastDrainMaxEvents!.Value);

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

    /// <summary>
    ///     The WRK-21 repro. A queue full of byte-heavy events (large string values — the payload
    ///     profile this gateway exists for) used to make <c>DrainEvents max_events = 0</c> build a
    ///     reply above the negotiated frame maximum: the writer rejected the frame, the exception
    ///     unwound the session, and the already-dequeued events were gone. The drain is now
    ///     byte-budgeted, so the reply fits, the truncation is reported in the reply's diagnostic
    ///     message, and the session keeps serving.
    /// </summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task DrainEvents_ByteHeavyQueue_ReplyIsBoundedAndSessionSurvives()
    {
        using CancellationTokenSource cancellation = new(TimeSpan.FromSeconds(60));
        using PipePair pipePair = await PipePair.CreateAsync(cancellation.Token);
        FakeRuntimeSession runtime = new()
        {
            SuppressDrainForBatchSize = 128,
            BackingQueue = CreateByteHeavyQueue(ByteHeavyEventCount, ByteHeavyEventPayloadBytes),
        };
        WorkerPipeSession session = CreatePipeSession(pipePair.WorkerStream, runtime);
        Task runTask = session.RunAsync(cancellation.Token);
        await CompleteGatewayHandshakeAsync(pipePair, NegotiatedMaxFrameBytes, cancellation.Token);

        await pipePair.GatewayWriter
            .WriteAsync(
                CreateControlCommandEnvelope(
                    "drain-heavy-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(ProtocolStatusCode.Ok, reply.ProtocolStatus.Code);

        // The whole queue is far larger than one frame, so the reply is a strict subset that fits.
        Assert.True(
            replyEnvelope.CalculateSize() <= NegotiatedMaxFrameBytes,
            $"DrainEvents reply serialized to {replyEnvelope.CalculateSize()} bytes, above the negotiated {NegotiatedMaxFrameBytes}.");
        Assert.InRange(reply.DrainEvents.Events.Count, 1, ByteHeavyEventCount - 1);
        Assert.Contains("remain", reply.DiagnosticMessage);
        Assert.Contains("repeat DrainEvents", reply.DiagnosticMessage);

        // The session is alive: it still answers a ping, and RunAsync has not unwound.
        await pipePair.GatewayWriter
            .WriteAsync(CreatePingCommandEnvelope("ping-after-drain", "still-here"), cancellation.Token);
        WorkerEnvelope pingReply = await ReadUntilAsync(
            pipePair.GatewayReader,
            WorkerEnvelope.BodyOneofCase.WorkerCommandReply,
            envelope => envelope.WorkerCommandReply.Reply.CorrelationId == "ping-after-drain",
            cancellation.Token);
        Assert.Equal("still-here", pingReply.WorkerCommandReply.Reply.DiagnosticMessage);
        Assert.False(runTask.IsCompleted, "The session must survive a byte-heavy DrainEvents.");

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

    /// <summary>
    ///     Verifies the byte-budgeted drain loses nothing: repeating DrainEvents until it comes back
    ///     empty recovers every enqueued event exactly once, in order, across the split replies. The
    ///     pre-fix drain removed events from the queue before sizing the reply, so a rejected frame
    ///     destroyed them — no-loss is the half of the P0 criterion a catch-only fix cannot deliver.
    /// </summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task DrainEvents_RepeatedCalls_RecoverAllEventsWithoutLoss()
    {
        using CancellationTokenSource cancellation = new(TimeSpan.FromSeconds(90));
        using PipePair pipePair = await PipePair.CreateAsync(cancellation.Token);
        FakeRuntimeSession runtime = new()
        {
            SuppressDrainForBatchSize = 128,
            BackingQueue = CreateByteHeavyQueue(RepeatedDrainEventCount, ByteHeavyEventPayloadBytes),
        };
        WorkerPipeSession session = CreatePipeSession(pipePair.WorkerStream, runtime);
        Task runTask = session.RunAsync(cancellation.Token);
        await CompleteGatewayHandshakeAsync(pipePair, NegotiatedMaxFrameBytes, cancellation.Token);

        List<ulong> recovered = new();
        int replyCount = 0;
        while (true)
        {
            string correlationId = $"drain-loop-{replyCount}";
            await pipePair.GatewayWriter
                .WriteAsync(
                    CreateControlCommandEnvelope(
                        correlationId,
                        MxCommandKind.DrainEvents,
                        command => command.DrainEvents = new DrainEventsCommand { MaxEvents = 0 }),
                    cancellation.Token);

            WorkerEnvelope replyEnvelope = await ReadUntilAsync(
                pipePair.GatewayReader,
                WorkerEnvelope.BodyOneofCase.WorkerCommandReply,
                envelope => envelope.WorkerCommandReply.Reply.CorrelationId == correlationId,
                cancellation.Token);
            replyCount++;

            MxCommandReply reply = replyEnvelope.WorkerCommandReply.Reply;
            Assert.True(
                replyEnvelope.CalculateSize() <= NegotiatedMaxFrameBytes,
                $"DrainEvents reply {replyCount} serialized to {replyEnvelope.CalculateSize()} bytes.");
            if (reply.DrainEvents.Events.Count == 0)
            {
                break;
            }

            foreach (MxEvent drained in reply.DrainEvents.Events)
            {
                recovered.Add(drained.WorkerSequence);
            }

            Assert.True(replyCount < 200, "DrainEvents made no progress across 200 replies.");
        }

        // More than one reply proves the drain really split; every event came back exactly once, in
        // enqueue order.
        Assert.True(replyCount > 2, $"Expected the byte cap to split the drain, saw {replyCount} replies.");
        Assert.Equal(RepeatedDrainEventCount, recovered.Count);
        for (int index = 0; index < recovered.Count; index++)
        {
            Assert.Equal((ulong)(index + 1), recovered[index]);
        }

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

    /// <summary>
    ///     Regression for the reserve-boundary budget bug. The gateway validator accepts a Worker
    ///     frame maximum as low as 1024 + 64 KiB, and just above that boundary a naive
    ///     subtract-then-guard budget collapses to ~1024 bytes — too small to move even one
    ///     byte-heavy event, so every drain reports truncation with the same head blocked and the
    ///     drain-until-empty loop never terminates. The budget is now a floor (never below half the
    ///     negotiated maximum), so a byte-heavy queue drains to empty even at the validator floor.
    /// </summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task DrainEvents_AtValidatorFloorFrameMax_MakesProgressAndTerminates()
    {
        // The lowest Worker.MaxMessageBytes GatewayOptionsValidator permits: the public gRPC floor
        // (1024) plus the 64 KiB envelope-overhead reserve. The naive budget would be exactly 1024
        // here; the floored budget is half of the frame max (~33 KiB).
        const uint validatorFloorFrameMax = 1024 + (64 * 1024);
        using CancellationTokenSource cancellation = new(TimeSpan.FromSeconds(60));
        using PipePair pipePair = await PipePair.CreateAsync(cancellation.Token);
        FakeRuntimeSession runtime = new()
        {
            SuppressDrainForBatchSize = 128,
            BackingQueue = CreateByteHeavyQueue(200, ByteHeavyEventPayloadBytes),
        };
        WorkerPipeSession session = CreatePipeSession(pipePair.WorkerStream, runtime);
        Task runTask = session.RunAsync(cancellation.Token);
        await CompleteGatewayHandshakeAsync(pipePair, validatorFloorFrameMax, cancellation.Token);

        int recovered = 0;
        int replyCount = 0;
        while (true)
        {
            string correlationId = $"floor-drain-{replyCount}";
            await pipePair.GatewayWriter
                .WriteAsync(
                    CreateControlCommandEnvelope(
                        correlationId,
                        MxCommandKind.DrainEvents,
                        command => command.DrainEvents = new DrainEventsCommand { MaxEvents = 0 }),
                    cancellation.Token);

            WorkerEnvelope replyEnvelope = await ReadUntilAsync(
                pipePair.GatewayReader,
                WorkerEnvelope.BodyOneofCase.WorkerCommandReply,
                envelope => envelope.WorkerCommandReply.Reply.CorrelationId == correlationId,
                cancellation.Token);
            replyCount++;

            MxCommandReply reply = replyEnvelope.WorkerCommandReply.Reply;
            Assert.True(
                replyEnvelope.CalculateSize() <= validatorFloorFrameMax,
                $"reply {replyCount} serialized to {replyEnvelope.CalculateSize()} bytes.");

            int drainedThisReply = reply.DrainEvents.Events.Count;
            if (drainedThisReply == 0)
            {
                break;
            }

            // The head is never reported as oversized at this frame max: the ~33 KiB floored budget
            // comfortably fits the ~1.8 KiB events, so each reply makes real progress.
            Assert.DoesNotContain("alone exceeds", reply.DiagnosticMessage);
            recovered += drainedThisReply;
            Assert.True(replyCount < 200, "DrainEvents made no progress at the validator floor frame max.");
        }

        Assert.Equal(200, recovered);

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

    /// <summary>
    ///     Verifies the control-reply write seam is not session-fatal on size (WRK-21 backstop). The
    ///     reply builders size their payloads, so this path needs a deliberately budget-blind drain
    ///     to reach — but that is the point: a future command or a sizing bug must degrade to an
    ///     error reply for that correlation, never to a dead session.
    /// </summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task ControlReplyTooLarge_WritesErrorReplyInsteadOfDying()
    {
        const uint tinyMaxFrameBytes = 4096;
        using CancellationTokenSource cancellation = new(TimeSpan.FromSeconds(15));
        using PipePair pipePair = await PipePair.CreateAsync(cancellation.Token);
        FakeRuntimeSession runtime = new()
        {
            SuppressDrainForBatchSize = 128,
            BackingQueue = CreateByteHeavyQueue(eventCount: 1, payloadBytes: 16 * 1024),
            IgnoreDrainByteBudget = true,
        };
        WorkerPipeSession session = CreatePipeSession(pipePair.WorkerStream, runtime);
        Task runTask = session.RunAsync(cancellation.Token);
        await CompleteGatewayHandshakeAsync(pipePair, tinyMaxFrameBytes, cancellation.Token);

        await pipePair.GatewayWriter
            .WriteAsync(
                CreateControlCommandEnvelope(
                    "drain-too-large",
                    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("drain-too-large", reply.CorrelationId);
        Assert.Equal(MxCommandKind.DrainEvents, reply.Kind);
        Assert.Equal(ProtocolStatusCode.InvalidRequest, reply.ProtocolStatus.Code);
        Assert.Contains("frame maximum", reply.ProtocolStatus.Message);
        Assert.False(runTask.IsCompleted, "An oversized control reply must not end the session.");

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

    /// <summary>
    ///     Verifies the same backstop on the STA command path: an oversized reply from a dispatched
    ///     command answers its correlation with an error reply instead of falling into the generic
    ///     catch that faults the whole session with MxaccessCommandFailed.
    /// </summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task CommandReplyTooLarge_WritesErrorReplyInsteadOfFaulting()
    {
        const uint tinyMaxFrameBytes = 4096;
        using CancellationTokenSource cancellation = new(TimeSpan.FromSeconds(15));
        using PipePair pipePair = await PipePair.CreateAsync(cancellation.Token);
        FakeRuntimeSession runtime = new()
        {
            DispatchReplyDiagnosticMessage = new string('x', 16 * 1024),
        };
        WorkerPipeSession session = CreatePipeSession(pipePair.WorkerStream, runtime);
        Task runTask = session.RunAsync(cancellation.Token);
        await CompleteGatewayHandshakeAsync(pipePair, tinyMaxFrameBytes, cancellation.Token);

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

        WorkerEnvelope replyEnvelope = await ReadUntilAsync(
            pipePair.GatewayReader,
            WorkerEnvelope.BodyOneofCase.WorkerCommandReply,
            envelope => envelope.WorkerCommandReply.Reply.CorrelationId == "command-too-large",
            cancellation.Token);

        MxCommandReply reply = replyEnvelope.WorkerCommandReply.Reply;
        Assert.Equal(MxCommandKind.Register, reply.Kind);
        Assert.Equal(ProtocolStatusCode.InvalidRequest, reply.ProtocolStatus.Code);

        // No fault, and the session still reports itself Ready rather than Faulted.
        await pipePair.GatewayWriter
            .WriteAsync(
                CreateControlCommandEnvelope(
                    "state-after-too-large",
                    MxCommandKind.GetSessionState,
                    command => command.GetSessionState = new GetSessionStateCommand()),
                cancellation.Token);
        WorkerEnvelope stateEnvelope = await ReadUntilAsync(
            pipePair.GatewayReader,
            WorkerEnvelope.BodyOneofCase.WorkerCommandReply,
            envelope => envelope.WorkerCommandReply.Reply.CorrelationId == "state-after-too-large",
            cancellation.Token);
        Assert.Equal(SessionState.Ready, stateEnvelope.WorkerCommandReply.Reply.SessionState.State);
        Assert.False(runTask.IsCompleted, "An oversized command reply must not fault the session.");

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

    /// <summary>
    ///     IPC-30. An event above the negotiated frame maximum is undeliverable end to end (the pipe
    ///     maximum sits only an envelope reserve above the public gRPC cap), so the session stays
    ///     fatal by design — but the death must be structured: a WorkerFault naming the event, with
    ///     no value payload in it, before the process exits. Silently dropping the event or
    ///     synthesizing a placeholder were both rejected.
    /// </summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task RunAsync_EventFrameTooLarge_WritesStructuredFaultThenEndsSession()
    {
        const uint tinyMaxFrameBytes = 4096;
        using CancellationTokenSource cancellation = new(TimeSpan.FromSeconds(15));
        using PipePair pipePair = await PipePair.CreateAsync(cancellation.Token);
        FakeRuntimeSession runtime = new();
        RecordingWorkerLogger logger = new();
        WorkerPipeSession session = CreatePipeSession(
            pipePair.WorkerStream,
            runtime,
            new WorkerPipeSessionOptions
            {
                HeartbeatInterval = TimeSpan.FromMilliseconds(100),
                HeartbeatGrace = TimeSpan.FromSeconds(5),
            },
            logger);
        runtime.EnqueueEvent(CreateOversizedWorkerEvent(sequence: 77, payloadBytes: 16 * 1024));
        Task runTask = session.RunAsync(cancellation.Token);
        await CompleteGatewayHandshakeAsync(pipePair, tinyMaxFrameBytes, cancellation.Token);

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

        WorkerFault fault = faultEnvelope.WorkerFault;
        Assert.Equal(WorkerFaultCategory.ProtocolViolation, fault.Category);
        Assert.Equal("EventDrain", fault.CommandMethod);
        Assert.Contains("77", fault.DiagnosticMessage);
        Assert.Contains("MaxMessageBytes", fault.DiagnosticMessage);
        // The identity is reported; the value payload never is.
        Assert.DoesNotContain(new string('x', 64), fault.DiagnosticMessage);

        Assert.Contains(
            logger.Events,
            entry => entry.EventName == "WorkerEventFrameTooLarge"
                && entry.Fields.TryGetValue("worker_sequence", out object? sequence)
                && sequence is ulong sequenceValue
                && sequenceValue == 77UL);

        // The session ends, and the fault frame parsed cleanly off the same stream above — the
        // rejected event never corrupted the wire.
        Task completedTask = await Task.WhenAny(runTask, Task.Delay(TimeSpan.FromSeconds(5), cancellation.Token));
        Assert.Same(runTask, completedTask);
        await Assert.ThrowsAsync<InvalidOperationException>(async () => await runTask);
    }

    /// <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>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [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. The watchdog skips 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>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [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>
    ///     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>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [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>
    ///     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>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [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>
    ///     The in-flight-command suppression on the <c>StaHung</c> watchdog 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>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [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>
    ///     WRK-27. An STA call outside the command dispatcher (the alarm poll) advertises itself on
    ///     the heartbeat snapshot's <c>StaCallInProgress</c> flag, and the watchdog grants it the same
    ///     grace-to-ceiling suppression as a dispatched command: stale STA activity within the ceiling
    ///     does not fault while the flag is set, but stale activity beyond the ceiling faults anyway.
    ///     This closes the 15 s-vs-75 s asymmetry between polls and commands.
    /// </summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task Watchdog_StaCallInProgress_SuppressedUntilCeiling()
    {
        using CancellationTokenSource cancellation = new(TimeSpan.FromSeconds(10));

        // Phase 1 — within the ceiling: stale beyond grace, empty correlation id, StaCallInProgress
        // set. The default 75 s ceiling is far beyond the 5 s staleness, so the watchdog must suppress.
        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,
                staCallInProgress: true));
            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);

            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)
                {
                    heartbeatsObserved++;
                }
            }

            Assert.True(
                heartbeatsObserved >= 2,
                $"Expected multiple heartbeats during the in-progress STA-call window; observed {heartbeatsObserved}.");

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

        // Phase 2 — beyond the ceiling: same StaCallInProgress flag, but staleness (5 s) exceeds the
        // 200 ms ceiling, so the watchdog must fire even with the poll in progress.
        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,
                staCallInProgress: true));
            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);

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

    /// <summary>
    ///     WRK-25. The event drain loop submits a whole drained batch through the writer's batch entry
    ///     point, so a burst of 128 events costs one flush, not 128 — the assertion the WRK-12
    ///     tracking claim needed to actually hold on the event hot path.
    /// </summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task EventBurst_DrainLoopCoalescesFlushes()
    {
        using CancellationTokenSource cancellation = new(TimeSpan.FromSeconds(15));
        using PipePair pipePair = await PipePair.CreateAsync(cancellation.Token);
        FakeRuntimeSession runtime = new();
        // A far-off heartbeat interval keeps every heartbeat after the first out of the measurement
        // window; the first beat is sent immediately on entering the message loop (see
        // RunHeartbeatLoopAsync) and is closed out explicitly below.
        FlushCountingPassthroughStream countingStream = new(pipePair.WorkerStream);
        WorkerPipeSession session = CreatePipeSession(
            countingStream,
            runtime,
            new WorkerPipeSessionOptions
            {
                HeartbeatInterval = TimeSpan.FromMinutes(5),
                HeartbeatGrace = TimeSpan.FromSeconds(30),
            });
        Task runTask = session.RunAsync(cancellation.Token);
        await CompleteGatewayHandshakeAsync(pipePair, cancellation.Token);

        // Close the pre-burst window on an explicit signal rather than a sleep. Three frames precede
        // the burst — WorkerHello, WorkerReady, and the immediate first heartbeat — and each is
        // flushed only after its bytes are already on the pipe, so having read a frame is no evidence
        // that its flush has been counted. Read the first beat (the last pre-burst frame), then wait
        // for the writer's deferred flush before sampling the baseline. A fixed delay left the first
        // beat's flush free to land after the sample under CI scheduling pressure, where it was
        // charged to the burst and the assertion below saw two flushes instead of one.
        await ReadUntilAsync(
            pipePair.GatewayReader,
            WorkerEnvelope.BodyOneofCase.WorkerHeartbeat,
            cancellation.Token);
        await countingStream.WaitForAllWritesFlushedAsync(cancellation.Token);
        int baselineFlushes = countingStream.FlushCount;

        // Enqueue a full 128-event batch atomically so the drain loop sees it as one batch.
        const int burst = 128;
        List<WorkerEvent> batch = new(burst);
        for (int index = 0; index < burst; index++)
        {
            batch.Add(CreateWorkerEvent(sequence: (ulong)(index + 1)));
        }

        runtime.EnqueueEvents(batch);

        // Drain all 128 events off the gateway side.
        for (int index = 0; index < burst; index++)
        {
            await ReadUntilAsync(
                pipePair.GatewayReader,
                WorkerEnvelope.BodyOneofCase.WorkerEvent,
                cancellation.Token);
        }

        // Take the same edge on the burst's own flush — the 128 frames are on the wire before the
        // writer flushes them — then assert on the recorded flush shape: exactly one flush beyond the
        // baseline, and that flush carried every event frame of the burst. Asserting the shape (not
        // just the count) is what makes the coalescing claim faithful: a split batch would show its
        // first flush carrying fewer than the whole burst.
        await countingStream.WaitForAllWritesFlushedAsync(cancellation.Token);
        IReadOnlyList<int> flushWriteCounts = countingStream.SnapshotFlushWriteCounts();
        Assert.Equal(baselineFlushes + 1, flushWriteCounts.Count);
        Assert.Equal(burst, flushWriteCounts[baselineFlushes]);

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

    /// <summary>
    ///     WRK-24. A GatewayHello negotiating a frame maximum below the worker floor faults at the
    ///     handshake with a fault frame rather than being adopted — mirroring the above-ceiling
    ///     handshake behavior — so a nonsensical tiny value never leaves a session that fails every
    ///     later frame. No message loop is entered.
    /// </summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task Handshake_GatewayHelloWithTinyMaxFrameBytes_FaultsAtHandshake()
    {
        WorkerFrameProtocolOptions options = CreateOptions();
        using MemoryStream inbound = new();
        await new WorkerFrameWriter(inbound, options)
            .WriteAsync(CreateGatewayHelloEnvelope(maxFrameBytes: 512));
        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.InvalidConfiguration, exception.ErrorCode);
        WorkerEnvelope fault = Assert.Single(ReadWrittenFrames(outbound, options));
        Assert.Equal(WorkerEnvelope.BodyOneofCase.WorkerFault, fault.BodyCase);
    }

    /// <summary>
    ///     Regression test: a long in-flight STA command that keeps pumping
    ///     must NOT self-fault as <c>StaHung</c>, and its reply must still be
    ///     delivered. The real fix makes <c>StaRuntime.PumpPendingMessages</c>
    ///     refresh <c>LastActivityUtc</c> on every wait iteration, so a healthy
    ///     <c>ReadBulk</c> holding the STA far longer than
    ///     <c>HeartbeatStuckCeiling</c> (75 s in production) keeps its activity
    ///     timestamp fresh. This test compresses the clock — a 100 ms ceiling
    ///     with a command in flight across a window many multiples longer — and
    ///     models the pump refresh by continuously advancing the snapshot's
    ///     <c>LastStaActivityUtc</c> while the command blocks. Contrast
    ///     <see cref="RunAsync_WhenStaActivityIsStaleBeyondCeilingWithCommandInFlight_WritesWatchdogFault"/>,
    ///     where a frozen timestamp beyond the ceiling correctly faults; here
    ///     the refreshed timestamp must keep the fault suppressed and let the
    ///     reply through the <c>Ready</c>-state gate.
    /// </summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task RunAsync_LongInFlightCommandThatKeepsPumping_DoesNotFaultAndDeliversReply()
    {
        using CancellationTokenSource cancellation = new(TimeSpan.FromSeconds(20));
        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.FromMilliseconds(50),
                HeartbeatStuckCeiling = TimeSpan.FromMilliseconds(100),
            });
        Task runTask = session.RunAsync(cancellation.Token);
        await CompleteGatewayHandshakeAsync(pipePair, cancellation.Token);

        // Kick off the long command; it blocks in DispatchAsync until released,
        // so its correlation id stays in flight in the heartbeat snapshot.
        await pipePair.GatewayWriter
            .WriteAsync(CreateCommandEnvelope("long-bulk-read"), cancellation.Token);
        Assert.True(
            runtime.DispatchStarted.Wait(TimeSpan.FromSeconds(5)),
            "The long command must reach the runtime and begin dispatch.");

        // Model the pump refreshing STA activity on each wait iteration: keep
        // the snapshot's LastStaActivityUtc current while the command is in
        // flight.
        using CancellationTokenSource pumpRefresh = new();
        Task refreshLoop = Task.Run(
            async () =>
            {
                while (!pumpRefresh.IsCancellationRequested)
                {
                    runtime.SetSnapshot(new WorkerRuntimeHeartbeatSnapshot(
                        DateTimeOffset.UtcNow,
                        pendingCommandCount: 1,
                        outboundEventQueueDepth: 0,
                        lastEventSequence: 0,
                        currentCommandCorrelationId: "long-bulk-read"));
                    await Task.Delay(TimeSpan.FromMilliseconds(20)).ConfigureAwait(false);
                }
            });

        // Inspect a bounded number of frames over a window many multiples of the
        // 100 ms ceiling (at least 30 heartbeats at 20 ms ~ 600 ms). None may be
        // a WorkerFault while activity is continuously refreshed.
        const int framesToInspect = 30;
        for (int index = 0; index < framesToInspect; index++)
        {
            WorkerEnvelope envelope = await pipePair.GatewayReader
                .ReadAsync(cancellation.Token);
            Assert.NotEqual(
                WorkerEnvelope.BodyOneofCase.WorkerFault,
                envelope.BodyCase);
        }

        // Stop refreshing and release the command; its reply must be delivered
        // because the session never faulted (state stayed Ready).
        pumpRefresh.Cancel();
        await refreshLoop;
        runtime.ReleaseDispatch();

        WorkerEnvelope reply = await ReadUntilAsync(
            pipePair.GatewayReader,
            WorkerEnvelope.BodyOneofCase.WorkerCommandReply,
            envelope => envelope.CorrelationId == "long-bulk-read",
            cancellation.Token);
        Assert.Equal(
            ProtocolStatusCode.Ok,
            reply.WorkerCommandReply.Reply.ProtocolStatus.Code);

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

    /// <summary>
    ///     <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>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [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>
    ///     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>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [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>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [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>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [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,
        // 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>
    ///     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>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [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>
    ///     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>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [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>
    ///     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>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [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);

        // The heartbeatWait CTS (5s cancel-after) already enforces the timing bound:
        // if the first heartbeat is not received within 5s, ReadUntilAsync throws
        // OperationCanceledException and the test fails.  A redundant wall-clock
        // elapsed < 5s assertion would add the same class of flakiness
        // corrected elsewhere, so it is omitted here.
        using CancellationTokenSource heartbeatWait = CancellationTokenSource
            .CreateLinkedTokenSource(cancellation.Token);
        heartbeatWait.CancelAfter(TimeSpan.FromSeconds(5));
        WorkerEnvelope heartbeat = await ReadUntilAsync(
            pipePair.GatewayReader,
            WorkerEnvelope.BodyOneofCase.WorkerHeartbeat,
            heartbeatWait.Token);

        Assert.Equal(WorkerEnvelope.BodyOneofCase.WorkerHeartbeat, heartbeat.BodyCase);

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

    /// <summary>
    ///     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>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [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)
    {
        return CreatePipeSession(stream, runtime, sessionOptions, logger: null);
    }

    private static WorkerPipeSession CreatePipeSession(
        Stream stream,
        FakeRuntimeSession runtime,
        WorkerPipeSessionOptions sessionOptions,
        ZB.MOM.WW.MxGateway.Worker.Bootstrap.IWorkerLogger? logger)
    {
        WorkerFrameProtocolOptions options = CreateOptions();
        return new WorkerPipeSession(
            new WorkerFrameReader(stream, options),
            new WorkerFrameWriter(stream, options),
            options,
            () => 1234,
            sessionOptions,
            () => runtime,
            logger);
    }

    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.
    private static WorkerEnvelope CreateGatewayHelloEnvelope(
        string nonce = Nonce,
        uint supportedProtocolVersion = GatewayContractInfo.WorkerProtocolVersion,
        ulong sequence = 1,
        uint maxFrameBytes = 0)
    {
        return new WorkerEnvelope
        {
            ProtocolVersion = GatewayContractInfo.WorkerProtocolVersion,
            SessionId = SessionId,
            Sequence = sequence,
            GatewayHello = new GatewayHello
            {
                SupportedProtocolVersion = supportedProtocolVersion,
                Nonce = nonce,
                GatewayVersion = "test-gateway",
                // 0 leaves the worker on its compile-time default; a non-zero value is adopted
                // during the handshake and becomes the frame maximum every later assertion is
                // measured against.
                MaxFrameBytes = maxFrameBytes,
            },
        };
    }

    // 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 WorkerEvent CreateOversizedWorkerEvent(ulong sequence, int payloadBytes)
    {
        WorkerEvent workerEvent = CreateWorkerEvent(sequence);
        workerEvent.Event.ItemHandle = 42;
        workerEvent.Event.RawStatus = new string('x', payloadBytes);
        return workerEvent;
    }

    /// <summary>
    /// Fills a real event queue with byte-heavy events — a large string field stands in for the
    /// array/string <c>MxValue</c> payloads that make a count-capped drain overshoot the frame
    /// maximum. The queue is real (not the fake's plain list) so the production byte-budgeting runs.
    /// </summary>
    /// <param name="eventCount">Number of events to enqueue.</param>
    /// <param name="payloadBytes">Size of each event's raw-status payload string.</param>
    /// <returns>The populated queue.</returns>
    private static MxAccessEventQueue CreateByteHeavyQueue(int eventCount, int payloadBytes)
    {
        MxAccessEventQueue queue = new(Math.Max(eventCount, 1));
        string payload = new string('x', payloadBytes);
        for (int index = 0; index < eventCount; index++)
        {
            queue.Enqueue(new MxEvent
            {
                SessionId = SessionId,
                Family = MxEventFamily.OnDataChange,
                ItemHandle = index,
                RawStatus = payload,
                OnDataChange = new OnDataChangeEvent(),
            });
        }

        return queue;
    }

    private static Task CompleteGatewayHandshakeAsync(
        PipePair pipePair,
        CancellationToken cancellationToken)
    {
        return CompleteGatewayHandshakeAsync(pipePair, maxFrameBytes: 0, cancellationToken);
    }

    private static async Task CompleteGatewayHandshakeAsync(
        PipePair pipePair,
        uint maxFrameBytes,
        CancellationToken cancellationToken)
    {
        await pipePair.GatewayWriter
            .WriteAsync(CreateGatewayHelloEnvelope(maxFrameBytes: maxFrameBytes), 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);
                }
            }
        }

        /// <inheritdoc />
        public void Information(string eventName, IReadOnlyDictionary<string, object?> fields)
        {
            Record(eventName, fields);
        }

        /// <inheritdoc />
        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; }
        }
    }

    // Wraps the worker side of the pipe and records the flush shape of the frames the writer emits —
    // how many stream writes each flush coalesced — so a test can assert the event drain loop turns a
    // burst into a single flush. Delegates every other operation to the inner stream; does not own the
    // inner stream's lifetime (PipePair disposes it).
    //
    // The writer flushes only after writing a whole drained batch, so the peer can read every frame of
    // that batch before the flush runs. Neither "I read the frames" nor any fixed sleep is evidence
    // that a flush has been counted; WaitForAllWritesFlushedAsync is the explicit edge a test must
    // take before sampling the counters.
    private sealed class FlushCountingPassthroughStream : Stream
    {
        private readonly Stream inner;
        private readonly object gate = new();
        private readonly List<int> flushWriteCounts = new();
        private readonly List<FlushWaiter> waiters = new();
        private int writesSinceLastFlush;

        /// <summary>Initializes the passthrough over the given inner stream.</summary>
        /// <param name="inner">The stream to delegate to.</param>
        public FlushCountingPassthroughStream(Stream inner)
        {
            this.inner = inner;
        }

        /// <summary>Gets the number of flushes observed so far.</summary>
        public int FlushCount
        {
            get
            {
                lock (gate)
                {
                    return flushWriteCounts.Count;
                }
            }
        }

        /// <summary>
        ///     Returns the number of stream writes coalesced into each flush, in flush order.
        /// </summary>
        /// <returns>A snapshot of the per-flush write counts.</returns>
        public IReadOnlyList<int> SnapshotFlushWriteCounts()
        {
            lock (gate)
            {
                return flushWriteCounts.ToArray();
            }
        }

        /// <summary>
        ///     Completes once every write issued so far has been flushed, giving the caller a
        ///     happens-before edge on the writer's deferred flush instead of a timing guess.
        /// </summary>
        /// <param name="cancellationToken">Token to abandon the wait.</param>
        /// <returns>A task that completes when no write is left unflushed.</returns>
        public Task WaitForAllWritesFlushedAsync(CancellationToken cancellationToken)
        {
            FlushWaiter waiter;
            lock (gate)
            {
                if (writesSinceLastFlush == 0)
                {
                    return Task.CompletedTask;
                }

                // Any flush drains every pending write, so the next flush is exactly the edge wanted.
                waiter = new FlushWaiter(flushWriteCounts.Count + 1);
                waiters.Add(waiter);
            }

            return AwaitFlushWaiterAsync(waiter, cancellationToken);
        }

        /// <inheritdoc />
        public override bool CanRead => inner.CanRead;

        /// <inheritdoc />
        public override bool CanSeek => inner.CanSeek;

        /// <inheritdoc />
        public override bool CanWrite => inner.CanWrite;

        /// <inheritdoc />
        public override long Length => inner.Length;

        /// <inheritdoc />
        public override long Position
        {
            get => inner.Position;
            set => inner.Position = value;
        }

        /// <inheritdoc />
        public override void Flush()
        {
            RecordFlush();
            inner.Flush();
        }

        /// <inheritdoc />
        public override Task FlushAsync(CancellationToken cancellationToken)
        {
            RecordFlush();
            return inner.FlushAsync(cancellationToken);
        }

        /// <inheritdoc />
        public override int Read(byte[] buffer, int offset, int count) => inner.Read(buffer, offset, count);

        /// <inheritdoc />
        public override Task<int> ReadAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
            => inner.ReadAsync(buffer, offset, count, cancellationToken);

        /// <inheritdoc />
        public override long Seek(long offset, SeekOrigin origin) => inner.Seek(offset, origin);

        /// <inheritdoc />
        public override void SetLength(long value) => inner.SetLength(value);

        /// <inheritdoc />
        public override void Write(byte[] buffer, int offset, int count)
        {
            RecordWrite();
            inner.Write(buffer, offset, count);
        }

        /// <inheritdoc />
        public override Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
        {
            RecordWrite();
            return inner.WriteAsync(buffer, offset, count, cancellationToken);
        }

        private static async Task AwaitFlushWaiterAsync(
            FlushWaiter waiter,
            CancellationToken cancellationToken)
        {
            using (cancellationToken.Register(() => waiter.Completion.TrySetCanceled()))
            {
                await waiter.Completion.Task.ConfigureAwait(false);
            }
        }

        // Counted at write issue, not completion: the peer can observe the bytes as soon as the write
        // is issued, so the pending count must already reflect the write by then.
        private void RecordWrite()
        {
            lock (gate)
            {
                writesSinceLastFlush++;
            }
        }

        private void RecordFlush()
        {
            List<FlushWaiter>? released = null;
            lock (gate)
            {
                flushWriteCounts.Add(writesSinceLastFlush);
                writesSinceLastFlush = 0;
                for (int index = waiters.Count - 1; index >= 0; index--)
                {
                    if (waiters[index].TargetFlushCount <= flushWriteCounts.Count)
                    {
                        released ??= new List<FlushWaiter>();
                        released.Add(waiters[index]);
                        waiters.RemoveAt(index);
                    }
                }
            }

            if (released is null)
            {
                return;
            }

            foreach (FlushWaiter waiter in released)
            {
                waiter.Completion.TrySetResult(true);
            }
        }

        // A pending WaitForAllWritesFlushedAsync call: completes once the recorded flush count
        // reaches TargetFlushCount.
        private sealed class FlushWaiter
        {
            /// <summary>Initializes a waiter released at the given flush ordinal.</summary>
            /// <param name="targetFlushCount">Flush count that releases the waiter.</param>
            public FlushWaiter(int targetFlushCount)
            {
                TargetFlushCount = targetFlushCount;
                Completion = new TaskCompletionSource<bool>(
                    TaskCreationOptions.RunContinuationsAsynchronously);
            }

            /// <summary>Gets the flush count at which this waiter completes.</summary>
            public int TargetFlushCount { get; }

            /// <summary>Gets the completion signaled when the target flush count is reached.</summary>
            public TaskCompletionSource<bool> Completion { 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 = TestNamedPipe.CreateServer(pipeName);
            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);
        }

        /// <inheritdoc />
        public void Dispose()
        {
            WorkerStream.Dispose();
            gatewayStream.Dispose();
        }
    }
}