diff --git a/clients/dotnet/AgentHostProtocol.slnx b/clients/dotnet/AgentHostProtocol.slnx index c9ea641c3..ab398c077 100644 --- a/clients/dotnet/AgentHostProtocol.slnx +++ b/clients/dotnet/AgentHostProtocol.slnx @@ -2,6 +2,8 @@ + + diff --git a/clients/dotnet/README.md b/clients/dotnet/README.md index b5d3839b8..9b8db87c1 100644 --- a/clients/dotnet/README.md +++ b/clients/dotnet/README.md @@ -51,7 +51,8 @@ Reducers.ApplyToSession(state, action); // mutates `state` in place ``` See [`examples/`](examples/) for runnable `ConnectWs` and `ReducersDemo` -console apps. +console apps, and the windowing examples described under +[Windowed delivery](#windowed-delivery-opt-in). ## Dependency injection @@ -87,6 +88,153 @@ state, each returned resource is delivered as a `SubscriptionEventSnapshot` before the host reports `Connected`. Consumers should replace that resource's local reducer state with the snapshot before processing later actions. +## Windowed delivery (opt-in) + +Windowing is off by default. It is offered per subscription, without changing +the protocol versions offered by initialization. A host that omits flow control +declines the offer; the returned handle reports `FlowControl.Mode == "ordinary"` +and `FallbackReason == "declined"`. Its snapshot is still the first event +(when an ordinary empty-state result omits its snapshot, there is no snapshot +event). Ordinary fallback retains the existing drop-oldest buffering behavior. +TCP subscriptions require acceptance and cannot fall back. + +```csharp +var config = new ClientConfig +{ + Windowing = new WindowingConfig + { + WindowBytes = 256 * 1024, + MaximumFrameBytes = 16 * 1024, + MaxBufferedSendBytes = 1024 * 1024, // per subscription + }, +}; +await using var client = AhpClient.Connect(transport, config); +await client.InitializeAsync("my-client"); +var (result, terminal) = await client.SubscribeWindowedAsync("ahp-terminal:/t1"); +await foreach (var evt in terminal.ReadEventsAsync(cancellationToken)) +{ + await RenderOrApplyAsync(evt, cancellationToken); +} // advancing (or disposing) the iterator releases the preceding event +``` + +The additive `IWindowedAhpClient` capability and `IAhpClient` extension methods +also work with clients returned by the DI factory. Existing `IAhpClient`, +`ITransport`, and serializer implementations do not need new interface members. +Per-host defaults use `HostConfig.ClientConfig.Windowing`; the multi-host +runtime offers the same `SubscribeWindowedAsync(host, uri, options)` API. + +For manual release, choose `ReceiveMode.Manual` and process delivery leases: + +```csharp +var (_, tcp) = await client.SubscribeWindowedAsync(tcpUri, + new WindowOptions { Receive = ReceiveMode.Manual, Coalescing = false }); +await foreach (var delivery in tcp.Deliveries.ReadAllAsync(cancellationToken)) +{ + using (delivery) + { + await WriteToBoundedSinkAsync(delivery.Event, cancellationToken); + } // Dispose() or Release() returns credit exactly once +} +``` + +There is one credit-owning handle per URI. Re-subscribing updates that same +handle. `AttachSubscription`, global event streams, and host observer streams +remain lossy observers and **never** release credit. The owning handle retains +every windowed event; it does not silently discard events when a consumer is slow. +New event variants cover snapshots, TCP data/EOF/reset, and unknown notifications. + +Send-side backpressure is exposed independently: + +```csharp +await tcp.WaitWritableAsync(cancellationToken); // advisory, not a reservation +await tcp.NotifyAsync("tcp/data", + new TcpDataParams { Channel = tcpUri, Data = Convert.ToBase64String(bytes) }, + cancellationToken); // completes after transport transmission +var status = tcp.FlowControl; +tcp.FlowControlChanged += (_, current) => UpdateStatus(current); +await foreach (var current in tcp.FlowControlChanges().ReadAllAsync(cancellationToken)) +{ + UpdateStatus(current); +} +``` + +Existing `NotifyAsync` and `DispatchAsync` use the same fair scheduler for an +enabled URI. An oversized single send is accepted when its subscription's queue +is empty; fragmentation and the negotiated window pace it. `MaxBufferedSendBytes` +bounds additional queued data, not the size of that first logical message. +Cancellation removes unsent messages, but an already-started fragmented message +finishes to avoid corrupting the channel. + +Receive coalescing is enabled by default, only for adjacent unconsumed terminal +output or chat deltas targeting the same part with matching metadata. It does +not introduce latency or delivery preferences. TCP sends can combine decoded +bytes; numbered dispatches, client-origin echoes, snapshots and EOF/reset never +combine. Register custom rules per connection: + +```csharp +config.Windowing.Coalescers.Add(new ProgressCoalescer()); // before Connect() + +sealed class ProgressCoalescer : IChannelMessageCoalescer +{ + public JsonRpcNotification? Coalesce(JsonRpcNotification earlier, + JsonRpcNotification later, CoalescingContext context) + { + // Example only when this application's progress events are replaceable. + return context.Direction == CoalescingDirection.Incoming && + earlier.Method == "root/progress" && + earlier.Params?.GetProperty("progressToken").GetString() == + later.Params?.GetProperty("progressToken").GetString() + ? later : null; + } +} +``` + +A rule must preserve channel/method and channel-required event semantics. +Custom coalescers must be fast and synchronous; the library validates their +output and faults invalid results. Merged receive deliveries preserve original +wire charges, not the serialization size of the merged event. + +All protocol sizes are **UTF-16 code units × 2 of the exact JSON text**, not +UTF-8 bytes. `WebSocketTransportOptions.MaxMessageBytes` defaults to 32 MiB of +complete inbound wire-message bytes; its growing 64 KiB receive buffer limits +only a read chunk. The transport-derived AHP frame default is +`2 * (MaxMessageBytes / 3)` = **22,369,620 UTF-16 bytes**. Unknown or unbounded +transport limits use **1,048,576**. Custom transports may implement the optional +`ITransportMessageLimits` capability; raw text/binary payloads are already +retained by `TransportMessage`. + +Managed windowed subscriptions retain their owner and unread deliveries across +reconnect, reset byte counters, and resume with snapshot/replay/live recovery. +Releasing old-connection deliveries never credits the new connection. Missing +subscriptions fail visibly. TCP channels close on connection loss and are never +automatically recreated. Low-level users open a new client and call +`ResumeAsync(ReconnectParams)` with the generated `Windows.Items` wrapper; +`GetWindowedSubscription(uri)` exposes its new owners. + +The receive window does not cap complete reassembled message size or total +memory. Incomplete fragments are credited as they are stored so messages larger +than a window can finish; consumer release holds the final fragment's credit. +Applications must still avoid retaining deliveries indefinitely. + +### Examples + +Two runnable examples use an in-memory demo host that follows the window rules, +so they need no server: + +- [`examples/WindowingBasic`](examples/WindowingBasic/Program.cs) opts a + terminal subscription into windowed delivery and shows the host pausing when + the window is full and resuming as the client returns credit. +- [`examples/WindowingBackpressure`](examples/WindowingBackpressure/Program.cs) + adds a slow renderer, a custom coalescer that keeps only the latest terminal + resize, and a paste that waits for `WaitWritableAsync()` instead of queueing. + +```sh +dotnet run --project examples/WindowingBasic +dotnet run --project examples/WindowingBackpressure +``` + +The demo host and in-memory transport are in [`examples/Shared`](examples/Shared). + ## Observability The client emits OpenTelemetry-native traces and metrics via `System.Diagnostics` diff --git a/clients/dotnet/TELEMETRY.md b/clients/dotnet/TELEMETRY.md index ad392ec2a..300b0500c 100644 --- a/clients/dotnet/TELEMETRY.md +++ b/clients/dotnet/TELEMETRY.md @@ -51,7 +51,37 @@ semantic conventions. | `ahp.client.reconnects` | Counter | `{reconnect}` | `ahp.outcome` | | `ahp.client.events.dropped` | Counter | `{event}` | `ahp.stream` (`subscription` \| `event` \| `state` \| `host-*`) | | `ahp.client.frames.malformed` | Counter | `{frame}` | — | +| `ahp.client.window.data.sent` | Counter | `By` | — | +| `ahp.client.window.data.received` | Counter | `By` | — | +| `ahp.client.window.credit.sent` | Counter | `By` | — | +| `ahp.client.window.credit.received` | Counter | `By` | — | +| `ahp.client.window.blocked.duration` | Histogram | `ms` | — | +| `ahp.client.window.buffered` | ObservableGauge | `By` | `direction` (`send` \| `receive`) | Metric names are lowercase-dotted per OpenTelemetry convention (the C# instrument fields are PascalCase). A consumer's own `ILogger` written inside one of the client's spans auto-correlates to that trace, so no logs are originated here. + +Window data counters measure the actual JSON message/fragment envelopes in +UTF-16 bytes. Credit counters measure the credit notifications' UTF-16 wire +size, not their cumulative `consumedBytes` value. The buffered gauge sums unsent +data and retained receive deliveries across live owning handles, including +completed unread deliveries retained across managed reconnect. Merged receive +items retain the sum of their original charges. The duration histogram records +blocked-to-writable (or closed) intervals, rather than polling samples. +Closed owners are explicitly removed from the process-wide gauge registry; +cleanup does not depend on a metrics listener collecting the gauge. + +Each flow-control transition adds an `ahp.flowcontrol.{kind}` Activity event +with `ahp.flowcontrol.bytes`, without payload content. When no ambient Activity +exists, an `ahp.flowcontrol` span hosts the event; no span is created when +tracing is disabled. + +`WindowedSubscription.FlowControl`, the `FlowControlChanged` event, and +`FlowControlChanges()` expose current limits, fallback reason, buffered and +outstanding amounts, and writability. `client.FlowControlEvents()` is a lossy, +observe-only connection stream carrying channel, kind and status. Kinds are +`negotiated`, `declined`, `blocked`, `unblocked`, `creditSent`, `creditReceived`, +`fragmentedMessageReceived`, `fragmentedMessageSent`, `coalesced`, `resumed`, +and `dropped`. `MultiHostClient.FlowControlEvents()` adds host attribution. +Reading these streams does not consume deliveries or return credit. diff --git a/clients/dotnet/examples/Shared/DemoSupport.cs b/clients/dotnet/examples/Shared/DemoSupport.cs new file mode 100644 index 000000000..3b94330c6 --- /dev/null +++ b/clients/dotnet/examples/Shared/DemoSupport.cs @@ -0,0 +1,37 @@ +// Helpers shared by the windowing examples. +#nullable enable + +using System.Collections.Generic; +using System.Runtime.CompilerServices; +using System.Text.Encodings.Web; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.AgentHostProtocol; + +namespace WindowingExamples; + +internal static class DemoSupport +{ + private static readonly JsonSerializerOptions QuoteOptions = + new() { Encoder = JavaScriptEncoder.UnsafeRelaxedJsonEscaping }; + + /// Quotes text the way JavaScript's JSON.stringify does. + public static string Quote(string text) => JsonSerializer.Serialize(text, QuoteOptions); + + /// + /// Takes deliveries in order and releases each one when the next one is + /// taken, or when the loop ends. This is the automatic release that + /// performs, but it also + /// exposes each delivery's charged size (), + /// which the examples print. Subscribe with . + /// + public static async IAsyncEnumerable TakeDeliveriesAsync( + this WindowedSubscription subscription, + [EnumeratorCancellation] CancellationToken cancellationToken = default) + { + await foreach (var delivery in subscription.Deliveries.ReadAllAsync(cancellationToken).ConfigureAwait(false)) + { + using (delivery) yield return delivery; + } + } +} diff --git a/clients/dotnet/examples/Shared/InMemoryTransport.cs b/clients/dotnet/examples/Shared/InMemoryTransport.cs new file mode 100644 index 000000000..6f9f74330 --- /dev/null +++ b/clients/dotnet/examples/Shared/InMemoryTransport.cs @@ -0,0 +1,62 @@ +// An in-memory ITransport pair for the windowing examples. +// +// Every message is a text frame built with TransportMessage.FromText, so the +// client sees the exact JSON text the host wrote and can charge each windowed +// message its UTF-16 size. Closing either half closes both. +#nullable enable + +using System.Threading; +using System.Threading.Channels; +using System.Threading.Tasks; +using Microsoft.AgentHostProtocol; + +namespace WindowingExamples; + +internal sealed class InMemoryTransport : ITransport +{ + private readonly Channel _inbox; + private readonly Channel _outbox; + + private InMemoryTransport(Channel inbox, Channel outbox) + { + _inbox = inbox; + _outbox = outbox; + } + + /// Creates a linked pair: what one half sends, the other half receives. + public static (InMemoryTransport Client, InMemoryTransport Host) CreatePair() + { + var toHost = Channel.CreateUnbounded(new UnboundedChannelOptions { SingleReader = true }); + var toClient = Channel.CreateUnbounded(new UnboundedChannelOptions { SingleReader = true }); + return (new InMemoryTransport(toClient, toHost), new InMemoryTransport(toHost, toClient)); + } + + public ValueTask SendAsync(TransportMessage message, CancellationToken cancellationToken = default) + { + cancellationToken.ThrowIfCancellationRequested(); + return _outbox.Writer.TryWrite(message) + ? default + : throw new AhpTransportException("closed"); + } + + public async ValueTask ReceiveAsync(CancellationToken cancellationToken = default) + { + try + { + return await _inbox.Reader.ReadAsync(cancellationToken).ConfigureAwait(false); + } + catch (ChannelClosedException error) + { + throw new AhpTransportException("closed", inner: error); + } + } + + public ValueTask CloseAsync(CancellationToken cancellationToken = default) + { + _outbox.Writer.TryComplete(); + _inbox.Writer.TryComplete(); + return default; + } + + public ValueTask DisposeAsync() => CloseAsync(); +} diff --git a/clients/dotnet/examples/Shared/WindowingDemoHost.cs b/clients/dotnet/examples/Shared/WindowingDemoHost.cs new file mode 100644 index 000000000..3edb5225f --- /dev/null +++ b/clients/dotnet/examples/Shared/WindowingDemoHost.cs @@ -0,0 +1,324 @@ +// A small in-memory AHP host for the windowing examples. +// +// It implements only what the examples need: `initialize`, a windowed +// `subscribe` to one terminal, and the host side of the window rules: +// +// - It counts every data message it sends, in UTF-16 bytes, and stops while +// `sent - consumed` is at or above the client's window. +// - It counts every data message the client sends, and returns +// `channel/credit` after it has processed each one. +// +// A real host applies the same rules to every windowed subscription. +#nullable enable + +using System.Text; +using System.Text.Encodings.Web; +using System.Threading; +using System.Threading.Channels; +using System.Threading.Tasks; +using Microsoft.AgentHostProtocol; + +namespace WindowingExamples; + +internal sealed class DemoHostOptions +{ + /// The largest window the host accepts for data that it sends. + public required long ClientWindowBytes { get; init; } + + /// The window the host offers for data that the client sends. + public required long HostWindowBytes { get; init; } + + /// How long the host takes to process one message from the client. + public TimeSpan InputDelay { get; init; } +} + +internal sealed class WindowingDemoHost : IAsyncDisposable +{ + private const long MaximumFrameBytes = 16_384; + + // Escape JSON the way JavaScript's JSON.stringify does, so the sizes match + // the TypeScript version of this example. + private static readonly SystemTextJsonAhpSerializer Json = new( + new JsonSerializerOptions { Encoder = JavaScriptEncoder.UnsafeRelaxedJsonEscaping }); + + private readonly DemoHostOptions _options; + private readonly InMemoryTransport _transport; + private readonly TaskCompletionSource _subscribed = new(TaskCreationOptions.RunContinuationsAsynchronously); + private readonly SemaphoreSlim _dataLock = new(1, 1); + private readonly object _gate = new(); + private readonly Channel<(long Bytes, StateAction Action)> _input = + Channel.CreateUnbounded<(long Bytes, StateAction Action)>(new UnboundedChannelOptions { SingleReader = true }); + private readonly StringBuilder _receivedInput = new(); + private readonly Task _run; + private readonly Task _processInput; + private ChannelFlowControl? _limits; + private string _channel = ""; + private long _serverSeq; + + // Data that the host sends to the client. + private long _sentBytes; + private long _consumedBytes; + private TaskCompletionSource _credit = new(TaskCreationOptions.RunContinuationsAsynchronously); + + // Data that the client sends to the host. + private long _receivedBytes; + private long _processedBytes; + + public WindowingDemoHost(DemoHostOptions options) + { + _options = options; + (var client, _transport) = InMemoryTransport.CreatePair(); + ClientTransport = client; + _run = RunAsync(); + _processInput = ProcessInputAsync(); + } + + /// The half of the transport to give to . + public ITransport ClientTransport { get; } + + /// Terminal input that the host has processed, in order. + public string ReceivedInput + { + get { lock (_gate) return _receivedInput.ToString(); } + } + + /// Completes after the host has accepted the subscription and sent the snapshot. + public Task WhenSubscribedAsync() => _subscribed.Task; + + /// Sends one terminal action, first waiting for credit if the client's window is full. + public async Task SendActionAsync(StateAction action, CancellationToken cancellationToken = default) + { + await _dataLock.WaitAsync(cancellationToken).ConfigureAwait(false); + try + { + long window; + lock (_gate) + window = (_limits ?? throw new InvalidOperationException("no windowed subscription")).ClientReceive.WindowBytes; + if (Outstanding() >= window) + { + Log($"window full ({Outstanding()} of {window} bytes outstanding); pausing output"); + while (true) + { + Task credit; + lock (_gate) + { + if (_sentBytes - _consumedBytes < window) break; + credit = _credit.Task; + } + await credit.WaitAsync(cancellationToken).ConfigureAwait(false); + } + long consumed; + lock (_gate) consumed = _consumedBytes; + Log($"credit received (consumedBytes={consumed}); resuming output"); + } + await SendDataAsync(new JsonRpcNotification + { + Method = "action", + Params = Json.SerializeToElement(new ActionEnvelope + { + Channel = _channel, + Action = action, + ServerSeq = ++_serverSeq, + }), + }).ConfigureAwait(false); + } + finally + { + _dataLock.Release(); + } + } + + public async ValueTask DisposeAsync() + { + await _transport.CloseAsync().ConfigureAwait(false); + await Task.WhenAll(_run, _processInput).ConfigureAwait(false); + _dataLock.Dispose(); + } + + private static void Log(string message) => Console.WriteLine($" host | {message}"); + + private static long Utf16Bytes(string text) => (long)text.Length * 2; + + private long Outstanding() + { + lock (_gate) return _sentBytes - _consumedBytes; + } + + // Callers hold _dataLock, so data messages keep their serverSeq order. + private async Task SendDataAsync(JsonRpcNotification notification) + { + var text = Json.EncodeMessage(new JsonRpcMessage { Notification = notification }).Text!; + lock (_gate) _sentBytes += Utf16Bytes(text); + await _transport.SendAsync(TransportMessage.FromText(text)).ConfigureAwait(false); + } + + private ValueTask SendControlAsync(JsonRpcMessage message) => + _transport.SendAsync(Json.EncodeMessage(message)); + + private static JsonRpcMessage Response(ulong id, T result) => new() + { + SuccessResponse = new JsonRpcSuccessResponse { Id = id, Result = Json.SerializeToElement(result) }, + }; + + private async Task RunAsync() + { + try + { + while (true) + { + TransportMessage frame; + try + { + frame = await _transport.ReceiveAsync().ConfigureAwait(false); + } + catch (AhpTransportException) + { + return; + } + if (frame.Frame != TransportFrame.Text) continue; + var message = Json.DecodeMessage(frame); + switch (message.Request?.Method ?? message.Notification?.Method) + { + case "initialize": + await SendControlAsync(Response(message.Request!.Id, new InitializeResult + { + ProtocolVersion = ProtocolVersion.Supported[0], + ServerSeq = 0, + Snapshots = new List(), + })).ConfigureAwait(false); + break; + case "subscribe": + await AcceptSubscriptionAsync(message.Request!).ConfigureAwait(false); + break; + case "channel/credit": + ReceiveCredit(Json.Deserialize(message.Notification!.Params!.Value)); + break; + case "dispatchAction": + ReceiveInput(frame.Text!, Json.Deserialize(message.Notification!.Params!.Value).Action); + break; + } + } + } + finally + { + _input.Writer.TryComplete(); + } + } + + private void ReceiveCredit(ChannelCreditParams credit) + { + lock (_gate) + { + // Credit is a running total, so ignore values that are not larger. + if (credit.ConsumedBytes <= _consumedBytes) return; + _consumedBytes = credit.ConsumedBytes; + _credit.TrySetResult(); + _credit = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + } + } + + private async Task AcceptSubscriptionAsync(JsonRpcRequest request) + { + var parameters = Json.Deserialize(request.Params!.Value); + var offered = parameters.FlowControl?.Receive + ?? throw new InvalidOperationException("the examples expect a windowed subscribe"); + await _dataLock.WaitAsync().ConfigureAwait(false); + try + { + _channel = parameters.Channel; + // A host can lower the limits that the client offered, but never raise them. + var limits = new ChannelFlowControl + { + ClientReceive = new ChannelReceiveLimits + { + WindowBytes = Math.Min(offered.WindowBytes, _options.ClientWindowBytes), + MaximumFrameBytes = Math.Min(offered.MaximumFrameBytes, MaximumFrameBytes), + }, + HostReceive = new ChannelReceiveLimits + { + WindowBytes = _options.HostWindowBytes, + MaximumFrameBytes = MaximumFrameBytes, + }, + }; + lock (_gate) _limits = limits; + Log($"accepted windowed delivery: client window {limits.ClientReceive.WindowBytes} bytes " + + $"(offered {offered.WindowBytes}), host window {limits.HostReceive.WindowBytes} bytes"); + await SendControlAsync(Response(request.Id, new SubscribeResult { FlowControl = limits })).ConfigureAwait(false); + + // The first data message on a windowed subscription is always the snapshot. + await SendDataAsync(new JsonRpcNotification + { + Method = "channel/snapshot", + Params = Json.SerializeToElement(new ChannelSnapshotParams + { + Channel = _channel, + Snapshot = new Snapshot + { + Resource = _channel, + State = new SnapshotState + { + Terminal = new TerminalState + { + Title = "demo shell", + Content = new List(), + Lifecycle = new TerminalLifecycleState(new TerminalRunningLifecycleState()), + Claim = new TerminalClaim(new TerminalClientClaim { ClientId = "windowing-example" }), + }, + }, + FromSeq = _serverSeq, + }, + }), + }).ConfigureAwait(false); + } + finally + { + _dataLock.Release(); + } + _subscribed.TrySetResult(); + } + + private void ReceiveInput(string text, StateAction action) + { + var bytes = Utf16Bytes(text); + lock (_gate) + { + _receivedBytes += bytes; + var limits = _limits!.HostReceive; + // The client can go past the window only with the message that fills it. + if (_receivedBytes - _processedBytes > limits.WindowBytes + limits.MaximumFrameBytes) + Log("client went past the host window"); + } + _input.Writer.TryWrite((bytes, action)); + } + + // Process input one message at a time, then return credit for it. + private async Task ProcessInputAsync() + { + await foreach (var (bytes, action) in _input.Reader.ReadAllAsync().ConfigureAwait(false)) + { + await Task.Delay(_options.InputDelay).ConfigureAwait(false); + long processed; + lock (_gate) + { + if (action.Value is TerminalInputAction input) _receivedInput.Append(input.Data); + _processedBytes += bytes; + processed = _processedBytes; + } + try + { + await SendControlAsync(new JsonRpcMessage + { + Notification = new JsonRpcNotification + { + Method = "channel/credit", + Params = Json.SerializeToElement(new ChannelCreditParams { Channel = _channel, ConsumedBytes = processed }), + }, + }).ConfigureAwait(false); + } + catch (AhpTransportException) + { + return; + } + } + } +} diff --git a/clients/dotnet/examples/WindowingBackpressure/Program.cs b/clients/dotnet/examples/WindowingBackpressure/Program.cs new file mode 100644 index 000000000..ff69fae51 --- /dev/null +++ b/clients/dotnet/examples/WindowingBackpressure/Program.cs @@ -0,0 +1,170 @@ +// Windowed delivery with backpressure and a custom coalescer. +// +// Part 1, receiving: the client renders slowly. While it is busy, the host +// fills the window and the client queues what arrives. Adjacent queued +// messages are combined before the renderer sees them: +// +// - The built-in coalescer joins adjacent `terminal/data` output. +// - A custom coalescer keeps only the last of adjacent `terminal/resized` +// actions, because only the final size matters. +// +// Part 2, sending: the client pastes a large block of text into the terminal. +// The host processes input slowly, so the client's send window fills. The +// client waits until the subscription is writable before it sends each chunk, +// instead of queueing the whole paste in memory. +// +// Run: dotnet run --project examples/WindowingBackpressure +#nullable enable + +using System; +using System.Collections.Generic; +using System.Linq; +using System.Threading.Tasks; +using Microsoft.AgentHostProtocol; +using WindowingExamples; + +const string Terminal = "ahp-terminal:/demo"; + +var host = new WindowingDemoHost(new DemoHostOptions +{ + ClientWindowBytes = 2048, + HostWindowBytes = 1024, + InputDelay = TimeSpan.FromMilliseconds(15), +}); +var config = new ClientConfig(); +config.Windowing.Coalescers.Add(new LatestResizeCoalescer()); +// Allow at most 2048 bytes of unsent data to wait behind the message being sent. +config.Windowing.MaxBufferedSendBytes = 2048; +var client = AhpClient.Connect(host.ClientTransport, config); +await client.InitializeAsync("windowing-example", ProtocolVersion.Supported); + +// Count what the client did. This stream is for monitoring only and can drop +// events; it never affects credit. +var counts = new SortedDictionary(StringComparer.Ordinal); +var counting = CountFlowControlEventsAsync(); + +// Manual receive mode exposes each delivery's charged size; TakeDeliveriesAsync +// still releases deliveries automatically. Coalescing is on by default. +var (_, subscription) = await client.SubscribeWindowedAsync(Terminal, + new WindowOptions { WindowBytes = 4096, Receive = ReceiveMode.Manual }); +Log($"subscribed: accepted window={subscription.FlowControl.Limits?.ClientReceive.WindowBytes} bytes"); + +// ─── Part 1: a slow renderer ──────────────────────────────────────────────── + +Console.WriteLine("\nPart 1: slow renderer"); +StateAction[] hostActions = +[ + Data("Building 120 modules\r\n"), + // The user drags the corner of the terminal window. + Resize(100, 30), + Resize(110, 32), + Resize(120, 36), + Resize(132, 40), + .. Enumerable.Range(1, 10) + .Select(step => step * 10) + .Select(percent => Data($"[{percent,3}%] compiled {percent * 120 / 100} modules\r\n")), + Data("Build finished\r\n"), +]; +var writer = WriteHostActionsAsync(); + +var renders = 0; +// Leaving this loop releases the last delivery but keeps the subscription open. +await foreach (var delivery in subscription.TakeDeliveriesAsync()) +{ + renders++; + if (delivery.Event is SubscriptionEventSnapshot) + { + Log($"render snapshot ({delivery.WireBytes} bytes)"); + } + else if (delivery.Event is SubscriptionEventAction { Envelope.Action.Value: TerminalResizedAction resize }) + { + Log($"render resize to {resize.Cols}x{resize.Rows} ({delivery.WireBytes} bytes)"); + } + else if (delivery.Event is SubscriptionEventAction { Envelope.Action.Value: TerminalDataAction data }) + { + var lines = data.Data.Split("\r\n", StringSplitOptions.RemoveEmptyEntries); + Log($"render {lines.Length} line(s) of output ({delivery.WireBytes} bytes): {DemoSupport.Quote(lines[^1])}"); + if (data.Data.Contains("Build finished", StringComparison.Ordinal)) break; + } + await Task.Delay(60); // The renderer is slow; the host fills the window meanwhile. +} +await writer; +Log($"rendered {renders} deliveries for {hostActions.Length + 1} host messages"); + +// ─── Part 2: pasting into a slow terminal ─────────────────────────────────── + +Console.WriteLine("\nPart 2: pasting into a slow terminal"); +var paste = Enumerable.Range(1, 8).Select(i => $"echo \"line {i} of the pasted script\"\n").ToArray(); +var waits = 0; +for (var i = 0; i < paste.Length; i++) +{ + // Wait for room in the host window and the local send queue, rather than + // queueing the whole paste. With one producer, this keeps sends within the limit. + var status = subscription.FlowControl; + if (!status.Writable) + { + Log($"not writable ({status.OutstandingSendBytes} bytes awaiting credit, {status.BufferedSendBytes} bytes queued); waiting"); + waits++; + await subscription.WaitWritableAsync(); + } + Log($"send chunk {i + 1} of {paste.Length}"); + // Completes when the message is written to the transport. If another + // producer filled the local send queue first, the .NET client waits for + // room in the queue rather than failing, so there is nothing to retry. + await client.DispatchAsync(Terminal, new StateAction(new TerminalInputAction { Data = paste[i] })); +} +// Wait until the host has processed, and credited, every chunk. +await foreach (var current in subscription.FlowControlChanges().ReadAllAsync()) +{ + if (current.OutstandingSendBytes == 0 && current.BufferedSendBytes == 0) break; +} +var intact = host.ReceivedInput == string.Concat(paste); +Log($"sent {paste.Length} chunks, waiting {waits} times; host received the paste intact: {(intact ? "true" : "false")}"); + +await subscription.DisposeAsync(); +await client.DisposeAsync(); +await host.DisposeAsync(); + +// Shutting down the client completes the event stream, so every event is counted. +await counting; +Console.WriteLine("\nFlow-control events:"); +foreach (var (kind, count) in counts) Console.WriteLine($" {kind}: {count}"); + +async Task CountFlowControlEventsAsync() +{ + await foreach (var flowEvent in client.FlowControlEvents().ReadAllAsync()) + { + counts[flowEvent.Kind] = counts.GetValueOrDefault(flowEvent.Kind) + 1; + } +} + +async Task WriteHostActionsAsync() +{ + await host.WhenSubscribedAsync(); + foreach (var action in hostActions) await host.SendActionAsync(action); +} + +static StateAction Data(string data) => new(new TerminalDataAction { Data = data }); + +static StateAction Resize(long cols, long rows) => new(new TerminalResizedAction { Cols = cols, Rows = rows }); + +static void Log(string message) => Console.WriteLine($" client | {message}"); + +// A custom coalescer runs on two adjacent messages that are still queued. It +// returns the combined message, or null to leave them separate. The client +// never passes it messages that the protocol requires to stay separate, such +// as actions with an origin. +internal sealed class LatestResizeCoalescer : IChannelMessageCoalescer +{ + public JsonRpcNotification? Coalesce(JsonRpcNotification earlier, JsonRpcNotification later, CoalescingContext context) + { + if (context.Direction != CoalescingDirection.Incoming || earlier.Method != "action") return null; + return IsResize(earlier) && IsResize(later) ? later : null; + } + + private static bool IsResize(JsonRpcNotification notification) => + notification.Params is { } parameters + && parameters.TryGetProperty("action", out var action) + && action.TryGetProperty("type", out var type) + && type.GetString() == ActionType.TerminalResized.Value; +} diff --git a/clients/dotnet/examples/WindowingBackpressure/WindowingBackpressure.csproj b/clients/dotnet/examples/WindowingBackpressure/WindowingBackpressure.csproj new file mode 100644 index 000000000..77766f1ee --- /dev/null +++ b/clients/dotnet/examples/WindowingBackpressure/WindowingBackpressure.csproj @@ -0,0 +1,21 @@ + + + + net8.0 + Exe + false + false + + Major + + + + + + + + + + + + diff --git a/clients/dotnet/examples/WindowingBasic/Program.cs b/clients/dotnet/examples/WindowingBasic/Program.cs new file mode 100644 index 000000000..7f71d06db --- /dev/null +++ b/clients/dotnet/examples/WindowingBasic/Program.cs @@ -0,0 +1,86 @@ +// Basic windowed delivery. +// +// The client offers windowed delivery on a terminal subscription. The host +// accepts a smaller window, then sends a snapshot and terminal output. When +// the window is full, the host waits; each time the client takes the next +// delivery, it releases the previous one and returns credit, so the host can +// continue. +// +// Run: dotnet run --project examples/WindowingBasic +#nullable enable + +using System; +using System.Threading.Tasks; +using Microsoft.AgentHostProtocol; +using WindowingExamples; + +const string Terminal = "ahp-terminal:/demo"; +string[] output = +[ + "$ npm test", + "> demo@1.0.0 test", + "> node --test", + "ok 1 - parses the configuration file", + "ok 2 - rejects an unknown option", + "ok 3 - reads input from stdin", + "ok 4 - writes the report", + "# pass 4, fail 0", +]; + +// The host accepts at most a 1024-byte window, so it fills after a few messages. +await using var host = new WindowingDemoHost(new DemoHostOptions { ClientWindowBytes = 1024, HostWindowBytes = 4096 }); +await using var client = AhpClient.Connect(host.ClientTransport); +await client.InitializeAsync("windowing-example", ProtocolVersion.Supported); + +// Windowing is off by default. Opt in for this subscription and offer a +// 4096-byte window. Coalescing is turned off so that every host message is a +// separate delivery; see WindowingBackpressure for coalescing. Manual receive +// mode exposes each delivery's charged size; TakeDeliveriesAsync still +// releases deliveries automatically. +var (_, subscription) = await client.SubscribeWindowedAsync(Terminal, new WindowOptions +{ + WindowBytes = 4096, + Coalescing = false, + Receive = ReceiveMode.Manual, +}); +var flowControl = subscription.FlowControl; +Log($"subscribed: mode={flowControl.Mode}, accepted window={flowControl.Limits?.ClientReceive.WindowBytes} bytes"); + +// The host writes terminal output as fast as the window allows. +var writer = WriteOutputAsync(); + +// With automatic release, taking the next delivery releases the previous one. +// The client then returns credit for it. +var lines = 0; +await foreach (var delivery in subscription.TakeDeliveriesAsync()) +{ + var buffered = subscription.FlowControl.BufferedReceiveBytes; + switch (delivery.Event) + { + case SubscriptionEventSnapshot: + Log($"snapshot ({delivery.WireBytes} bytes; {buffered} bytes not yet released)"); + break; + case SubscriptionEventAction { Envelope.Action.Value: TerminalDataAction data }: + Log($"output {DemoSupport.Quote(data.Data.TrimEnd())} ({delivery.WireBytes} bytes; {buffered} bytes not yet released)"); + lines++; + break; + } + await Task.Delay(20); // Simulate rendering the delivery. + if (lines == output.Length) break; // Leaving the loop releases the last delivery. +} +await subscription.DisposeAsync(); // Unsubscribe. + +Log($"received all {lines} lines"); +await writer; +// Leaving the program disposes the client, which shuts it down, and then the host. + +async Task WriteOutputAsync() +{ + await host.WhenSubscribedAsync(); + foreach (var line in output) + { + await host.SendActionAsync(new StateAction(new TerminalDataAction { Data = line + "\r\n" })); + } +} + +static void Log(string message) => Console.WriteLine($" client | {message}"); diff --git a/clients/dotnet/examples/WindowingBasic/WindowingBasic.csproj b/clients/dotnet/examples/WindowingBasic/WindowingBasic.csproj new file mode 100644 index 000000000..77766f1ee --- /dev/null +++ b/clients/dotnet/examples/WindowingBasic/WindowingBasic.csproj @@ -0,0 +1,21 @@ + + + + net8.0 + Exe + false + false + + Major + + + + + + + + + + + + diff --git a/clients/dotnet/src/AgentHostProtocol.Abstractions/Generated/State.generated.cs b/clients/dotnet/src/AgentHostProtocol.Abstractions/Generated/State.generated.cs index 9d5ab51ce..63b8dc9e0 100644 --- a/clients/dotnet/src/AgentHostProtocol.Abstractions/Generated/State.generated.cs +++ b/clients/dotnet/src/AgentHostProtocol.Abstractions/Generated/State.generated.cs @@ -8013,11 +8013,10 @@ public override SnapshotState Read(ref Utf8JsonReader reader, Type typeToConvert { result.Chat = root.Deserialize(AhpJsonTypeInfo.Get(options)); } - else if (root.TryGetProperty("lifecycle", out _)) + else if (root.TryGetProperty("lifecycle", out _) && root.TryGetProperty("provider", out _)) { - // SessionState is discriminated on its required `lifecycle` field. - // (It no longer carries `summary`; that field was removed when the - // session state was flattened.) + // SessionState is discriminated on its required `lifecycle` and + // `provider` fields; TerminalState also has `lifecycle`. result.Session = root.Deserialize(AhpJsonTypeInfo.Get(options)); } else if (root.TryGetProperty("instanceId", out _) && diff --git a/clients/dotnet/src/AgentHostProtocol.Abstractions/Transport/ITransportMessageLimits.cs b/clients/dotnet/src/AgentHostProtocol.Abstractions/Transport/ITransportMessageLimits.cs new file mode 100644 index 000000000..f3545bbd0 --- /dev/null +++ b/clients/dotnet/src/AgentHostProtocol.Abstractions/Transport/ITransportMessageLimits.cs @@ -0,0 +1,10 @@ +#nullable enable + +namespace Microsoft.AgentHostProtocol; + +/// Optional transport capability describing its complete inbound wire-message limit. +public interface ITransportMessageLimits +{ + /// Maximum wire bytes per message, or null when unknown or unlimited. + long? MaximumInboundMessageBytes { get; } +} diff --git a/clients/dotnet/src/AgentHostProtocol/AhpClient.cs b/clients/dotnet/src/AgentHostProtocol/AhpClient.cs index e90eb47d0..b9b74417e 100644 --- a/clients/dotnet/src/AgentHostProtocol/AhpClient.cs +++ b/clients/dotnet/src/AgentHostProtocol/AhpClient.cs @@ -7,6 +7,7 @@ using System.Collections.Generic; using System.Diagnostics; using System.Diagnostics.CodeAnalysis; +using System.Linq; using System.Text.Json; using System.Threading; using System.Threading.Channels; @@ -72,6 +73,8 @@ public static KeepAlivePolicy Enabled(TimeSpan? interval = null, TimeSpan? timeo /// Tuning knobs for an . public sealed class ClientConfig { + /// Defaults for subscriptions explicitly opting into windowed delivery. + public WindowingConfig Windowing { get; set; } = new(); /// /// How long waits for a /// response. Zero disables the timeout. Defaults to 30 seconds. @@ -112,6 +115,7 @@ internal static ClientConfig Snapshot(ClientConfig? config) ? source.SubscriptionBufferCapacity : 256, KeepAlive = source.KeepAlive ?? KeepAlivePolicy.Disabled, + Windowing = (source.Windowing ?? new WindowingConfig()).Snapshot(), }; } } @@ -156,7 +160,7 @@ public enum ConnectionState /// All public methods are safe to call from multiple threads. /// /// -public sealed class AhpClient : IAhpClient +public sealed partial class AhpClient : IWindowedAhpClient { // ── State that lives for the client lifetime ────────────────────────── @@ -240,6 +244,7 @@ internal int SubscriptionCount // Lifecycle private readonly TaskCompletionSource _doneTcs = new(TaskCreationOptions.RunContinuationsAsynchronously); + private readonly TaskCompletionSource _teardownTcs = new(TaskCreationOptions.RunContinuationsAsynchronously); private readonly CancellationTokenSource _lifetimeCts = new(); private int _shutdownStarted; // 0 = running, 1 = shut down private Exception? _closeErr; @@ -254,6 +259,8 @@ private sealed class OutboundMessage public JsonRpcMessage Message { get; } public TaskCompletionSource? Sent { get; } public CancellationToken CancellationToken { get; } + public bool IsBarrier { get; init; } + public WindowResponse? WindowRequest { get; init; } public OutboundMessage( JsonRpcMessage message, @@ -329,11 +336,26 @@ public static AhpClient Connect( /// public async Task ShutdownAsync(CancellationToken cancellationToken = default) { - await ShutdownWithErrorAsync(null).ConfigureAwait(false); + try + { + if (Volatile.Read(ref _shutdownStarted) == 0) + await DrainOutboundAsync(cancellationToken).ConfigureAwait(false); + } + catch (Exception) when (!cancellationToken.IsCancellationRequested) { } + finally { await ShutdownWithErrorAsync(null).ConfigureAwait(false); } + await _teardownTcs.Task.ConfigureAwait(false); // Wait for both background tasks to exit. await Task.WhenAll(_readerTask, _writerTask).WaitAsync(cancellationToken).ConfigureAwait(false); } + // An uninstalled connection attempt has no application sends to drain. + internal async Task AbortAsync() + { + await ShutdownWithErrorAsync(null).ConfigureAwait(false); + await _teardownTcs.Task.ConfigureAwait(false); + await Task.WhenAll(_readerTask, _writerTask).ConfigureAwait(false); + } + /// public async ValueTask DisposeAsync() { @@ -360,6 +382,7 @@ private async Task ShutdownWithErrorAsync(Exception? cause, bool fromKeepAlive = Volatile.Write(ref _closeErr, cause); _lifetimeCts.Cancel(); + EndWindows(cause); // Signal the done task so Done-waiters unblock. _doneTcs.TrySetResult(true); @@ -400,6 +423,7 @@ private async Task ShutdownWithErrorAsync(Exception? cause, bool fromKeepAlive = var shutdownEx = cause is null ? new AhpClientClosedException() : new AhpClientClosedException($"ahp: client shut down: {cause.Message}"); + while (_outbound.Reader.TryRead(out var queued)) queued.Sent?.TrySetException(shutdownEx); foreach (var kv in _pending) { @@ -439,6 +463,17 @@ private async Task ShutdownWithErrorAsync(Exception? cause, bool fromKeepAlive = _stateListeners.Clear(); } foreach (var st in allStateListeners) st.Close(); + _teardownTcs.TrySetResult(true); + } + + private async Task DrainOutboundAsync(CancellationToken cancellationToken) + { + using var timeout = new CancellationTokenSource(TimeSpan.FromSeconds(2)); + using var linked = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, timeout.Token, _lifetimeCts.Token); + var sent = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + await _outbound.Writer.WriteAsync(new OutboundMessage(new JsonRpcMessage(), sent) { IsBarrier = true }, linked.Token).ConfigureAwait(false); + _writerWake.Release(); + await sent.Task.WaitAsync(linked.Token).ConfigureAwait(false); } // ── Connection state ────────────────────────────────────────────────── @@ -533,49 +568,65 @@ await ShutdownWithErrorAsync( private async Task RunWriterAsync() { + OutboundMessage? active = null; try { - await foreach (var item in _outbound.Reader.ReadAllAsync(_lifetimeCts.Token).ConfigureAwait(false)) + var controlBurst = 0; + while (!_lifetimeCts.IsCancellationRequested) { - if (item.CancellationToken.IsCancellationRequested) - { - item.Sent?.TrySetCanceled(item.CancellationToken); - continue; - } - - var frame = _serializer.EncodeMessage(item.Message); - if (item.CancellationToken.IsCancellationRequested) - { - item.Sent?.TrySetCanceled(item.CancellationToken); - continue; - } - try + TransportMessage frame; + Action? sent = null; + active = null; + if ((controlBurst >= 8 || !_outbound.Reader.TryPeek(out _)) && TryWindowFrame(out frame, out sent)) { - // Canceling ClientWebSocket.SendAsync aborts the shared connection, - // so message cancellation only prevents writes that have not started. - await _transport.SendAsync(frame, _lifetimeCts.Token).ConfigureAwait(false); - item.Sent?.TrySetResult(true); + controlBurst = 0; } - catch (OperationCanceledException) when (_lifetimeCts.IsCancellationRequested) + else if (_outbound.Reader.TryRead(out var item)) { - item.Sent?.TrySetException(new AhpClientClosedException()); - return; + active = item; + controlBurst++; + if (item.IsBarrier) + { + item.Sent?.TrySetResult(true); + continue; + } + if (item.CancellationToken.IsCancellationRequested) + { + item.Sent?.TrySetCanceled(item.CancellationToken); + continue; + } + frame = _serializer.EncodeMessage(item.Message); + if (item.CancellationToken.IsCancellationRequested) + { + item.Sent?.TrySetCanceled(item.CancellationToken); + continue; + } + if (item.WindowRequest is { } response && !response.TryStart()) + { + item.Sent?.TrySetException(response.Failure ?? new AhpClientClosedException()); + continue; + } } - catch (Exception ex) + else { - item.Sent?.TrySetException(ex); - await ShutdownWithErrorAsync(new AhpTransportException("io", $"ahp: transport send: {ex.Message}", ex)).ConfigureAwait(false); - return; + await _writerWake.WaitAsync(_lifetimeCts.Token).ConfigureAwait(false); + continue; } + await _transport.SendAsync(frame, _lifetimeCts.Token).ConfigureAwait(false); + sent?.Invoke(); + active?.Sent?.TrySetResult(true); } } catch (OperationCanceledException) when (_lifetimeCts.IsCancellationRequested) { + active?.Sent?.TrySetException(new AhpClientClosedException()); // Normal teardown canceled the channel enumeration. } catch (Exception ex) { - await ShutdownWithErrorAsync(new AhpTransportException("io", $"ahp: writer: {ex.Message}", ex)).ConfigureAwait(false); + active?.Sent?.TrySetException(ex); + await ShutdownWithErrorAsync(ex is AhpTransportException ? ex + : new AhpTransportException("io", $"ahp: writer: {ex.Message}", ex)).ConfigureAwait(false); } } @@ -617,24 +668,28 @@ private async Task RunReaderAsync() } catch { + lock (_windowGate) + if (_windows.Values.Any(window => window.Limits is not null)) + throw Protocol("malformed JSON on a windowed connection"); // Skip malformed frames; protocol resync is the server's responsibility. AhpTelemetry.MalformedFrames.Add(1); continue; } AhpTelemetry.MessagesReceived.Add(1); - Dispatch(msg); + Dispatch(msg, msg.Notification is not null && RequiresWindowSizes ? RawWindowBytes(frame) : 0); } } catch (Exception ex) when (ex is not OperationCanceledException) { - await ShutdownWithErrorAsync(new AhpTransportException("io", $"ahp: reader: {ex.Message}", ex)).ConfigureAwait(false); + await ShutdownWithErrorAsync(ex is AhpTransportException ? ex + : new AhpTransportException("protocol", $"ahp: reader: {ex.Message}", ex)).ConfigureAwait(false); } } // ── Dispatch ────────────────────────────────────────────────────────── - private void Dispatch(JsonRpcMessage msg) + private void Dispatch(JsonRpcMessage msg, long wireBytes) { if (msg.SuccessResponse is not null) { @@ -647,7 +702,8 @@ private void Dispatch(JsonRpcMessage msg) } else if (msg.Notification is not null) { - HandleNotification(msg.Notification); + if (!HandleWindowNotification(msg.Notification, wireBytes)) + HandleNotification(msg.Notification); } else if (msg.Request is not null) { @@ -660,6 +716,8 @@ private void Dispatch(JsonRpcMessage msg) private void Deliver(ulong id, JsonElement result, AhpRpcException? rpcError) { + if (_windowResponses.TryRemove(id, out var hook)) + hook.Apply(rpcError is null ? result : null, rpcError); if (_pending.TryRemove(id, out var tcs)) { if (rpcError is not null) @@ -734,6 +792,7 @@ private async Task EnqueueReplyAsync(JsonRpcMessage msg) try { await _outbound.Writer.WriteAsync(new OutboundMessage(msg)).ConfigureAwait(false); + _writerWake.Release(); } catch { /* shutting down — best effort */ } } @@ -894,6 +953,9 @@ private void FanOut(string channel, SubscriptionEvent ev) } JsonElement? paramsEl = parameters is null ? null : _serializer.SerializeToElement(parameters); + WindowOptions? requestedOptions = null; + if (parameters is SubscribeParams subscribe) _windowRequestOptions.TryGetValue(subscribe, out requestedOptions); + PrepareWindowResponse(id, method, paramsEl, requestedOptions, requestCts.Token); var req = new JsonRpcMessage { @@ -933,6 +995,7 @@ private void FanOut(string channel, SubscriptionEvent ev) } catch (Exception ex) { + AbandonWindowRequest(id, ex); // Distinguish caller cancellation (the caller's token) from a request // timeout (the configured default-timeout fired its linked token) from a // genuine error — the metric and the span status differ for each. @@ -987,12 +1050,27 @@ public async Task NotifyAsync( private async Task SendMessageAsync(JsonRpcMessage msg, CancellationToken cancellationToken) { + if (msg.Notification is { Params: { } p } notification && p.TryGetProperty("channel", out var channelEl) + && channelEl.GetString() is { } channel && GetWindowedSubscription(channel) is { } owner) + { + if (notification.Method is "channel/frame" or "channel/credit" or "channel/snapshot") + throw new InvalidOperationException("Window framing and credit are owned by the client."); + if ((owner.Limits is not null || owner.Negotiating) && WindowedDelivery.CountsAgainstWindow(notification.Method)) + { + await SendWindowAsync(owner, notification, cancellationToken).ConfigureAwait(false); + return; + } + } var sentTcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); - var item = new OutboundMessage(msg, sentTcs, cancellationToken); + var item = new OutboundMessage(msg, sentTcs, cancellationToken) + { + WindowRequest = msg.Request is { } request && _windowResponses.TryGetValue(request.Id, out var response) ? response : null, + }; try { await _outbound.Writer.WriteAsync(item, cancellationToken).ConfigureAwait(false); + _writerWake.Release(); } catch (OperationCanceledException) { @@ -1190,6 +1268,7 @@ public Subscription AttachSubscription(string uri) public async Task UnsubscribeAsync(string uri, CancellationToken cancellationToken = default) { Guard.ThrowIfNull(uri, nameof(uri)); + StopWindow(uri); List subs; lock (_subsLock) { diff --git a/clients/dotnet/src/AgentHostProtocol/Hosts/HostClientHandle.cs b/clients/dotnet/src/AgentHostProtocol/Hosts/HostClientHandle.cs index 733e50902..1f9ccfd57 100644 --- a/clients/dotnet/src/AgentHostProtocol/Hosts/HostClientHandle.cs +++ b/clients/dotnet/src/AgentHostProtocol/Hosts/HostClientHandle.cs @@ -57,6 +57,18 @@ private AhpClient CheckAlive() /// public void CheckAliveOrThrow() => CheckAlive(); + /// Returns the authoritative windowed owner on this connection, if any. + public WindowedSubscription? GetWindowedSubscription(string channel) => CheckAlive().GetWindowedSubscription(channel); + + /// Waits for channel credit and queue capacity on this connection. + public ValueTask WaitWritableAsync(string channel, CancellationToken cancellationToken = default) => + (CheckAlive().GetWindowedSubscription(channel) + ?? throw new InvalidOperationException("The channel has no windowed owner.")).WaitWritableAsync(cancellationToken); + + /// Sends a notification through the connection's flow-controlled scheduler. + public Task NotifyAsync(string method, T parameters, CancellationToken cancellationToken = default) => + CheckAlive().NotifyAsync(method, parameters, cancellationToken); + /// /// Dispatches an action through this connection on , /// refusing (throwing) if the host was removed or the connection has been diff --git a/clients/dotnet/src/AgentHostProtocol/Hosts/IMultiHostWindowedClient.cs b/clients/dotnet/src/AgentHostProtocol/Hosts/IMultiHostWindowedClient.cs new file mode 100644 index 000000000..ebd02b096 --- /dev/null +++ b/clients/dotnet/src/AgentHostProtocol/Hosts/IMultiHostWindowedClient.cs @@ -0,0 +1,24 @@ +#nullable enable + +using System.Threading; +using System.Threading.Tasks; + +namespace Microsoft.AgentHostProtocol.Hosts; + +public interface IMultiHostWindowedClient : IMultiHostClient +{ + Task<(SubscribeResult Result, WindowedSubscription Sub)> SubscribeWindowedAsync( + HostId host, string uri, WindowOptions? options = null, CancellationToken cancellationToken = default); + WindowedSubscription? GetWindowedSubscription(HostId host, string uri); + System.Threading.Channels.ChannelReader FlowControlEvents(); +} + +public static class MultiHostWindowedExtensions +{ + public static Task<(SubscribeResult Result, WindowedSubscription Sub)> SubscribeWindowedAsync( + this IMultiHostClient client, HostId host, string uri, WindowOptions? options = null, + CancellationToken cancellationToken = default) => + client is IMultiHostWindowedClient windowed + ? windowed.SubscribeWindowedAsync(host, uri, options, cancellationToken) + : throw new NotSupportedException("This runtime does not implement windowed delivery."); +} diff --git a/clients/dotnet/src/AgentHostProtocol/Hosts/MultiHostClient.cs b/clients/dotnet/src/AgentHostProtocol/Hosts/MultiHostClient.cs index 2a629c551..a3f90a47a 100644 --- a/clients/dotnet/src/AgentHostProtocol/Hosts/MultiHostClient.cs +++ b/clients/dotnet/src/AgentHostProtocol/Hosts/MultiHostClient.cs @@ -17,6 +17,33 @@ namespace Microsoft.AgentHostProtocol.Hosts; internal sealed class HostEntry : IDisposable { + internal readonly ConcurrentDictionary WindowOwners = new(StringComparer.Ordinal); + internal void RegisterWindowOwner(WindowedSubscription owner) + { + lock (_gate) + { + WindowOwners[owner.Uri] = owner; + if (!_subscriptions.Contains(owner.Uri)) _subscriptions.Add(owner.Uri); + } + } + + internal void ForgetWindowOwner(WindowedSubscription owner) + { + lock (_gate) + if (WindowOwners.TryGetValue(owner.Uri, out var current) && ReferenceEquals(current, owner)) + { + WindowOwners.TryRemove(owner.Uri, out _); + _subscriptions.Remove(owner.Uri); + } + } + internal void ApplyWindowRoot(RootState root) + { + lock (_gate) + { + _agents = root.Agents is { } agents ? new List(agents) : new List(); + _activeSessions = root.ActiveSessions; + } + } public HostId Id { get; } public HostConfig Config { get; } public string ClientId { get; set; } @@ -437,7 +464,7 @@ public bool ApplyRootAction(StateAction action) /// Multi-host registry + reconnect supervisor. Manages N independent AHP hosts, /// fans in their inbound events, and supervises reconnects per-host policy. /// -public sealed class MultiHostClient : IMultiHostClient +public sealed partial class MultiHostClient : IMultiHostWindowedClient { private readonly ConcurrentDictionary _hosts = new(StringComparer.Ordinal); private volatile IClientIdStore _store; @@ -580,6 +607,7 @@ public async Task AddHostAsync( if (!_hosts.TryAdd(ownedConfig.Id.ToString(), entry)) { entry.Dispose(); + DropWindowOwners(entry); throw new DuplicateHostException(ownedConfig.Id); } @@ -595,6 +623,7 @@ public async Task AddHostAsync( { entry.LifetimeCts.Cancel(); entry.Dispose(); + DropWindowOwners(entry); } throw; } @@ -660,6 +689,7 @@ public async Task RemoveHostAsync(HostId id, CancellationToken cancellationToken FinishPerHostListeners(id.ToString()); entry!.LifetimeCts.Cancel(); + DropWindowOwners(entry); await entry.ConnectionGate.WaitAsync(CancellationToken.None).ConfigureAwait(false); try { @@ -727,6 +757,7 @@ public async Task ShutdownAsync(CancellationToken cancellationToken = default) var entries = new List(_hosts.Values); _hosts.Clear(); + FinishFlowEvents(); // Finish per-host listener streams for every host so their consumers' // `await foreach` loops exit (mirrors the perResourceListeners finish in @@ -736,6 +767,7 @@ public async Task ShutdownAsync(CancellationToken cancellationToken = default) foreach (var entry in entries) { entry.LifetimeCts.Cancel(); + DropWindowOwners(entry); // Wake a parked (failed/disabled) supervisor so it observes the // cancellation and exits rather than blocking on its manual-reconnect // wait forever. @@ -805,6 +837,7 @@ private async Task OpenHostAsync(HostEntry entry, CancellationToken cancellation await openTask.ConfigureAwait(false); return; } + await attempt.AbortClientAsync().ConfigureAwait(false); if (openTask.IsCompleted) { await openTask.ConfigureAwait(false); @@ -838,7 +871,24 @@ private async Task OpenHostCoreAsync( transport, entry.Config.ClientConfig, null); + attempt.TrackClient(client, cancellationToken); transport = null; + client.RetainManagedWindows = true; + client.WindowHostId = entry.Id; + client.ManagedFlowEvent = BroadcastFlowEvent; + client.ManagedWindowEvent = (owner, notification, value) => ApplyManagedWindowEvent(entry, owner, notification, value); + foreach (var pair in entry.WindowOwners) + { + if (pair.Value.Closed) + entry.ForgetWindowOwner(pair.Value); + else if (pair.Key.StartsWith("ahp-tcp:", StringComparison.Ordinal)) + { + pair.Value.End(new HostNotConnectedException(entry.Id)); + entry.WindowOwners.TryRemove(pair.Key, out _); + entry.RemoveSubscription(pair.Key); + } + else if (!client.BindManagedWindow(pair.Value)) entry.ForgetWindowOwner(pair.Value); + } // Register before the first handshake request so notifications that // race initialize/reconnect are buffered rather than discarded. var stream = client.CreateEventStream(); @@ -856,9 +906,29 @@ private async Task OpenHostCoreAsync( ReconnectResult? reconnectResult = null; try { - reconnectResult = await client.ReconnectAsync( - snap.ClientId, snap.ServerSeq, subscriptions, cancellationToken) - .ConfigureAwait(false); + var items = new List(); + foreach (var pair in entry.WindowOwners) + if (pair.Value.ResumeWindowed) + items.Add(new ChannelDeliveryResumeParams + { + Channel = pair.Key, + FlowControl = new SubscriptionFlowControlOptions { Receive = pair.Value.Offer }, + LastAppliedServerSeq = pair.Value.LastAppliedServerSeq, + }); + reconnectResult = await client.ResumeAsync(new ReconnectParams + { + Channel = ProtocolVersion.RootResourceUri, + ClientId = snap.ClientId, + LastSeenServerSeq = snap.ServerSeq, + Subscriptions = new List(subscriptions), + Windows = items.Count == 0 ? null : new ChannelDeliveryResumeOptions { Items = items }, + }, cancellationToken).ConfigureAwait(false); + foreach (var pair in entry.WindowOwners) + if (pair.Value.Closed) + { + entry.WindowOwners.TryRemove(pair.Key, out _); + entry.RemoveSubscription(pair.Key); + } } catch (Exception) when (!cancellationToken.IsCancellationRequested) { @@ -881,7 +951,7 @@ private async Task OpenHostCoreAsync( () => { if (summaries is not null) entry.SeedSessionSummaries(summaries); - ApplyReconnectReplay(entry, replay); + ApplyReconnectReplay(entry, replay, snap.ServerSeq); }); installed = true; } @@ -919,9 +989,20 @@ private async Task OpenHostCoreAsync( var result = await client.InitializeAsync( entry.ClientId, entry.Config.ProtocolVersions, - subscriptions, + subscriptions is null ? null : new List(subscriptions).FindAll(uri => + !entry.WindowOwners.ContainsKey(uri) && !uri.StartsWith("ahp-tcp:", StringComparison.Ordinal)), cancellationToken) .ConfigureAwait(false); + foreach (var pair in entry.WindowOwners) + { + if (pair.Value.Closed) continue; + try { await client.SubscribeWindowedAsync(pair.Key, pair.Value.Options, cancellationToken).ConfigureAwait(false); } + catch (Exception error) when (!cancellationToken.IsCancellationRequested && !client.Completion.IsCompleted) + { + pair.Value.End(error, true); + client.DropManagedWindow(pair.Value); + } + } // Extract the root-state snapshot (agents + activeSessions) that the // server returned for the root channel, mirroring the @@ -953,7 +1034,7 @@ private async Task OpenHostCoreAsync( { if (client is not null) { - try { await client.ShutdownAsync(CancellationToken.None).ConfigureAwait(false); } catch { } + try { await client.AbortAsync().ConfigureAwait(false); } catch { } } else if (transport is not null) { @@ -1016,9 +1097,9 @@ private void CompleteOpenHost(HostEntry entry, EventStream stream) /// subscribed before the drop observe the actions missed while disconnected. /// URIs in Missing are pruned from the reconnect subscription set. /// - private ulong ApplyReconnectReplay(HostEntry entry, ReconnectReplayResult replay) + private ulong ApplyReconnectReplay(HostEntry entry, ReconnectReplayResult replay, long? previousServerSeq = null) { - var initialSeq = entry.Snapshot().ServerSeq; + var initialSeq = previousServerSeq ?? entry.Snapshot().ServerSeq; var previousSeq = initialSeq; if (replay.Actions is { } seqScan) { @@ -1032,12 +1113,13 @@ private ulong ApplyReconnectReplay(HostEntry entry, ReconnectReplayResult replay } } - var generation = entry.ApplyConnected(null, initialSeq); + var generation = entry.ApplyConnected(null, Math.Max(initialSeq, entry.Snapshot().ServerSeq)); if (replay.Actions is { } actions) { foreach (var env in actions) { + if (entry.CurrentClient?.GetWindowedSubscription(env.Channel) is not null) continue; var evt = new SubscriptionEventAction(env); ApplyEventToHostState(entry, evt); if (env.Channel == ProtocolVersion.RootResourceUri) @@ -1074,9 +1156,12 @@ private ulong ApplyReconnectSnapshot(HostEntry entry, ReconnectSnapshotResult re { if (snapshot.FromSeq > lastSeq) lastSeq = snapshot.FromSeq; } - var generation = entry.ApplyConnected(ExtractRootSnapshot(result.Snapshots), lastSeq); + var generation = entry.ApplyConnected( + entry.CurrentClient?.GetWindowedSubscription(ProtocolVersion.RootResourceUri) is null + ? ExtractRootSnapshot(result.Snapshots) : null, lastSeq); foreach (var snapshot in result.Snapshots) { + if (entry.CurrentClient?.GetWindowedSubscription(snapshot.Resource) is not null) continue; var evt = new SubscriptionEventSnapshot(snapshot); var hostEv = new HostSubscriptionEvent(entry.Id, snapshot.Resource, evt); List>? channels; @@ -1155,6 +1240,7 @@ private async Task PumpEventsAsync(HostEntry entry, EventStream stream) // normal-shutdown path. await foreach (var ev in stream.Events.ReadAllAsync(entry.LifetimeCts.Token).ConfigureAwait(false)) { + if (entry.CurrentClient?.GetWindowedSubscription(ev.Channel) is not null) continue; // Update per-host observable state BEFORE broadcasting so any // observer reading the next snapshot sees the post-event state // (mirrors the ordering in Swift HostRuntime.handleEvent). @@ -1858,6 +1944,11 @@ public async Task UnsubscribeAsync( if (ReferenceEquals(client, entry.CurrentClient)) { entry.RemoveSubscription(uri); + if (entry.WindowOwners.TryGetValue(uri, out var window)) + { + window.End(); + entry.WindowOwners.TryRemove(uri, out _); + } return; } } @@ -1893,16 +1984,44 @@ internal sealed class OpenHostAttempt private const int Active = 0; private const int Committing = 1; private const int Abandoned = 2; + private readonly object _gate = new(); private int _state; + private AhpClient? _client; - public bool TryAbandon() => - Interlocked.CompareExchange(ref _state, Abandoned, Active) == Active; + public bool TryAbandon() + { + lock (_gate) + { + if (_state != Active) return false; + _state = Abandoned; + return true; + } + } + + public Task AbortClientAsync() + { + AhpClient? client; + lock (_gate) client = _client; + return client?.AbortAsync() ?? Task.CompletedTask; + } + + public void TrackClient(AhpClient client, CancellationToken cancellationToken) + { + lock (_gate) + { + if (_state == Abandoned) throw new OperationCanceledException(cancellationToken); + _client = client; + } + } public void BeginCommit(CancellationToken cancellationToken) { - cancellationToken.ThrowIfCancellationRequested(); - if (Interlocked.CompareExchange(ref _state, Committing, Active) != Active) - throw new OperationCanceledException(cancellationToken); + lock (_gate) + { + cancellationToken.ThrowIfCancellationRequested(); + if (_state != Active) throw new OperationCanceledException(cancellationToken); + _state = Committing; + } } } } diff --git a/clients/dotnet/src/AgentHostProtocol/Hosts/MultiHostWindowed.cs b/clients/dotnet/src/AgentHostProtocol/Hosts/MultiHostWindowed.cs new file mode 100644 index 000000000..aaf84a580 --- /dev/null +++ b/clients/dotnet/src/AgentHostProtocol/Hosts/MultiHostWindowed.cs @@ -0,0 +1,101 @@ +#nullable enable + +using System.Threading; +using System.Threading.Channels; +using System.Threading.Tasks; + +namespace Microsoft.AgentHostProtocol.Hosts; + +public sealed partial class MultiHostClient +{ + private readonly object _flowGate = new(); + private readonly List> _flowReaders = new(); + + public ChannelReader FlowControlEvents() + { + var channel = Channel.CreateBounded( + new BoundedChannelOptions(256) { FullMode = BoundedChannelFullMode.DropOldest }); + lock (_flowGate) _flowReaders.Add(channel); + return channel.Reader; + } + + private void BroadcastFlowEvent(FlowControlEvent value) + { + lock (_flowGate) + foreach (var reader in _flowReaders) reader.Writer.TryWrite(value); + } + + private void FinishFlowEvents() + { + lock (_flowGate) + foreach (var reader in _flowReaders) reader.Writer.TryComplete(); + } + + public WindowedSubscription? GetWindowedSubscription(HostId host, string uri) + { + if (!_hosts.TryGetValue(host.ToString(), out var entry)) throw new UnknownHostException(host); + return entry.WindowOwners.TryGetValue(uri, out var owner) ? owner : null; + } + + public async Task<(SubscribeResult Result, WindowedSubscription Sub)> SubscribeWindowedAsync( + HostId host, string uri, WindowOptions? options = null, CancellationToken cancellationToken = default) + { + if (!_hosts.TryGetValue(host.ToString(), out var entry)) throw new UnknownHostException(host); + while (true) + { + var client = await entry.WaitForClientAsync(cancellationToken).ConfigureAwait(false); + var result = await client.SubscribeWindowedAsync(uri, options, cancellationToken).ConfigureAwait(false); + await entry.ConnectionGate.WaitAsync(cancellationToken).ConfigureAwait(false); + try + { + if (entry.LifetimeCts.IsCancellationRequested) + { + result.Sub.End(new HostShutDownException(host), true); + throw new HostShutDownException(host); + } + if (ReferenceEquals(client, entry.CurrentClient)) + { + lock (result.Sub.Gate) + if (!result.Sub.Closed) + { + result.Sub.OnClosed = entry.ForgetWindowOwner; + entry.RegisterWindowOwner(result.Sub); + } + return result; + } + } + finally { entry.ConnectionGate.Release(); } + } + } + + private void ApplyManagedWindowEvent(HostEntry entry, WindowedSubscription owner, + JsonRpcNotification notification, SubscriptionEvent value) + { + var summaryTouched = ApplyEventToHostState(entry, value); + var rootTouched = value is SubscriptionEventAction rootAction + && owner.Uri == ProtocolVersion.RootResourceUri && entry.ApplyRootAction(rootAction.Envelope.Action); + if (value is SubscriptionEventAction action) entry.AdvanceServerSeq(action.Envelope.ServerSeq); + if (value is SubscriptionEventSnapshot snapshot) + { + entry.AdvanceServerSeq(snapshot.Snapshot.FromSeq); + if (owner.Uri == ProtocolVersion.RootResourceUri && snapshot.Snapshot.State.Root is { } root) + { + entry.ApplyWindowRoot(root); + rootTouched = true; + } + } + var ev = new HostSubscriptionEvent(entry.Id, owner.Uri, value); + List> readers; + lock (_subsLock) readers = new List>(_subChannels); + foreach (var reader in readers) reader.Writer.TryWrite(ev); + BroadcastPerResourceEvent(entry.Id, owner.Uri, value); + if (summaryTouched || rootTouched) NotifyPerHostSnapshot(entry); + if (summaryTouched) NotifyPerHostSummaries(entry); + } + + private static void DropWindowOwners(HostEntry entry) + { + foreach (var pair in entry.WindowOwners) pair.Value.End(new HostShutDownException(entry.Id)); + entry.WindowOwners.Clear(); + } +} diff --git a/clients/dotnet/src/AgentHostProtocol/Telemetry.cs b/clients/dotnet/src/AgentHostProtocol/Telemetry.cs index 2ff70674a..bb4dc8a1e 100644 --- a/clients/dotnet/src/AgentHostProtocol/Telemetry.cs +++ b/clients/dotnet/src/AgentHostProtocol/Telemetry.cs @@ -11,6 +11,7 @@ using System.Diagnostics; using System.Diagnostics.Metrics; using System.Reflection; +using System.Collections.Concurrent; namespace Microsoft.AgentHostProtocol; @@ -74,4 +75,51 @@ public static class AhpTelemetry internal static readonly Counter MalformedFrames = Meter.CreateCounter(AhpTelemetryNames.FramesMalformed, unit: AhpTelemetryNames.FramesMalformedUnit, description: AhpTelemetryNames.FramesMalformedDescription); + + internal static readonly Counter WindowDataBytesSent = Meter.CreateCounter("ahp.client.window.data.sent", "By"); + internal static readonly Counter WindowDataBytesReceived = Meter.CreateCounter("ahp.client.window.data.received", "By"); + internal static readonly Counter WindowCreditBytesSent = Meter.CreateCounter("ahp.client.window.credit.sent", "By"); + internal static readonly Counter WindowCreditBytesReceived = Meter.CreateCounter("ahp.client.window.credit.received", "By"); + internal static readonly Histogram WindowBlockedDuration = Meter.CreateHistogram("ahp.client.window.blocked.duration", "ms"); + private static readonly ConcurrentDictionary> Windows = new(); + private static long _nextWindow; + private static readonly ObservableGauge BufferedWindowBytes = + Meter.CreateObservableGauge("ahp.client.window.buffered", ObserveWindowBytes, "By"); + + internal static long TrackWindow(WindowedSubscription subscription) + { + foreach (var pair in Windows) + if (!pair.Value.TryGetTarget(out _)) Windows.TryRemove(pair.Key, out _); + var key = System.Threading.Interlocked.Increment(ref _nextWindow); + Windows.TryAdd(key, new WeakReference(subscription)); + return key; + } + + internal static void UntrackWindow(long key) => Windows.TryRemove(key, out _); + internal static int TrackedWindowCount => Windows.Count; + + private static IEnumerable> ObserveWindowBytes() + { + long send = 0, receive = 0; + foreach (var pair in Windows) + { + if (!pair.Value.TryGetTarget(out var owner) || owner.Closed) + { + Windows.TryRemove(pair.Key, out _); + continue; + } + var status = owner.FlowControl; + send += status.BufferedSendBytes; + receive += status.BufferedReceiveBytes; + } + yield return new Measurement(send, new KeyValuePair("direction", "send")); + yield return new Measurement(receive, new KeyValuePair("direction", "receive")); + } + + internal static void FlowEvent(FlowControlEvent value) + { + using var activity = Activity.Current is null ? ActivitySource.StartActivity("ahp.flowcontrol") : null; + (Activity.Current ?? activity)?.AddEvent(new ActivityEvent("ahp.flowcontrol." + value.Kind, + tags: new ActivityTagsCollection { { "ahp.flowcontrol.bytes", value.Bytes } })); + } } diff --git a/clients/dotnet/src/AgentHostProtocol/WebSocketTransport.cs b/clients/dotnet/src/AgentHostProtocol/WebSocketTransport.cs index 64f6474ba..41ac00b4e 100644 --- a/clients/dotnet/src/AgentHostProtocol/WebSocketTransport.cs +++ b/clients/dotnet/src/AgentHostProtocol/WebSocketTransport.cs @@ -38,7 +38,7 @@ public sealed class WebSocketTransportOptions /// or to wrap /// an existing connection. /// -public sealed class WebSocketTransport : ITransport +public sealed class WebSocketTransport : ITransport, ITransportMessageLimits { private readonly WebSocket _ws; private readonly SemaphoreSlim _sendLock = new(1, 1); @@ -46,6 +46,8 @@ public sealed class WebSocketTransport : ITransport private int _disposed; internal long MaxMessageBytes => _maxMessageBytes; + /// Complete inbound WebSocket message cap, not the receive-buffer size. + public long? MaximumInboundMessageBytes => _maxMessageBytes > 0 ? _maxMessageBytes : null; // Receive buffer: 64 KiB initial, grows as needed. private byte[] _receiveBuffer = new byte[64 * 1024]; diff --git a/clients/dotnet/src/AgentHostProtocol/WindowOptions.cs b/clients/dotnet/src/AgentHostProtocol/WindowOptions.cs new file mode 100644 index 000000000..c850ea7f3 --- /dev/null +++ b/clients/dotnet/src/AgentHostProtocol/WindowOptions.cs @@ -0,0 +1,135 @@ +#nullable enable + +using System.Threading; +using System.Threading.Tasks; + +namespace Microsoft.AgentHostProtocol; + +/// Connection defaults; configuring them does not enable windowing. +public sealed class WindowingConfig +{ + public long WindowBytes { get; set; } = 262144; + /// UTF-16 bytes; null derives a safe limit from the transport. + public long? MaximumFrameBytes { get; set; } + public long MaxBufferedSendBytes { get; set; } = 1048576; + public IList Coalescers { get; } = new List(); + + internal WindowingConfig Snapshot() + { + Validate(WindowBytes, 1, nameof(WindowBytes)); + if (MaximumFrameBytes is { } maximum) Validate(maximum, 1024, nameof(MaximumFrameBytes)); + Validate(MaxBufferedSendBytes, 1, nameof(MaxBufferedSendBytes)); + var copy = new WindowingConfig + { + WindowBytes = WindowBytes, + MaximumFrameBytes = MaximumFrameBytes, + MaxBufferedSendBytes = MaxBufferedSendBytes, + }; + foreach (var coalescer in Coalescers) + copy.Coalescers.Add(coalescer ?? throw new ArgumentException("A coalescer cannot be null.")); + return copy; + } + + internal static void Validate(long value, long minimum, string name) + { + if (value < minimum || value > 9007199254740991) + throw new ArgumentOutOfRangeException(name); + } +} + +public enum ReceiveMode { Auto, Manual } + +/// Per-subscription opt-in; unspecified limits inherit connection defaults. +public sealed class WindowOptions +{ + public long? WindowBytes { get; set; } + public long? MaximumFrameBytes { get; set; } + public ReceiveMode Receive { get; set; } = ReceiveMode.Auto; + public bool Coalescing { get; set; } = true; + + internal WindowOptions Snapshot() => new() + { + WindowBytes = WindowBytes, + MaximumFrameBytes = MaximumFrameBytes, + Receive = Receive, + Coalescing = Coalescing, + }; +} + +/// Additive client capability; existing IAhpClient implementations need not implement it. +public interface IWindowedAhpClient : IAhpClient +{ + Task<(SubscribeResult Result, WindowedSubscription Sub)> SubscribeWindowedAsync( + string uri, WindowOptions? options = null, CancellationToken cancellationToken = default); + WindowedSubscription? GetWindowedSubscription(string uri); + Task ResumeAsync(ReconnectParams parameters, CancellationToken cancellationToken = default); + System.Threading.Channels.ChannelReader FlowControlEvents(); +} + +public static class WindowedClientExtensions +{ + private static IWindowedAhpClient Capability(IAhpClient client) => client as IWindowedAhpClient + ?? throw new NotSupportedException("This client does not implement windowed delivery."); + + public static Task<(SubscribeResult Result, WindowedSubscription Sub)> SubscribeWindowedAsync( + this IAhpClient client, string uri, WindowOptions? options = null, + CancellationToken cancellationToken = default) => + client is IWindowedAhpClient windowed + ? windowed.SubscribeWindowedAsync(uri, options, cancellationToken) + : throw new NotSupportedException("This client does not implement windowed delivery."); + + public static WindowedSubscription? GetWindowedSubscription(this IAhpClient client, string uri) => + Capability(client).GetWindowedSubscription(uri); + + public static Task ResumeAsync(this IAhpClient client, ReconnectParams parameters, + CancellationToken cancellationToken = default) => Capability(client).ResumeAsync(parameters, cancellationToken); + + public static System.Threading.Channels.ChannelReader FlowControlEvents(this IAhpClient client) => + Capability(client).FlowControlEvents(); +} + +public enum CoalescingDirection { Incoming, Outgoing } + +public sealed class CoalescingContext +{ + public string Channel { get; } + public CoalescingDirection Direction { get; } + internal CoalescingContext(string channel, CoalescingDirection direction) + { + Channel = channel; + Direction = direction; + } +} + +/// Combines adjacent queued notifications; null means they cannot be combined. +public interface IChannelMessageCoalescer +{ + JsonRpcNotification? Coalesce(JsonRpcNotification earlier, JsonRpcNotification later, CoalescingContext context); +} + +public sealed class FlowControlStatus +{ + public string Mode { get; internal init; } = "ordinary"; + public string? FallbackReason { get; internal init; } + public ChannelFlowControl? Limits { get; internal init; } + public long OutstandingSendBytes { get; internal init; } + public long BufferedSendBytes { get; internal init; } + public long BufferedReceiveBytes { get; internal init; } + public bool Writable { get; internal init; } +} + +public sealed class FlowControlEvent +{ + public string Channel { get; } + public string Kind { get; } + public long Bytes { get; } + public FlowControlStatus Status { get; } + public Hosts.HostId? HostId { get; internal init; } + internal FlowControlEvent(string channel, string kind, FlowControlStatus status, long bytes = 0) + { + Channel = channel; + Kind = kind; + Status = status; + Bytes = bytes; + } +} diff --git a/clients/dotnet/src/AgentHostProtocol/WindowedDeliveryEngine.cs b/clients/dotnet/src/AgentHostProtocol/WindowedDeliveryEngine.cs new file mode 100644 index 000000000..a00664c59 --- /dev/null +++ b/clients/dotnet/src/AgentHostProtocol/WindowedDeliveryEngine.cs @@ -0,0 +1,981 @@ +#nullable enable + +using System.Collections.Concurrent; +using System.Linq; +using System.Text; +using System.Threading; +using System.Threading.Channels; +using System.Threading.Tasks; + +namespace Microsoft.AgentHostProtocol; + +public sealed partial class AhpClient +{ + private sealed class WindowResponse + { + internal readonly IReadOnlyList Owners; + internal Action Apply = null!; + internal readonly Action Cleanup; + private int _phase; + internal bool Abandoned; + internal Exception? Failure; + + internal WindowResponse(IReadOnlyList owners, Action cleanup) + { + Owners = owners; + Cleanup = cleanup; + } + + internal bool TryStart() => Interlocked.CompareExchange(ref _phase, 1, 0) == 0; + internal bool Abandon(Exception error) + { + Failure = error; + Abandoned = true; + var unsent = Interlocked.CompareExchange(ref _phase, 2, 0) == 0; + Cleanup(error); + return unsent; + } + } + + private readonly object _windowGate = new(); + private readonly Dictionary _windows = new(StringComparer.Ordinal); + private readonly ConcurrentDictionary _windowResponses = new(); + private readonly LinkedList _closedWindowOrder = new(); + private readonly Dictionary> _closedWindows = new(StringComparer.Ordinal); + private readonly System.Runtime.CompilerServices.ConditionalWeakTable _windowRequestOptions = new(); + private readonly List> _flowEventReaders = new(); + private readonly WindowWakeSignal _writerWake = new(); + private int _windowTurn; + private WindowSend? _windowInFlight; + private bool RequiresWindowSizes + { + get { lock (_windowGate) return _windows.Values.Any(owner => owner.Limits is not null); } + } + internal WindowingConfig WindowConfig => _cfg.Windowing; + internal JsonElement WindowParameters(T parameters) => _serializer.SerializeToElement(parameters); + internal Action? ManagedWindowEvent { get; set; } + internal bool RetainManagedWindows { get; set; } + internal Hosts.HostId? WindowHostId { get; set; } + internal Action? ManagedFlowEvent { get; set; } + + public WindowedSubscription? GetWindowedSubscription(string uri) + { + lock (_windowGate) return _windows.TryGetValue(uri, out var owner) ? owner : null; + } + + public ChannelReader FlowControlEvents() + { + var channel = Channel.CreateBounded( + new BoundedChannelOptions(256) { FullMode = BoundedChannelFullMode.DropOldest }); + lock (_windowGate) _flowEventReaders.Add(channel); + return channel.Reader; + } + + internal void WindowEvent(WindowedSubscription owner, string kind, long bytes = 0) + { + var value = new FlowControlEvent(owner.Uri, kind, owner.FlowControl, bytes) { HostId = owner.HostId }; + lock (_windowGate) + foreach (var reader in _flowEventReaders) reader.Writer.TryWrite(value); + AhpTelemetry.FlowEvent(value); + ManagedFlowEvent?.Invoke(value); + } + + private ChannelReceiveLimits WindowOffer(WindowOptions options) + { + if (options.Receive is not ReceiveMode.Auto and not ReceiveMode.Manual) + throw new ArgumentOutOfRangeException(nameof(options)); + var transportDefault = _transport is ITransportMessageLimits { MaximumInboundMessageBytes: > 0 } limited + ? 2 * (limited.MaximumInboundMessageBytes!.Value / 3) + : 1048576; + var result = new ChannelReceiveLimits + { + WindowBytes = options.WindowBytes ?? WindowConfig.WindowBytes, + MaximumFrameBytes = options.MaximumFrameBytes ?? WindowConfig.MaximumFrameBytes ?? transportDefault, + }; + WindowingConfig.Validate(result.WindowBytes, 1, nameof(options.WindowBytes)); + WindowingConfig.Validate(result.MaximumFrameBytes, 1024, nameof(options.MaximumFrameBytes)); + return result; + } + + public async Task<(SubscribeResult Result, WindowedSubscription Sub)> SubscribeWindowedAsync( + string uri, WindowOptions? options = null, CancellationToken cancellationToken = default) + { + Guard.ThrowIfNull(uri, nameof(uri)); + cancellationToken.ThrowIfCancellationRequested(); + var copy = (options ?? new WindowOptions()).Snapshot(); + var offer = WindowOffer(copy); + WindowedSubscription owner; + lock (_windowGate) + { + if (_windows.TryGetValue(uri, out var closing) && closing.Closed) + throw new InvalidOperationException("The subscription is closing."); + if (!_windows.TryGetValue(uri, out owner!)) + { + owner = new WindowedSubscription(this, uri, copy, offer); + _windows.Add(uri, owner); + } + } + var request = new SubscribeParams + { + Channel = uri, + FlowControl = new SubscriptionFlowControlOptions { Receive = offer }, + }; + _windowRequestOptions.Add(request, copy); + SubscribeResult result; + try + { + result = await RequestAsync("subscribe", request, cancellationToken).ConfigureAwait(false) + ?? throw Protocol("subscribe returned no result"); + } + catch (Exception error) + { + CleanupUnacceptedWindow(owner, error); + throw; + } + finally { _windowRequestOptions.Remove(request); } + return (result, owner); + } + + public async Task ResumeAsync(ReconnectParams parameters, CancellationToken cancellationToken = default) + { + try + { + var result = await RequestAsync("reconnect", parameters, cancellationToken).ConfigureAwait(false) + ?? throw Protocol("reconnect returned no result"); + AhpTelemetry.Reconnects.Add(1, new KeyValuePair(AhpTelemetryNames.AttrOutcome, AhpTelemetryNames.OutcomeOk)); + return result; + } + catch + { + AhpTelemetry.Reconnects.Add(1, new KeyValuePair(AhpTelemetryNames.AttrOutcome, AhpTelemetryNames.OutcomeError)); + throw; + } + } + + internal bool BindManagedWindow(WindowedSubscription owner) + { + lock (owner.Gate) + { + if (owner.Closed) return false; + owner.Client = this; + owner.HostId = WindowHostId; + owner.Epoch++; + owner.Closed = false; + owner.CloseError = null; + owner.Negotiating = true; + owner.Limits = null; + owner.Received = owner.Released = owner.CreditSent = owner.Sent = owner.PeerConsumed = 0; + owner.CreditTimer = false; + owner.BufferedReceive -= owner.FragmentBytes; + owner.Fragments = null; + owner.FragmentBytes = 0; + owner.Offer = WindowOffer(owner.Options); + } + lock (_windowGate) _windows[owner.Uri] = owner; + return true; + } + + private void PrepareWindowResponse(ulong id, string method, JsonElement? parameters, WindowOptions? requestedOptions, + CancellationToken cancellationToken) + { + if (parameters is not { } p) return; + if (method == "subscribe" && p.TryGetProperty("channel", out var channelEl) && channelEl.GetString() is { } uri) + { + var owner = GetWindowedSubscription(uri); + ChannelReceiveLimits? offer = null; + if (p.TryGetProperty("flowControl", out var flow)) + { + offer = _serializer.Deserialize(flow).Receive; + WindowingConfig.Validate(offer.WindowBytes, 1, "windowBytes"); + WindowingConfig.Validate(offer.MaximumFrameBytes, 1024, "maximumFrameBytes"); + if (owner is null) + { + owner = new WindowedSubscription(this, uri, new WindowOptions(), offer); + lock (_windowGate) _windows[uri] = owner; + } + } + if (owner is null) return; + var target = owner; + var pendingOptions = requestedOptions; + var initial = owner.Negotiating; + var response = new WindowResponse(new[] { target }, error => CleanupUnacceptedWindow(target, error)); + response.Apply = (result, error) => + { + if (result is null) + { + response.Cleanup(error ?? new AhpClientClosedException("Subscribe failed.")); + return; + } + if (target.Closed || !ReferenceEquals(GetWindowedSubscription(uri), target)) + { + if (GetWindowedSubscription(uri) is null) _ = UnsubscribeOwnerAsync(target); + return; + } + var decoded = _serializer.Deserialize(result.Value); + if (decoded.FlowControl is not null && decoded.Snapshot is not null) + throw Protocol("windowed subscribe returned a snapshot"); + InstallWindow(target, decoded.FlowControl, offer, true); + lock (target.Gate) + { + if (pendingOptions is not null) target.Options = pendingOptions; + } + if (decoded.Snapshot is { } snapshot) + ReceiveWindowMessage(target, SnapshotNotification(snapshot), 0, 0); + if (initial && (!_pending.ContainsKey(id) || cancellationToken.IsCancellationRequested)) + CloseWindowOwner(target); + }; + _windowResponses[id] = response; + } + else if (method == "reconnect") + { + var windows = p.TryGetProperty("windows", out var windowsEl) + ? _serializer.Deserialize(windowsEl) : null; + var requested = new Dictionary(StringComparer.Ordinal); + foreach (var item in windows?.Items ?? new List()) + { + if (item.Channel.StartsWith("ahp-tcp:", StringComparison.Ordinal)) + throw Protocol("TCP subscriptions cannot reconnect"); + if (!p.GetProperty("subscriptions").EnumerateArray().Any(x => x.GetString() == item.Channel) + || requested.ContainsKey(item.Channel)) + throw Protocol("invalid reconnect window subscriptions"); + var owner = GetWindowedSubscription(item.Channel); + if (owner is null) + { + owner = new WindowedSubscription(this, item.Channel, new WindowOptions(), item.FlowControl.Receive); + lock (_windowGate) _windows[item.Channel] = owner; + } + requested.Add(item.Channel, (owner, item.FlowControl.Receive)); + } + List ordinary; + lock (_windowGate) ordinary = _windows.Values.Where(x => !x.ResumeWindowed && x.EverAccepted && !x.Closed).ToList(); + if (requested.Count == 0 && ordinary.Count == 0) return; + var targets = requested.Values.Select(x => x.Owner).Concat(ordinary).ToList(); + var response = new WindowResponse(targets, error => + { + foreach (var target in targets) CleanupUnacceptedWindow(target, error); + }); + response.Apply = (result, error) => + { + if (result is null) + { + response.Cleanup(error ?? new AhpClientClosedException("Reconnect failed.")); + return; + } + var retained = new HashSet(StringComparer.Ordinal); + if (result.Value.TryGetProperty("windows", out var returned)) + { + foreach (var item in _serializer.Deserialize(returned).Items) + { + if (!requested.TryGetValue(item.Channel, out var state) || !retained.Add(item.Channel)) + throw Protocol("unsolicited reconnect window"); + if (state.Owner.Closed || !ReferenceEquals(GetWindowedSubscription(item.Channel), state.Owner)) + { + if (GetWindowedSubscription(item.Channel) is null) _ = UnsubscribeOwnerAsync(state.Owner); + continue; + } + InstallWindow(state.Owner, item.FlowControl, state.Offer, item.Recovery == ChannelRecoveryKind.Snapshot); + WindowEvent(state.Owner, "resumed"); + } + } + foreach (var item in requested) + if (!retained.Contains(item.Key)) + { + item.Value.Owner.End(Protocol("subscription was not retained"), true); + WindowEvent(item.Value.Owner, "dropped"); + RemoveWindow(item.Value.Owner); + } + RecoverOrdinaryWindows(_serializer.Deserialize(result.Value)); + }; + _windowResponses[id] = response; + } + } + + private void InstallWindow(WindowedSubscription owner, ChannelFlowControl? limits, ChannelReceiveLimits? offer, bool snapshot) + { + if (limits is not null) + { + if (offer is null) throw Protocol("unsolicited window acceptance"); + ValidateLimits(limits.ClientReceive); + ValidateLimits(limits.HostReceive); + if (limits.ClientReceive.WindowBytes > offer.WindowBytes || limits.ClientReceive.MaximumFrameBytes > offer.MaximumFrameBytes) + throw Protocol("accepted limits exceed the offer"); + } + else if (owner.Uri.StartsWith("ahp-tcp:", StringComparison.Ordinal)) + throw Protocol("TCP requires windowed delivery"); + lock (owner.Gate) + { + if (owner.Limits is null || limits is null) + { + owner.Epoch++; + owner.Sent = owner.PeerConsumed = owner.Received = owner.Released = owner.CreditSent = 0; + owner.Fragments = null; + owner.FragmentBytes = 0; + } + ForgetClosedWindow(owner.Uri); + owner.Offer = offer ?? owner.Offer; + owner.Limits = limits; + owner.ResumeWindowed = limits is not null; + owner.Fallback = limits is null && offer is not null ? "declined" : null; + owner.Negotiating = false; + owner.EverAccepted = true; + owner.AwaitingSnapshot = limits is not null && snapshot; + } + owner.Publish(); + WindowEvent(owner, limits is null && offer is not null ? "declined" : "negotiated"); + _writerWake.Release(); + } + + private static void ValidateLimits(ChannelReceiveLimits limits) + { + if (limits.WindowBytes < 1 || limits.MaximumFrameBytes < 1024 + || limits.WindowBytes > 9007199254740991 || limits.MaximumFrameBytes > 9007199254740991) + throw Protocol("invalid window limits"); + } + + internal static long RawWindowBytes(TransportMessage message) => + checked((long)(message.Text ?? new UTF8Encoding(false, true).GetString(message.Binary.ToArray())).Length * 2); + private static AhpTransportException Protocol(string message) => new("protocol", "ahp: " + message); + + private JsonRpcNotification SnapshotNotification(Snapshot snapshot) => new() + { + Method = "channel/snapshot", + Params = _serializer.SerializeToElement(new ChannelSnapshotParams { Channel = snapshot.Resource, Snapshot = snapshot }), + }; + + private bool HandleWindowNotification(JsonRpcNotification notification, long bytes) + { + if (notification.Params is not { } parameters + || !parameters.TryGetProperty("channel", out var channel) + || channel.GetString() is not { } uri) return false; + var owner = GetWindowedSubscription(uri); + if (owner is null) + { + if (IsClosedWindow(uri)) return true; + if (notification.Method is "channel/frame" or "channel/snapshot" or "channel/credit") + throw Protocol("window message without a subscription"); + return false; + } + if (owner.Closed) return true; + if (owner.Negotiating && notification.Method is not "channel/credit") + { + if (IsClosedWindow(uri)) return true; + if (notification.Method is "channel/frame" or "channel/snapshot" || owner.ResumeWindowed) + throw Protocol("data before negotiation completed"); + return false; + } + if (notification.Method == "channel/credit") + { + var consumed = _serializer.Deserialize(parameters).ConsumedBytes; + lock (owner.Gate) + { + if (owner.Limits is null) throw Protocol("credit on an ordinary subscription"); + if (consumed < 0 || consumed > owner.Sent) throw Protocol("credit exceeds sent bytes"); + if (consumed > owner.PeerConsumed) owner.PeerConsumed = consumed; + } + AhpTelemetry.WindowCreditBytesReceived.Add(bytes); + WindowEvent(owner, "creditReceived", bytes); + owner.Publish(); + _writerWake.Release(); + return true; + } + var counted = owner.Limits is not null && WindowedDelivery.CountsAgainstWindow(notification.Method); + if (counted) + { + if (bytes > owner.Limits!.ClientReceive.MaximumFrameBytes) throw Protocol("data exceeds maximumFrameBytes"); + lock (owner.Gate) + { + owner.Received = checked(owner.Received + bytes); + owner.BufferedReceive = checked(owner.BufferedReceive + bytes); + } + AhpTelemetry.WindowDataBytesReceived.Add(bytes); + } + if (notification.Method == "channel/frame") + { + if (!counted) throw Protocol("fragment on an ordinary subscription"); + var fragment = _serializer.Deserialize(parameters); + if (fragment.Data.Length == 0 || char.IsLowSurrogate(fragment.Data[0]) || char.IsHighSurrogate(fragment.Data[fragment.Data.Length - 1])) + throw Protocol("invalid fragment boundary"); + string? joined = null; + long total; + lock (owner.Gate) + { + owner.Fragments ??= new StringBuilder(); + owner.Fragments.Append(fragment.Data); + owner.FragmentBytes = checked(owner.FragmentBytes + bytes); + total = owner.FragmentBytes; + if (fragment.Final == true) + { + joined = owner.Fragments.ToString(); + owner.Fragments = null; + owner.FragmentBytes = 0; + } + } + if (joined is null) + ReleaseWindowBytes(owner, bytes, owner.Epoch, true); + else + { + WindowEvent(owner, "fragmentedMessageReceived", total); + var message = _serializer.DecodeMessage(TransportMessage.FromText(joined)); + var inner = message.Notification ?? throw Protocol("fragment did not contain a notification"); + if (inner.Method == "channel/frame" || !WindowedDelivery.CountsAgainstWindow(inner.Method) + || inner.Params?.GetProperty("channel").GetString() != uri) + throw Protocol("invalid joined fragment channel or method"); + ReceiveWindowMessage(owner, inner, total, bytes); + } + return true; + } + if (counted && owner.Fragments is not null) throw Protocol("data interleaved with fragments"); + if (notification.Method == "tcp/reset") + { + lock (owner.Gate) + { + owner.Fragments = null; + foreach (var send in owner.Sends) send.Fail(Protocol("TCP reset")); + owner.Sends.Clear(); + owner.BufferedSend = 0; + foreach (var delivery in owner.Queue) owner.BufferedReceive -= delivery.WireBytes; + owner.Queue.Clear(); + owner.BufferedReceive -= owner.FragmentBytes; + owner.FragmentBytes = 0; + } + } + if (notification.Method == "tcp/data") + { + if (owner.ReceiveEof) throw Protocol("TCP data after EOF"); + ValidateTcpData(parameters.GetProperty("data").GetString()!); + } + if (notification.Method == "tcp/eof") owner.ReceiveEof = true; + ReceiveWindowMessage(owner, notification, counted ? bytes : 0, counted ? bytes : 0); + if (notification.Method == "tcp/reset") + { + owner.End(); + _ = UnsubscribeOwnerAsync(owner); + } + return true; + } + + internal void RecoverOrdinaryWindows(ReconnectResult result) + { + List owners; + lock (_windowGate) owners = _windows.Values.Where(x => !x.ResumeWindowed && !x.Closed).ToList(); + foreach (var owner in owners) + { + owner.Negotiating = false; + owner.Publish(); + if (result.Value is ReconnectReplayResult replay) + foreach (var action in replay.Actions.Where(x => x.Channel == owner.Uri)) + ReceiveWindowMessage(owner, new JsonRpcNotification + { + Method = "action", + Params = _serializer.SerializeToElement(action), + }, 0, 0); + if (result.Value is ReconnectSnapshotResult snapshot) + foreach (var state in snapshot.Snapshots.Where(x => x.Resource == owner.Uri)) + ReceiveWindowMessage(owner, SnapshotNotification(state), 0, 0); + } + } + + private SubscriptionEvent WindowEventValue(JsonRpcNotification notification) + { + var p = notification.Params ?? throw Protocol("notification has no params"); + return notification.Method switch + { + "action" => new SubscriptionEventAction(_serializer.Deserialize(p)), + "channel/snapshot" => new SubscriptionEventSnapshot(_serializer.Deserialize(p).Snapshot), + "tcp/data" => new SubscriptionEventTcpData(_serializer.Deserialize(p)), + "tcp/eof" => new SubscriptionEventTcpEof(_serializer.Deserialize(p)), + "tcp/reset" => new SubscriptionEventTcpReset(_serializer.Deserialize(p)), + "root/sessionAdded" => new SubscriptionEventSessionAdded(_serializer.Deserialize(p)), + "root/sessionRemoved" => new SubscriptionEventSessionRemoved(_serializer.Deserialize(p)), + "root/sessionSummaryChanged" => new SubscriptionEventSessionSummaryChanged(_serializer.Deserialize(p)), + "root/progress" => new SubscriptionEventProgress(_serializer.Deserialize(p)), + "auth/required" => new SubscriptionEventAuthRequired(_serializer.Deserialize(p)), + _ => new SubscriptionEventNotification(notification), + }; + } + + [System.Diagnostics.CodeAnalysis.SuppressMessage("Reliability", "CA2000:Dispose objects before losing scope", + Justification = "The reliable queue transfers delivery ownership to its consumer; coalescing replaces leases without releasing their credit.")] + private void ReceiveWindowMessage(WindowedSubscription owner, JsonRpcNotification notification, long bytes, long credit, bool observe = true) + { + if (owner.AwaitingSnapshot && WindowedDelivery.CountsAgainstWindow(notification.Method)) + { + if (notification.Method != "channel/snapshot") throw Protocol("first data is not a snapshot"); + owner.AwaitingSnapshot = false; + } + + if (notification.Method == "channel/snapshot" + && notification.Params!.Value.GetProperty("snapshot").GetProperty("resource").GetString() != owner.Uri) + throw Protocol("snapshot resource mismatch"); + var value = WindowEventValue(notification); + var epoch = owner.Epoch; + var delivery = new ReceivedDelivery(value, notification, bytes, credit, epoch, + () => owner.Client.ReleaseDelivery(owner, bytes, credit, epoch)); + lock (owner.Gate) + { + if (value is SubscriptionEventAction action) owner.LastSeq = Math.Max(owner.LastSeq, action.Envelope.ServerSeq); + if (value is SubscriptionEventSnapshot snapshot) owner.LastSeq = Math.Max(owner.LastSeq, snapshot.Snapshot.FromSeq); + if (owner.Limits is not null && owner.Options.Coalescing && owner.Queue.Last is { } tail && tail.Value.Epoch == epoch) + { + var merged = Coalesce(tail.Value.Notification, notification, owner, CoalescingDirection.Incoming); + if (merged is not null) + { + var total = checked(bytes + tail.Value.WireBytes); + var owed = checked(credit + tail.Value.CreditBytes); + delivery = new ReceivedDelivery(WindowEventValue(merged), merged, total, owed, epoch, + () => owner.Client.ReleaseDelivery(owner, total, owed, epoch)); + owner.Queue.RemoveLast(); + WindowEvent(owner, "coalesced", total); + } + } + if (owner.Limits is null && owner.Queue.Count >= _cfg.SubscriptionBufferCapacity) + owner.Queue.RemoveFirst(); + owner.Queue.AddLast(delivery); + owner.Available.TrySetResult(true); + owner.Available = WindowedSubscription.Signal(); + } + if (observe) + { + FanOut(owner.Uri, value); + ManagedWindowEvent?.Invoke(owner, notification, value); + } + owner.Publish(); + } + + private void ReleaseWindowBytes(WindowedSubscription owner, long bytes, long epoch, bool immediate = false) + { + if (bytes == 0) return; + bool flush; + lock (owner.Gate) + { + if (owner.Closed || owner.Epoch != epoch || owner.Limits is null || owner.Negotiating) return; + owner.Released = checked(owner.Released + bytes); + var threshold = Math.Max(1, Math.Min(16384, owner.Limits.ClientReceive.WindowBytes / 4)); + flush = immediate || owner.Released - owner.CreditSent >= threshold + || owner.Received - owner.CreditSent >= owner.Limits.ClientReceive.WindowBytes; + if (!flush && !owner.CreditTimer) + { + owner.CreditTimer = true; + _ = FlushCreditLaterAsync(owner, epoch); + } + + } + if (flush) _writerWake.Release(); + owner.Publish(); + } + + private void ReleaseDelivery(WindowedSubscription owner, long bytes, long credit, long epoch) + { + lock (owner.Gate) owner.BufferedReceive -= bytes; + ReleaseWindowBytes(owner, credit, epoch); + owner.Publish(); + } + + private async Task FlushCreditLaterAsync(WindowedSubscription owner, long epoch) + { + using var cancellation = CancellationTokenSource.CreateLinkedTokenSource(_lifetimeCts.Token); + lock (owner.Gate) + { + if (owner.Closed || owner.Epoch != epoch) return; + owner.CreditCancellation = cancellation; + } + try + { + await TimeProviderCompatibility.DelayAsync(_cfg.TimeProvider, TimeSpan.FromMilliseconds(5), cancellation.Token).ConfigureAwait(false); + _writerWake.Release(); + } + catch (OperationCanceledException) when (cancellation.IsCancellationRequested) { } + finally + { + lock (owner.Gate) + if (ReferenceEquals(owner.CreditCancellation, cancellation)) + { + owner.CreditCancellation = null; + owner.CreditTimer = false; + } + } + } + + private string EncodeWindow(JsonRpcNotification notification) => + _serializer.EncodeMessage(new JsonRpcMessage { Notification = notification }).Text + ?? throw Protocol("serializer must encode text"); + + private JsonRpcNotification? Coalesce(JsonRpcNotification older, JsonRpcNotification newer, + WindowedSubscription owner, CoalescingDirection direction) + { + if (older.Method != newer.Method || !WindowedDelivery.CountsAgainstWindow(older.Method) + || older.Method is "channel/snapshot" or "dispatchAction" or "tcp/eof" + || older.Params?.TryGetProperty("origin", out _) == true || newer.Params?.TryGetProperty("origin", out _) == true) + return null; + foreach (var coalescer in WindowConfig.Coalescers) + { + var combined = coalescer.Coalesce(older, newer, new CoalescingContext(owner.Uri, direction)); + if (combined is null) continue; + if (combined.Method != older.Method || combined.Params?.GetProperty("channel").GetString() != owner.Uri) + throw Protocol("custom coalescer changed method or channel"); + return combined; + } + if (direction == CoalescingDirection.Incoming && older.Method == "action") + { + var first = _serializer.Deserialize(older.Params!.Value); + var last = _serializer.Deserialize(newer.Params!.Value); + if (first.RejectionReason is not null || last.RejectionReason is not null) return null; + StateAction? action = null; + if (first.Action.Value is TerminalDataAction a && last.Action.Value is TerminalDataAction b) + action = new StateAction(b with { Data = a.Data + b.Data }); + if (first.Action.Value is ChatDeltaAction c && last.Action.Value is ChatDeltaAction d + && c.TurnId == d.TurnId && c.PartId == d.PartId + && MetadataEqual(older.Params.Value.GetProperty("action"), newer.Params.Value.GetProperty("action"))) + action = new StateAction(d with { Content = c.Content + d.Content }); + return action is null ? null : new JsonRpcNotification + { + Method = "action", + Params = _serializer.SerializeToElement(last with { Action = action }), + }; + } + if (direction == CoalescingDirection.Outgoing && older.Method == "tcp/data") + { + var first = Convert.FromBase64String(older.Params!.Value.GetProperty("data").GetString()!); + var last = Convert.FromBase64String(newer.Params!.Value.GetProperty("data").GetString()!); + var combined = new byte[first.Length + last.Length]; + first.CopyTo(combined, 0); + last.CopyTo(combined, first.Length); + return new JsonRpcNotification + { + Method = "tcp/data", + Params = _serializer.SerializeToElement(new TcpDataParams + { + Channel = owner.Uri, + Data = Convert.ToBase64String(combined), + }), + }; + } + return null; + } + + private static bool MetadataEqual(JsonElement first, JsonElement last) + { + var hasFirst = first.TryGetProperty("_meta", out var a); + var hasLast = last.TryGetProperty("_meta", out var b); + return hasFirst == hasLast && (!hasFirst || JsonEqual(a, b)); + } + + internal static bool JsonEqual(JsonElement a, JsonElement b) + { + if (a.ValueKind != b.ValueKind) return false; + if (a.ValueKind == JsonValueKind.Object) + return a.EnumerateObject().Count() == b.EnumerateObject().Count() + && a.EnumerateObject().All(p => b.TryGetProperty(p.Name, out var value) && JsonEqual(p.Value, value)); + if (a.ValueKind == JsonValueKind.Array) + return a.GetArrayLength() == b.GetArrayLength() && a.EnumerateArray().Zip(b.EnumerateArray(), JsonEqual).All(x => x); + return a.ToString() == b.ToString(); + } + + private async Task SendWindowAsync(WindowedSubscription owner, JsonRpcNotification notification, CancellationToken token) + { + var text = EncodeWindow(notification); + var completion = WindowedSubscription.Signal(); + var item = new WindowSend(notification, text, completion, token); + while (true) + { + Task wait; + token.ThrowIfCancellationRequested(); + if (_lifetimeCts.IsCancellationRequested) throw new AhpClientClosedException(); + lock (owner.Gate) + { + if (owner.Closed) throw owner.CloseError ?? new AhpClientClosedException(); + if (owner.BufferedSend == 0 || owner.BufferedSend + item.Bytes <= WindowConfig.MaxBufferedSendBytes) + { + if (notification.Method == "tcp/data") + { + if (owner.SendEof) throw new InvalidOperationException("TCP data cannot follow EOF."); + ValidateTcpData(notification.Params!.Value.GetProperty("data").GetString()!); + } + if (notification.Method == "tcp/eof") owner.SendEof = true; + if (owner.Options.Coalescing && owner.Sends.Last is { } tail && !tail.Value.Started + && tail.Value.Waiters.All(x => x.Token == token)) + { + var combined = Coalesce(tail.Value.Notification, notification, owner, CoalescingDirection.Outgoing); + if (combined is not null) + { + var encoded = EncodeWindow(combined); + var size = (long)encoded.Length * 2; + if (owner.BufferedSend - tail.Value.Bytes + size <= WindowConfig.MaxBufferedSendBytes) + { + owner.BufferedSend += size - tail.Value.Bytes; + tail.Value.Notification = combined; + tail.Value.Text = encoded; + tail.Value.Bytes = size; + tail.Value.Waiters.Add((completion, token)); + WindowEvent(owner, "coalesced", size); + break; + } + } + } + owner.Sends.AddLast(item); + owner.BufferedSend += item.Bytes; + break; + } + wait = owner.Writable.Task; + } + await Task.WhenAny(wait, Completion).WaitAsync(token).ConfigureAwait(false); + } + owner.Publish(); + _writerWake.Release(); + using (token.Register(() => _writerWake.Release())) + await completion.Task.WaitAsync(token).ConfigureAwait(false); + } + + private bool TryWindowFrame(out TransportMessage frame, out Action sent) + { + List owners; + lock (_windowGate) owners = new List(_windows.Values); + foreach (var owner in owners) + { + lock (owner.Gate) + { + if (owner.Closed || owner.Negotiating || owner.Released <= owner.CreditSent) continue; + var total = owner.Released; + var notification = new JsonRpcNotification + { + Method = "channel/credit", + Params = _serializer.SerializeToElement(new ChannelCreditParams { Channel = owner.Uri, ConsumedBytes = total }), + }; + frame = TransportMessage.FromText(EncodeWindow(notification)); + owner.CreditSent = total; + var creditBytes = RawWindowBytes(frame); + sent = () => + { + AhpTelemetry.WindowCreditBytesSent.Add(creditBytes); + WindowEvent(owner, "creditSent", creditBytes); + }; + return true; + } + } + for (var i = 0; i < owners.Count; i++) + { + var index = (_windowTurn + i) % owners.Count; + var owner = owners[index]; + lock (owner.Gate) + { + if (owner.Closed || owner.Negotiating) continue; + while (owner.Sends.First is { } canceled && !canceled.Value.Started + && canceled.Value.Waiters.All(x => x.Token.IsCancellationRequested)) + { + owner.BufferedSend -= canceled.Value.Bytes; + foreach (var waiter in canceled.Value.Waiters) waiter.Completion.TrySetCanceled(waiter.Token); + owner.Sends.RemoveFirst(); + if (canceled.Value.Notification.Method == "tcp/eof") owner.SendEof = false; + owner.Publish(); + } + if (owner.Sends.First is not { } node + || (owner.Limits is not null && owner.Sent - owner.PeerConsumed >= owner.Limits.HostReceive.WindowBytes)) continue; + var item = node.Value; + string output; + bool final; + if (owner.Limits is null || (!item.Started && item.Bytes <= owner.Limits.HostReceive.MaximumFrameBytes)) + { + output = item.Text; + final = true; + } + else + { + var remaining = owner.Limits.HostReceive.WindowBytes - (owner.Sent - owner.PeerConsumed); + output = EncodeFragment(owner, item, remaining, out var count); + item.Offset += count; + final = item.Offset == item.Text.Length; + item.FragmentWireBytes += (long)output.Length * 2; + } + item.Started = true; + _windowInFlight = item; + var charge = checked((long)output.Length * 2); + if (owner.Limits is not null) owner.Sent = checked(owner.Sent + charge); + if (owner.Sent > 9007199254740991) throw Protocol("window counter overflow"); + if (final) + { + owner.Sends.RemoveFirst(); + owner.BufferedSend -= item.Bytes; + } + _windowTurn = (index + 1) % owners.Count; + frame = TransportMessage.FromText(output); + sent = () => + { + AhpTelemetry.WindowDataBytesSent.Add(charge); + _windowInFlight = null; + if (final) + { + foreach (var waiter in item.Waiters) waiter.Completion.TrySetResult(true); + if (item.FragmentWireBytes > 0) WindowEvent(owner, "fragmentedMessageSent", item.FragmentWireBytes); + } + owner.Publish(); + }; + return true; + } + } + frame = null!; + sent = null!; + return false; + } + + private string EncodeFragment(WindowedSubscription owner, WindowSend item, long remaining, out int count) + { + string Encode(int length) => EncodeWindow(new JsonRpcNotification + { + Method = "channel/frame", + Params = _serializer.SerializeToElement(new ChannelFrameParams + { + Channel = owner.Uri, + Data = item.Text.Substring(item.Offset, length), + Final = item.Offset + length == item.Text.Length ? true : null, + }), + }); + int Boundary(int length) => length > 0 && item.Offset + length < item.Text.Length + && char.IsHighSurrogate(item.Text[item.Offset + length - 1]) + && char.IsLowSurrogate(item.Text[item.Offset + length]) ? length - 1 : length; + var smallest = char.IsHighSurrogate(item.Text[item.Offset]) ? 2 : 1; + var minimum = Encode(smallest); + var limit = Math.Min(owner.Limits!.HostReceive.MaximumFrameBytes, Math.Max(remaining, (long)minimum.Length * 2)); + if ((long)minimum.Length * 2 > limit) throw Protocol("maximumFrameBytes cannot carry a character"); + var low = smallest; + var high = item.Text.Length - item.Offset; + count = smallest; + var output = minimum; + while (low <= high) + { + var midpoint = low + (high - low) / 2; + var candidateCount = Boundary(midpoint); + if (candidateCount == 0) { low = midpoint + 1; continue; } + var candidate = Encode(candidateCount); + if ((long)candidate.Length * 2 <= limit) + { + count = candidateCount; + output = candidate; + low = midpoint + 1; + } + else high = midpoint - 1; + } + return output; + } + + internal void CloseWindowOwner(WindowedSubscription owner) + { + owner.End(); + _ = UnsubscribeOwnerAsync(owner); + } + + internal void DropManagedWindow(WindowedSubscription owner) => RemoveWindow(owner); + + private static void ValidateTcpData(string data) + { + if (data.Length == 0 || Convert.ToBase64String(Convert.FromBase64String(data)) != data) + throw Protocol("invalid TCP base64"); + } + + private void RemoveWindow(WindowedSubscription owner) + { + lock (_windowGate) + if (_windows.TryGetValue(owner.Uri, out var current) && ReferenceEquals(current, owner)) + _windows.Remove(owner.Uri); + RememberClosedWindow(owner.Uri); + } + + private void CleanupUnacceptedWindow(WindowedSubscription owner, Exception error) + { + if (owner.EverAccepted) return; + if (_windowResponses.Values.Any(x => !x.Abandoned && x.Owners.Contains(owner))) return; + owner.End(error, true); + RemoveWindow(owner); + } + + private void AbandonWindowRequest(ulong id, Exception error) + { + if (_windowResponses.TryGetValue(id, out var response) && response.Abandon(error)) + _windowResponses.TryRemove(id, out _); + } + + private void StopWindow(string uri) + { + var owner = GetWindowedSubscription(uri); + if (owner is not null) + { + owner.End(); + RemoveWindow(owner); + } + else if (_windowResponses.Values.Any(x => x.Owners.Any(owner => owner.Uri == uri))) + RememberClosedWindow(uri); + foreach (var pending in _windowResponses) + if (pending.Value.Owners.Any(x => x.Uri == uri)) + { + var error = new AhpClientClosedException("The subscription was unsubscribed."); + AbandonWindowRequest(pending.Key, error); + if (_pending.TryRemove(pending.Key, out var completion)) completion.TrySetException(error); + } + } + + internal int PendingWindowResponseCount => _windowResponses.Count; + internal int ClosedWindowCount { get { lock (_windowGate) return _closedWindows.Count; } } + + private void RememberClosedWindow(string uri) + { + lock (_windowGate) + { + if (_closedWindows.TryGetValue(uri, out var previous)) _closedWindowOrder.Remove(previous); + _closedWindows[uri] = _closedWindowOrder.AddLast(uri); + if (_closedWindows.Count > 1024 && _closedWindowOrder.First is { } oldest) + { + _closedWindows.Remove(oldest.Value); + _closedWindowOrder.RemoveFirst(); + } + } + } + + private bool IsClosedWindow(string uri) { lock (_windowGate) return _closedWindows.ContainsKey(uri); } + + private void ForgetClosedWindow(string uri) + { + lock (_windowGate) + if (_closedWindows.TryGetValue(uri, out var node)) + { + _closedWindowOrder.Remove(node); + _closedWindows.Remove(uri); + } + } + + private async Task UnsubscribeOwnerAsync(WindowedSubscription owner) + { + try { await UnsubscribeAsync(owner.Uri).ConfigureAwait(false); } + catch (AhpException) { } + } + + private void EndWindows(Exception? cause) + { + _windowInFlight?.Fail(cause ?? new AhpClientClosedException()); + List owners; + lock (_windowGate) + { + owners = new List(_windows.Values); + foreach (var reader in _flowEventReaders) reader.Writer.TryComplete(); + } + foreach (var owner in owners) + { + if (RetainManagedWindows && !owner.Closed && !owner.Uri.StartsWith("ahp-tcp:", StringComparison.Ordinal)) + { + lock (owner.Gate) + { + owner.Epoch++; + owner.Negotiating = true; + owner.BufferedReceive -= owner.FragmentBytes; + owner.FragmentBytes = 0; + owner.Fragments = null; + foreach (var item in owner.Sends) item.Fail(cause ?? new AhpClientClosedException()); + owner.Sends.Clear(); + owner.BufferedSend = 0; + } + } + else owner.End(cause, owner.Uri.StartsWith("ahp-tcp:", StringComparison.Ordinal)); + owner.Publish(); + } + _windowResponses.Clear(); + _writerWake.Release(); + } +} diff --git a/clients/dotnet/src/AgentHostProtocol/WindowedSubscription.cs b/clients/dotnet/src/AgentHostProtocol/WindowedSubscription.cs new file mode 100644 index 000000000..2049edefc --- /dev/null +++ b/clients/dotnet/src/AgentHostProtocol/WindowedSubscription.cs @@ -0,0 +1,334 @@ +#nullable enable + +using System.Runtime.CompilerServices; +using System.Threading; +using System.Threading.Channels; +using System.Threading.Tasks; + +namespace Microsoft.AgentHostProtocol; + +/// An owned delivery. Release after processing; release and disposal are idempotent. +public sealed class ReceivedDelivery : IDisposable +{ + private Action? _release; + public SubscriptionEvent Event { get; } + public JsonRpcNotification Notification { get; } + public long WireBytes { get; } + internal long CreditBytes { get; } + internal long Epoch { get; } + + internal ReceivedDelivery(SubscriptionEvent value, JsonRpcNotification notification, + long bytes, long creditBytes, long epoch, Action release) + { + Event = value; + Notification = notification; + WireBytes = bytes; + CreditBytes = creditBytes; + Epoch = epoch; + _release = release; + } + + public void Release() => Interlocked.Exchange(ref _release, null)?.Invoke(); + public void Dispose() => Release(); +} + +/// The sole credit-owning reader for a channel. Observational streams do not release credit. +public sealed class WindowedSubscription : IDisposable, IAsyncDisposable +{ + internal readonly object Gate = new(); + internal readonly LinkedList Queue = new(); + internal TaskCompletionSource Available = Signal(); + internal TaskCompletionSource Writable = Signal(); + internal readonly TaskCompletionSource Finished = Signal(); + internal readonly List> StatusReaders = new(); + internal AhpClient Client; + internal WindowOptions Options; + internal Hosts.HostId? HostId; + internal ChannelReceiveLimits Offer; + internal ChannelFlowControl? Limits; + internal string? Fallback; + internal long Epoch; + internal long Received; + internal long Released; + internal long CreditSent; + internal long Sent; + internal long PeerConsumed; + internal long BufferedSend; + internal long BufferedReceive; + internal long LastSeq; + internal bool AwaitingSnapshot; + internal bool Negotiating = true; + internal bool EverAccepted; + internal bool ResumeWindowed; + internal bool Closed; + internal bool SendEof; + internal bool ReceiveEof; + internal Exception? CloseError; + internal bool CreditTimer; + internal CancellationTokenSource? CreditCancellation; + internal Action? OnClosed; + internal readonly long TelemetryKey; + internal System.Text.StringBuilder? Fragments; + internal long FragmentBytes; + internal readonly LinkedList Sends = new(); + internal long? BlockedAt; + private readonly DeliveryReader _reader; + + public string Uri { get; } + public FlowControlStatus FlowControl + { + get + { + lock (Gate) + return new FlowControlStatus + { + Mode = Limits is null ? "ordinary" : "windowed", + FallbackReason = Fallback, + Limits = Limits, + OutstandingSendBytes = Sent - PeerConsumed, + BufferedSendBytes = BufferedSend, + BufferedReceiveBytes = BufferedReceive, + Writable = !Closed && !Negotiating && (Limits is null || Sent - PeerConsumed < Limits.HostReceive.WindowBytes) + && BufferedSend < Client.WindowConfig.MaxBufferedSendBytes, + }; + } + } + public long LastAppliedServerSeq { get { lock (Gate) return LastSeq; } } + public event EventHandler? FlowControlChanged; + /// Manual-mode reader. Use ReadEventsAsync for automatic release on advance. + public ChannelReader Deliveries => + Options.Receive == ReceiveMode.Manual ? _reader + : throw new InvalidOperationException("Use ReadEventsAsync for an automatic subscription."); + + internal WindowedSubscription(AhpClient client, string uri, WindowOptions options, ChannelReceiveLimits offer) + { + Client = client; + HostId = client.WindowHostId; + Uri = uri; + Options = options; + Offer = offer; + _reader = new DeliveryReader(this); + TelemetryKey = AhpTelemetry.TrackWindow(this); + } + + internal static TaskCompletionSource Signal() => new(TaskCreationOptions.RunContinuationsAsynchronously); + + public ChannelReader FlowControlChanges() + { + var channel = Channel.CreateBounded(new BoundedChannelOptions(1) + { + FullMode = BoundedChannelFullMode.DropOldest, + }); + lock (Gate) StatusReaders.Add(channel); + channel.Writer.TryWrite(FlowControl); + return channel.Reader; + } + + public async IAsyncEnumerable ReadEventsAsync( + [EnumeratorCancellation] CancellationToken cancellationToken = default) + { + if (Options.Receive != ReceiveMode.Auto) + throw new InvalidOperationException("Use Deliveries for a manual subscription."); + await foreach (var delivery in _reader.ReadAllAsync(cancellationToken).ConfigureAwait(false)) + { + using (delivery) yield return delivery.Event; + } + } + + public async ValueTask WaitWritableAsync(CancellationToken cancellationToken = default) + { + while (true) + { + Task wait; + lock (Gate) + { + if (Closed) throw CloseError ?? new AhpClientClosedException(); + if (Client.Completion.IsCompleted) throw new AhpClientClosedException(); + if (FlowControl.Writable) return; + wait = Writable.Task; + } + await wait.WaitAsync(cancellationToken).ConfigureAwait(false); + } + } + + public Task NotifyAsync(string method, T parameters, CancellationToken cancellationToken = default) + { + if (Closed) throw CloseError ?? new AhpClientClosedException(); + var encoded = Client.WindowParameters(parameters); + if (!encoded.TryGetProperty("channel", out var channel) || channel.GetString() != Uri) + throw new ArgumentException("The notification must target this subscription.", nameof(parameters)); + return Client.NotifyAsync(method, encoded, cancellationToken); + } + public Task DispatchAsync(StateAction action, long? clientSeq = null, + CancellationToken cancellationToken = default) + { + if (Closed) throw CloseError ?? new AhpClientClosedException(); + return Client.DispatchAsync(Uri, action, clientSeq, cancellationToken); + } + public Task UnsubscribeAsync(CancellationToken cancellationToken = default) + { + if (ReferenceEquals(Client.GetWindowedSubscription(Uri), this) && !Client.Completion.IsCompleted) + return Client.UnsubscribeAsync(Uri, cancellationToken); + End(); + return Task.CompletedTask; + } + public async ValueTask DisposeAsync() => await UnsubscribeAsync().ConfigureAwait(false); + public void Dispose() + { + if (!Closed && ReferenceEquals(Client.GetWindowedSubscription(Uri), this)) Client.CloseWindowOwner(this); + } + + internal void Publish() + { + var status = FlowControl; + List> readers; + lock (Gate) + { + Writable.TrySetResult(true); + Writable = Signal(); + readers = new List>(StatusReaders); + if (!status.Writable && !Closed && !Negotiating && BlockedAt is null) + { + BlockedAt = System.Diagnostics.Stopwatch.GetTimestamp(); + Client.WindowEvent(this, "blocked"); + } + else if ((status.Writable || Closed) && BlockedAt is { } start) + { + BlockedAt = null; + AhpTelemetry.WindowBlockedDuration.Record(Compatibility.GetElapsedTime(start).TotalMilliseconds); + Client.WindowEvent(this, "unblocked"); + } + } + foreach (var reader in readers) reader.Writer.TryWrite(status); + if (FlowControlChanged is { } changed) + foreach (EventHandler observer in changed.GetInvocationList()) + { + try { observer(this, status); } + catch (Exception) { /* Observational callbacks must not interrupt the protocol driver. */ } + } + } + + internal void End(Exception? error = null, bool discard = false) + { + lock (Gate) + { + if (Closed) return; + Closed = true; + Epoch++; + CreditTimer = false; + CreditCancellation?.Cancel(); + CloseError = error; + if (discard) + { + foreach (var delivery in Queue) BufferedReceive -= delivery.WireBytes; + Queue.Clear(); + BufferedReceive -= FragmentBytes; + FragmentBytes = 0; + Fragments = null; + } + Available.TrySetResult(true); + Writable.TrySetResult(true); + foreach (var reader in StatusReaders) reader.Writer.TryComplete(error); + foreach (var send in Sends) send.Fail(error ?? new AhpClientClosedException()); + Sends.Clear(); + BufferedSend = 0; + if (Queue.Count == 0) CompleteReader(); + } + AhpTelemetry.UntrackWindow(TelemetryKey); + OnClosed?.Invoke(this); + Publish(); + } + + internal ChannelReader InternalDeliveries => _reader; + + private void CompleteReader() + { + if (CloseError is { } error) Finished.TrySetException(error); + else Finished.TrySetResult(true); + } + + private sealed class DeliveryReader(WindowedSubscription owner) : ChannelReader + { + public override Task Completion => owner.Finished.Task; + public override bool TryRead(out ReceivedDelivery item) + { + lock (owner.Gate) + { + if (owner.Queue.First is { } first) + { + item = first.Value; + owner.Queue.RemoveFirst(); + if (owner.Closed && owner.Queue.Count == 0) owner.CompleteReader(); + return true; + } + item = null!; + return false; + } + } + public override async ValueTask WaitToReadAsync(CancellationToken cancellationToken = default) + { + while (true) + { + Task wait; + lock (owner.Gate) + { + if (owner.Queue.Count > 0) return true; + if (owner.Closed) + { + if (owner.CloseError is { } error) throw error; + return false; + } + wait = owner.Available.Task; + } + await wait.WaitAsync(cancellationToken).ConfigureAwait(false); + } + } + } +} + +/// A notification without a dedicated event variant, including future methods. +public sealed class SubscriptionEventNotification(JsonRpcNotification notification) : SubscriptionEvent +{ + public JsonRpcNotification Notification { get; } = notification; +} +public sealed class SubscriptionEventTcpData(TcpDataParams parameters) : SubscriptionEvent +{ + public TcpDataParams Params { get; } = parameters; +} +public sealed class SubscriptionEventTcpEof(TcpEofParams parameters) : SubscriptionEvent +{ + public TcpEofParams Params { get; } = parameters; +} +public sealed class SubscriptionEventTcpReset(TcpResetParams parameters) : SubscriptionEvent +{ + public TcpResetParams Params { get; } = parameters; +} + +internal sealed class WindowSend +{ + internal JsonRpcNotification Notification; + internal string Text; + internal int Offset; + internal bool Started; + internal long Bytes; + internal long FragmentWireBytes; + internal readonly List<(TaskCompletionSource Completion, CancellationToken Token)> Waiters = new(); + internal WindowSend(JsonRpcNotification notification, string text, TaskCompletionSource completion, CancellationToken token) + { + Notification = notification; + Text = text; + Bytes = (long)text.Length * 2; + Waiters.Add((completion, token)); + } + + internal void Fail(Exception error) { foreach (var waiter in Waiters) waiter.Completion.TrySetException(error); } +} + +internal sealed class WindowWakeSignal +{ + private readonly Channel _channel = Channel.CreateBounded( + new BoundedChannelOptions(1) { FullMode = BoundedChannelFullMode.DropWrite }); + internal void Release() => _channel.Writer.TryWrite(true); + internal async Task WaitAsync(CancellationToken cancellationToken) => + _ = await _channel.Reader.ReadAsync(cancellationToken).ConfigureAwait(false); +} diff --git a/clients/dotnet/tests/AgentHostProtocol.AotSmoke/Program.cs b/clients/dotnet/tests/AgentHostProtocol.AotSmoke/Program.cs index d4aed706c..b9ed0e364 100644 --- a/clients/dotnet/tests/AgentHostProtocol.AotSmoke/Program.cs +++ b/clients/dotnet/tests/AgentHostProtocol.AotSmoke/Program.cs @@ -76,12 +76,73 @@ public static async Task Main() "Generated AHP metadata resolver did not provide ActionEnvelope metadata."); await RunTransportScenarioAsync(serializer); + await RunWindowScenarioAsync(serializer); Console.WriteLine( "Native AOT smoke passed: initialize, reconnect, ping, subscribe/action/reducer/unsubscribe, " + "custom request/notification metadata, and typed/raw/null inbound request results."); } + static async Task RunWindowScenarioAsync(SystemTextJsonAhpSerializer serializer) + { + const string channel = "ahp-test:/aot-window"; + using var timeout = new CancellationTokenSource(TimeSpan.FromSeconds(10)); + var (clientTransport, server) = DuplexTransport.CreatePair(); + await using var serverOwner = server; + await using var client = AhpClient.Connect(clientTransport, serializer: serializer); + var limits = new ChannelReceiveLimits { WindowBytes = 256, MaximumFrameBytes = 1024 }; + var response = RespondToRequestAsync(server, serializer, "subscribe", + new SubscribeResult { FlowControl = new ChannelFlowControl { ClientReceive = limits, HostReceive = limits } }, + timeout.Token); + var (_, subscription) = await client.SubscribeWindowedAsync(channel, + new WindowOptions { Receive = ReceiveMode.Manual }, timeout.Token); + await response; + var snapshot = serializer.EncodeMessage(new JsonRpcMessage + { + Notification = new JsonRpcNotification + { + Method = "channel/snapshot", + Params = serializer.SerializeToElement(new ChannelSnapshotParams + { + Channel = channel, + Snapshot = new Snapshot + { + Resource = channel, + FromSeq = 1, + State = new SnapshotState { Empty = new EmptyState() }, + }, + }), + }, + }).Text!; + var split = snapshot.Length / 2; + for (var part = 0; part < 2; part++) + { + await server.SendAsync(serializer.EncodeMessage(new JsonRpcMessage + { + Notification = new JsonRpcNotification + { + Method = "channel/frame", + Params = serializer.SerializeToElement(new ChannelFrameParams + { + Channel = channel, + Data = part == 0 ? snapshot.Substring(0, split) : snapshot.Substring(split), + Final = part == 1 ? true : null, + }), + }, + }), timeout.Token); + if (part == 0) + { + var credit = await ReceiveMessageAsync(server, serializer, timeout.Token); + Ensure(credit.Notification?.Method == "channel/credit", "Nonfinal fragment was not credited."); + } + } + using var delivery = await subscription.Deliveries.ReadAsync(timeout.Token); + Ensure(delivery.Event is SubscriptionEventSnapshot, "Window snapshot was not reassembled."); + delivery.Release(); + var finalCredit = await ReceiveMessageAsync(server, serializer, timeout.Token); + Ensure(finalCredit.Notification?.Method == "channel/credit", "Final fragment was not credited on release."); + } + static async Task RunTransportScenarioAsync(SystemTextJsonAhpSerializer serializer) { const string envelopeChannel = "ahp-session:/native-aot"; diff --git a/clients/dotnet/tests/AgentHostProtocol.Tests/WindowedClientTests.cs b/clients/dotnet/tests/AgentHostProtocol.Tests/WindowedClientTests.cs new file mode 100644 index 000000000..1a60e7671 --- /dev/null +++ b/clients/dotnet/tests/AgentHostProtocol.Tests/WindowedClientTests.cs @@ -0,0 +1,341 @@ +#nullable enable + +using System.Diagnostics; +using System.Diagnostics.Metrics; +using System.Linq; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.AgentHostProtocol.Hosts; +using Xunit; + +namespace Microsoft.AgentHostProtocol.Tests; + +public sealed class WindowedClientTests +{ + private const string Uri = "ahp-test:/window"; + private const string SnapshotText = "{\"jsonrpc\":\"2.0\",\"method\":\"channel/snapshot\",\"params\":{\"channel\":\"ahp-test:/window\",\"snapshot\":{\"resource\":\"ahp-test:/window\",\"state\":{},\"fromSeq\":7}}}"; + private static readonly SystemTextJsonAhpSerializer Ser = SystemTextJsonAhpSerializer.Default; + + [Fact] + public async Task BlockedChannelDoesNotHoldOtherChannelsOrCanceledSends() => + await CheckFairnessAsync(); + + private static async Task CheckFairnessAsync() + { + var transport = new CorpusTransport(); + RespondAutomatically(transport); + await using var client = AhpClient.Connect(transport); + var (_, blocked) = await client.SubscribeWindowedAsync(Uri, + new WindowOptions { WindowBytes = 512, MaximumFrameBytes = 1024 }).ConfigureAwait(false); + var (_, other) = await client.SubscribeWindowedAsync("ahp-test:/other").ConfigureAwait(false); + await blocked.NotifyAsync("future/data", new TcpDataParams { Channel = Uri, Data = new string('x', 220) }).ConfigureAwait(false); + Assert.False(blocked.FlowControl.Writable); + using var cancellation = new CancellationTokenSource(); + var canceled = blocked.NotifyAsync("future/canceled", new TcpEofParams { Channel = Uri }, cancellation.Token); + await WindowedDeliveryCorpusTests.UntilAsync(() => blocked.FlowControl.BufferedSendBytes > 0).ConfigureAwait(false); + cancellation.Cancel(); + await Assert.ThrowsAnyAsync(() => canceled).ConfigureAwait(false); + await other.NotifyAsync("future/data", new TcpEofParams { Channel = other.Uri }).WaitAsync(TimeSpan.FromSeconds(1)).ConfigureAwait(false); + await WindowedDeliveryCorpusTests.UntilAsync(() => blocked.FlowControl.BufferedSendBytes == 0).ConfigureAwait(false); + Assert.DoesNotContain(transport.Snapshot(), x => x.Json.GetProperty("method").GetString() == "future/canceled"); + var writable = blocked.WaitWritableAsync().AsTask(); + Assert.False(writable.IsCompleted); + transport.Inject(Ser.Serialize(new JsonRpcNotification + { + Method = "channel/credit", + Params = Ser.SerializeToElement(new ChannelCreditParams { Channel = Uri, ConsumedBytes = blocked.Sent }), + })); + await writable.WaitAsync(TimeSpan.FromSeconds(1)).ConfigureAwait(false); + } + + [Fact] + public async Task MissingWindowClosesItsRetainedQueue() => + await CheckMissingWindowAsync(); + + private static async Task CheckMissingWindowAsync() + { + var transport = new CorpusTransport(); + RespondAutomatically(transport, keepWindows: false); + await using var client = AhpClient.Connect(transport); + var (_, owner) = await client.SubscribeWindowedAsync(Uri, new WindowOptions { Receive = ReceiveMode.Manual }).ConfigureAwait(false); + transport.Inject(SnapshotText); + await WindowedDeliveryCorpusTests.UntilAsync(() => owner.LastAppliedServerSeq == 7).ConfigureAwait(false); + await client.ResumeAsync(new ReconnectParams + { + Channel = ProtocolVersion.RootResourceUri, + ClientId = "missing-window", + LastSeenServerSeq = 7, + Subscriptions = new List { Uri }, + Windows = new ChannelDeliveryResumeOptions + { + Items = new List + { + new() { Channel = Uri, FlowControl = new SubscriptionFlowControlOptions { Receive = owner.Offer } }, + }, + }, + }).ConfigureAwait(false); + Assert.Null(client.GetWindowedSubscription(Uri)); + await Assert.ThrowsAsync(() => owner.Deliveries.ReadAsync().AsTask()).ConfigureAwait(false); + Assert.Equal(0, owner.FlowControl.BufferedReceiveBytes); + } + + [Fact] + public async Task TcpOwnerClosesAndIsExcludedFromManagedReconnect() => + await CheckTcpReconnectAsync(); + + private static async Task CheckTcpReconnectAsync() + { + var transports = new List(); + await using var multi = new MultiHostClient(); + await multi.AddHostAsync(new HostConfig + { + Id = "tcp-host", + ReconnectPolicy = ReconnectPolicy.Disabled, + TransportFactory = (_, _) => + { + var transport = new CorpusTransport(); + RespondAutomatically(transport); + lock (transports) transports.Add(transport); + return Task.FromResult(transport); + }, + }).ConfigureAwait(false); + var (_, tcp) = await multi.SubscribeWindowedAsync("tcp-host", "ahp-tcp:/test", + new WindowOptions { Receive = ReceiveMode.Manual }).ConfigureAwait(false); + var generation = multi.Host("tcp-host")!.Generation; + await multi.ReconnectAsync("tcp-host").ConfigureAwait(false); + await WindowedDeliveryCorpusTests.UntilAsync(() => multi.Host("tcp-host")!.Generation > generation).ConfigureAwait(false); + Assert.True(tcp.Closed); + var request = transports[1].Snapshot().First(x => x.Json.GetProperty("method").GetString() == "reconnect").Json; + Assert.DoesNotContain(request.GetProperty("params").GetProperty("subscriptions").EnumerateArray(), + x => x.GetString() == tcp.Uri); + Assert.False(request.GetProperty("params").TryGetProperty("windows", out _)); + } + + [Theory] + [InlineData(null, 1048576)] + [InlineData(33554432L, 22369620)] + public async Task TransportDefaultsAreEncodingSafe(long? cap, long expected) => + await CheckDefaultAsync(cap, expected); + + private static async Task CheckDefaultAsync(long? cap, long expected) + { + var transport = new CorpusTransport { MaximumInboundMessageBytes = cap }; + await using var client = AhpClient.Connect(transport); + RespondAutomatically(transport); + var (_, owner) = await client.SubscribeWindowedAsync(Uri).ConfigureAwait(false); + Assert.Equal(expected, owner.Offer.MaximumFrameBytes); + Assert.Equal(262144, owner.Offer.WindowBytes); + } + + [Fact] + public async Task AutomaticIteratorReleasesOnlyOnAdvanceAndDisposal() => + await CheckAutomaticAsync(); + + private static async Task CheckAutomaticAsync() + { + var transport = new CorpusTransport(); + RespondAutomatically(transport); + await using var client = AhpClient.Connect(transport); + var (_, owner) = await client.SubscribeWindowedAsync(Uri).ConfigureAwait(false); + var observers = client.CreateEventStream(); + using (observers) + { + transport.Inject(SnapshotText); + await using var iterator = owner.ReadEventsAsync().GetAsyncEnumerator(); + Assert.True(await iterator.MoveNextAsync().ConfigureAwait(false)); + Assert.IsType(iterator.Current); + Assert.DoesNotContain(transport.Snapshot(), x => x.Json.GetProperty("method").GetString() == "channel/credit"); + var waiting = iterator.MoveNextAsync().AsTask(); + await WindowedDeliveryCorpusTests.UntilAsync(() => HasCredit(transport, (long)SnapshotText.Length * 2)).ConfigureAwait(false); + Assert.False(waiting.IsCompleted); + transport.Inject("{\"jsonrpc\":\"2.0\",\"method\":\"future/data\",\"params\":{\"channel\":\"ahp-test:/window\",\"value\":1}}"); + Assert.True(await waiting.ConfigureAwait(false)); + Assert.IsType(iterator.Current); + } + await WindowedDeliveryCorpusTests.UntilAsync(() => owner.FlowControl.BufferedReceiveBytes == 0).ConfigureAwait(false); + } + + [Theory] + [InlineData("snapshot")] + [InlineData("replay")] + [InlineData("live")] + public async Task ManagedEpochRetainsUnreadItemsWithoutCreditingNewConnection(string recovery) => + await CheckEpochAsync(recovery); + + private static async Task CheckEpochAsync(string recovery) + { + var first = new CorpusTransport(); + RespondAutomatically(first); + await using var oldClient = AhpClient.Connect(first); + oldClient.RetainManagedWindows = true; + var (_, owner) = await oldClient.SubscribeWindowedAsync(Uri, new WindowOptions { Receive = ReceiveMode.Manual }).ConfigureAwait(false); + first.Inject(SnapshotText); + await WindowedDeliveryCorpusTests.UntilAsync(() => owner.Queue.Count == 1).ConfigureAwait(false); + await oldClient.ShutdownAsync().ConfigureAwait(false); + var transport = new CorpusTransport(); + RespondAutomatically(transport, recovery); + await using var client = AhpClient.Connect(transport); + client.BindManagedWindow(owner); + var result = await client.ResumeAsync(new ReconnectParams + { + Channel = ProtocolVersion.RootResourceUri, + ClientId = "window-client", + LastSeenServerSeq = 7, + Subscriptions = new List { Uri }, + Windows = new ChannelDeliveryResumeOptions + { + Items = new List + { + new() { Channel = Uri, FlowControl = new SubscriptionFlowControlOptions { Receive = owner.Offer }, LastAppliedServerSeq = owner.LastAppliedServerSeq }, + }, + }, + }).ConfigureAwait(false); + Assert.NotNull(result); + var old = await owner.Deliveries.ReadAsync().ConfigureAwait(false); + old.Dispose(); + Assert.DoesNotContain(transport.Snapshot(), x => x.Json.GetProperty("method").GetString() == "channel/credit"); + Assert.Equal(0, owner.FlowControl.BufferedReceiveBytes); + if (recovery == "snapshot") transport.Inject(SnapshotText); + else transport.Inject("{\"jsonrpc\":\"2.0\",\"method\":\"future/data\",\"params\":{\"channel\":\"ahp-test:/window\"}}"); + using var current = await owner.Deliveries.ReadAsync().ConfigureAwait(false); + current.Release(); + await WindowedDeliveryCorpusTests.UntilAsync(() => HasCredit(transport, current.WireBytes)).ConfigureAwait(false); + } + + [Fact] + public async Task MultiHostKeepsOwnerAndOptionsAcrossReconnect() => + await CheckMultiHostAsync(); + + private static async Task CheckMultiHostAsync() + { + var transports = new List(); + await using var multi = new MultiHostClient(); + await multi.AddHostAsync(new HostConfig + { + Id = "window-host", + ReconnectPolicy = ReconnectPolicy.Disabled, + TransportFactory = (_, _) => + { + var transport = new CorpusTransport(); + RespondAutomatically(transport); + lock (transports) transports.Add(transport); + return Task.FromResult(transport); + }, + }).ConfigureAwait(false); + var (_, owner) = await multi.SubscribeWindowedAsync("window-host", Uri, + new WindowOptions { Receive = ReceiveMode.Manual, WindowBytes = 8192 }).ConfigureAwait(false); + transports[0].Inject(SnapshotText); + await WindowedDeliveryCorpusTests.UntilAsync(() => owner.LastAppliedServerSeq == 7).ConfigureAwait(false); + var generation = multi.Host("window-host")!.Generation; + await multi.ReconnectAsync("window-host").ConfigureAwait(false); + await WindowedDeliveryCorpusTests.UntilAsync(() => multi.Host("window-host")!.Generation > generation).ConfigureAwait(false); + Assert.Same(owner, multi.GetWindowedSubscription("window-host", Uri)); + var request = transports[1].Snapshot().First(x => x.Json.GetProperty("method").GetString() == "reconnect").Json; + Assert.Equal(7, request.GetProperty("params").GetProperty("windows").GetProperty("items")[0].GetProperty("lastAppliedServerSeq").GetInt64()); + Assert.Equal(8192, owner.Offer.WindowBytes); + using var old = await owner.Deliveries.ReadAsync().ConfigureAwait(false); + old.Release(); + Assert.DoesNotContain(transports[1].Snapshot(), x => x.Json.GetProperty("method").GetString() == "channel/credit"); + } + + [Fact] + public async Task CustomCoalescerAndTelemetryAreObservable() => + await CheckTelemetryAsync(); + + private static async Task CheckTelemetryAsync() + { + var measurements = new System.Collections.Concurrent.ConcurrentDictionary(StringComparer.Ordinal); + using var meter = new MeterListener + { + InstrumentPublished = (instrument, listener) => + { + if (instrument.Meter.Name == AhpTelemetry.Name) listener.EnableMeasurementEvents(instrument); + }, + }; + meter.SetMeasurementEventCallback((instrument, value, _, _) => + measurements.AddOrUpdate(instrument.Name, value, (_, previous) => previous + value)); + meter.Start(); + var events = new System.Collections.Concurrent.ConcurrentBag(); + using var activities = new ActivityListener + { + ShouldListenTo = source => source.Name == AhpTelemetry.Name, + Sample = (ref ActivityCreationOptions _) => ActivitySamplingResult.AllData, + ActivityStopped = activity => + { + foreach (var ev in activity.Events) events.Add(ev.Name); + }, + }; + ActivitySource.AddActivityListener(activities); + var transport = new CorpusTransport(); + RespondAutomatically(transport); + var config = new ClientConfig(); + config.Windowing.Coalescers.Add(new TestCoalescer()); + await using var client = AhpClient.Connect(transport, config); + var flowEvents = client.FlowControlEvents(); + var (_, owner) = await client.SubscribeWindowedAsync(Uri, new WindowOptions { Receive = ReceiveMode.Manual }).ConfigureAwait(false); + transport.Inject(SnapshotText); + using var snapshot = await owner.Deliveries.ReadAsync().ConfigureAwait(false); + snapshot.Release(); + const string text = "{\"jsonrpc\":\"2.0\",\"method\":\"root/progress\",\"params\":{\"channel\":\"ahp-test:/window\",\"progressToken\":\"p\"}}"; + transport.Inject(text); + transport.Inject(text); + await WindowedDeliveryCorpusTests.UntilAsync(() => + { + lock (owner.Gate) return owner.Queue.First?.Value.WireBytes == text.Length * 4L; + }).ConfigureAwait(false); + meter.RecordObservableInstruments(); + Assert.True(measurements.ContainsKey("ahp.client.window.buffered")); + using var merged = await owner.Deliveries.ReadAsync().ConfigureAwait(false); + merged.Release(); + await WindowedDeliveryCorpusTests.UntilAsync(() => measurements.ContainsKey("ahp.client.window.credit.sent")).ConfigureAwait(false); + Assert.True(measurements["ahp.client.window.data.received"] > 0); + Assert.Contains(events, x => x == "ahp.flowcontrol.coalesced"); + Assert.True(flowEvents.TryRead(out _)); + } + + private sealed class TestCoalescer : IChannelMessageCoalescer + { + public JsonRpcNotification? Coalesce(JsonRpcNotification earlier, JsonRpcNotification later, CoalescingContext context) => + context.Direction == CoalescingDirection.Incoming && earlier.Method == "root/progress" ? later : null; + } + + private static bool HasCredit(CorpusTransport transport, long amount) => transport.Snapshot().Any(x => + x.Json.GetProperty("method").GetString() == "channel/credit" + && x.Json.GetProperty("params").GetProperty("consumedBytes").GetInt64() >= amount); + + private static void RespondAutomatically(CorpusTransport transport, string recovery = "replay", bool keepWindows = true) + { + transport.OnSent = request => + { + if (!request.TryGetProperty("id", out var id)) return; + var method = request.GetProperty("method").GetString(); + object result; + if (method == "subscribe") + { + var offer = Ser.Deserialize(request.GetProperty("params")).FlowControl!.Receive; + result = new SubscribeResult { FlowControl = new ChannelFlowControl { ClientReceive = offer, HostReceive = offer } }; + } + else if (method == "initialize") + result = new InitializeResult { ProtocolVersion = "1.0.0", ServerSeq = 0, Snapshots = new List() }; + else if (method == "reconnect") + { + var parameters = Ser.Deserialize(request.GetProperty("params")); + var items = parameters.Windows?.Items.Where(_ => keepWindows).Select(x => new ResumedChannelSubscription + { + Channel = x.Channel, + FlowControl = new ChannelFlowControl { ClientReceive = x.FlowControl.Receive, HostReceive = x.FlowControl.Receive }, + Recovery = recovery switch { "snapshot" => ChannelRecoveryKind.Snapshot, "live" => ChannelRecoveryKind.Live, _ => ChannelRecoveryKind.Replay }, + }).ToList(); + result = new ReconnectReplayResult + { + Type = ReconnectResultType.Replay, + Actions = new List(), + Missing = new List(), + Windows = items is null ? null : new ChannelDeliveryResumeResult { Items = items }, + }; + } + else result = new Dictionary(); + transport.Inject("{\"jsonrpc\":\"2.0\",\"id\":" + id.GetRawText() + ",\"result\":" + Ser.Serialize(result) + "}"); + }; + } +} diff --git a/clients/dotnet/tests/AgentHostProtocol.Tests/WindowedDeliveryCorpusTests.cs b/clients/dotnet/tests/AgentHostProtocol.Tests/WindowedDeliveryCorpusTests.cs new file mode 100644 index 000000000..b9ee8ee2c --- /dev/null +++ b/clients/dotnet/tests/AgentHostProtocol.Tests/WindowedDeliveryCorpusTests.cs @@ -0,0 +1,314 @@ +#nullable enable + +using System.IO; +using System.Linq; +using System.Threading; +using System.Threading.Channels; +using System.Threading.Tasks; +using Xunit; + +namespace Microsoft.AgentHostProtocol.Tests; + +public sealed class WindowedDeliveryCorpusTests +{ + private static readonly JsonSerializerOptions StatusJson = new() { PropertyNamingPolicy = JsonNamingPolicy.CamelCase }; + private static string Corpus + { + get + { + var directory = new DirectoryInfo(AppContext.BaseDirectory); + while (directory is not null) + { + var candidate = Path.Combine(directory.FullName, "types", "test-cases", "windowed-delivery"); + if (Directory.Exists(candidate)) return candidate; + directory = directory.Parent; + } + throw new DirectoryNotFoundException("Shared windowed delivery corpus was not found."); + } + } + + public static IEnumerable Scenarios() => + Directory.GetFiles(Path.Combine(Corpus, "scenarios"), "*.json").OrderBy(x => x, StringComparer.Ordinal) + .Select(path => new object[] { Path.GetFileName(path), path }); + + [Fact] + public void AllRawSizeCases() + { + using var document = JsonDocument.Parse(File.ReadAllText(Path.Combine(Corpus, "sizes.json"))); + var cases = document.RootElement.ValueKind == JsonValueKind.Array ? document.RootElement : document.RootElement.GetProperty("cases"); + foreach (var entry in cases.EnumerateArray()) + { + var text = entry.GetProperty("text").GetString()!; + var expected = entry.GetProperty("utf16Bytes").GetInt64(); + Assert.Equal(expected, AhpClient.RawWindowBytes(TransportMessage.FromText(text))); + Assert.Equal(expected, AhpClient.RawWindowBytes(TransportMessage.FromBinary(System.Text.Encoding.UTF8.GetBytes(text)))); + } + } + + [Theory] + [MemberData(nameof(Scenarios))] + public async Task EveryScenario(string name, string path) + { + _ = name; + await RunScenarioAsync(path); + } + + private static async Task RunScenarioAsync(string path) + { + using var document = JsonDocument.Parse(File.ReadAllText(path)); + var scenario = document.RootElement; + var defaults = scenario.GetProperty("windowing"); + var config = new ClientConfig + { + Windowing = new WindowingConfig + { + WindowBytes = defaults.GetProperty("windowBytes").GetInt64(), + MaximumFrameBytes = defaults.GetProperty("maximumFrameBytes").GetInt64(), + MaxBufferedSendBytes = defaults.TryGetProperty("maxBufferedSendBytes", out var sendLimit) ? sendLimit.GetInt64() : 1048576, + }, + }; + var transport = new CorpusTransport(); + await using var client = AhpClient.Connect(transport, config); + var handles = new Dictionary(StringComparer.Ordinal); + var modes = new Dictionary(StringComparer.Ordinal); + var deliveries = new Dictionary>(StringComparer.Ordinal); + var ordinary = new Dictionary>(StringComparer.Ordinal); + var requests = new List(); + var ordinarySnapshotRead = new HashSet(StringComparer.Ordinal); + var answered = new HashSet(); + var readIndex = 0; + var currentStep = ""; + try + { + foreach (var step in scenario.GetProperty("steps").EnumerateArray()) + { + var op = step.GetProperty("op").GetString(); + currentStep = step.GetRawText(); + var uri = step.TryGetProperty("channel", out var channel) ? channel.GetString()! : ""; + switch (op) + { + case "subscribe": + var manual = step.TryGetProperty("receive", out var receive) && receive.GetString() == "manual"; + modes[uri] = manual; + if (!deliveries.ContainsKey(uri)) deliveries[uri] = new List(); + if (step.TryGetProperty("window", out var window)) + { + var task = client.SubscribeWindowedAsync(uri, new WindowOptions + { + WindowBytes = window.TryGetProperty("windowBytes", out var w) ? w.GetInt64() : null, + MaximumFrameBytes = window.TryGetProperty("maximumFrameBytes", out var f) ? f.GetInt64() : null, + Receive = manual ? ReceiveMode.Manual : ReceiveMode.Auto, + Coalescing = !step.TryGetProperty("coalescing", out var c) || c.GetBoolean(), + }); + requests.Add(task); + handles[uri] = client.GetWindowedSubscription(uri)!; + } + else + { + ordinary[uri] = client.SubscribeAsync(uri); + requests.Add(ordinary[uri]); + } + break; + case "respond": + await UntilAsync(() => transport.Snapshot().Any(x => + x.Json.TryGetProperty("id", out var id) && !answered.Contains(id.GetUInt64()) + && x.Json.GetProperty("method").GetString() == step.GetProperty("to").GetString())).ConfigureAwait(false); + var pending = transport.Snapshot().First(x => x.Json.TryGetProperty("id", out var id) + && !answered.Contains(id.GetUInt64()) + && x.Json.GetProperty("method").GetString() == step.GetProperty("to").GetString()); + var requestId = pending.Json.GetProperty("id").GetUInt64(); + answered.Add(requestId); + transport.Inject("{\"jsonrpc\":\"2.0\",\"id\":" + requestId.ToString(System.Globalization.CultureInfo.InvariantCulture) + + ",\"result\":" + step.GetProperty("result").GetRawText() + "}"); + break; + case "hostSends": + transport.Inject(step.GetProperty("text").GetString()!); + break; + case "hostCreditsAll": + var total = transport.Snapshot().Where(x => + !x.Json.TryGetProperty("id", out _) && x.Json.GetProperty("method").GetString() is { } method + && WindowedDelivery.CountsAgainstWindow(method) + && x.Json.GetProperty("params").GetProperty("channel").GetString() == uri) + .Sum(x => (long)x.Text.Length * 2); + transport.Inject(SystemTextJsonAhpSerializer.Default.Serialize(new JsonRpcNotification + { + Method = "channel/credit", + Params = SystemTextJsonAhpSerializer.Default.SerializeToElement(new ChannelCreditParams + { + Channel = uri, + ConsumedBytes = total, + }), + })); + break; + case "expectSent": + var match = step.GetProperty("match"); + var assembly = new System.Text.StringBuilder(); + while (true) + { + await UntilAsync(() => transport.Snapshot().Length > readIndex).ConfigureAwait(false); + var actual = transport.Snapshot()[readIndex++]; + var actualJson = actual.Json; + if (actualJson.GetProperty("method").GetString() == "channel/frame" + && step.TryGetProperty("reassemble", out var reassemble) && reassemble.GetBoolean()) + { + var p = actualJson.GetProperty("params"); + Assert.True((long)actual.Text.Length * 2 <= handles[p.GetProperty("channel").GetString()!].Limits!.HostReceive.MaximumFrameBytes); + assembly.Append(p.GetProperty("data").GetString()); + if (!p.TryGetProperty("final", out var final) || !final.GetBoolean()) continue; + using var joined = JsonDocument.Parse(assembly.ToString()); + actualJson = joined.RootElement.Clone(); + } + if (Matches(match, actualJson)) + { + if (step.TryGetProperty("absent", out var absent)) + foreach (var missing in absent.EnumerateArray()) + Assert.False(ContainsPath(actualJson, missing.GetString()!)); + break; + } + Assert.Equal("channel/credit", actualJson.GetProperty("method").GetString()); + } + break; + case "expectNotSent": + await Task.Delay(step.GetProperty("withinMs").GetInt32()).ConfigureAwait(false); + Assert.DoesNotContain(transport.Snapshot().Skip(readIndex), x => Matches(step.GetProperty("match"), x.Json)); + break; + case "next": + if (!modes[uri] && deliveries[uri].Count > 0) deliveries[uri][deliveries[uri].Count - 1].Release(); + ReceivedDelivery delivery; + if (handles.TryGetValue(uri, out var handle)) + { + await UntilAsync(() => + { + lock (handle.Gate) + return handle.Queue.First is { } first + && Matches(step.GetProperty("expect"), SystemTextJsonAhpSerializer.Default.SerializeToElement(first.Value.Notification)); + }).ConfigureAwait(false); + delivery = await handle.InternalDeliveries.ReadAsync().ConfigureAwait(false); + } + else + { + var subscribed = await ordinary[uri].WaitAsync(TimeSpan.FromSeconds(1)).ConfigureAwait(false); + if (ordinarySnapshotRead.Add(uri) && subscribed.Result.Snapshot is { } snapshot) + { + var notification = new JsonRpcNotification + { + Method = "channel/snapshot", + Params = SystemTextJsonAhpSerializer.Default.SerializeToElement(new ChannelSnapshotParams + { + Channel = uri, + Snapshot = snapshot, + }), + }; + delivery = new ReceivedDelivery(new SubscriptionEventSnapshot(snapshot), notification, 0, 0, 0, () => { }); + } + else + { + var value = await subscribed.Sub.Events.ReadAsync().AsTask().WaitAsync(TimeSpan.FromSeconds(1)).ConfigureAwait(false); + var action = Assert.IsType(value); + delivery = new ReceivedDelivery(value, new JsonRpcNotification + { + Method = "action", + Params = SystemTextJsonAhpSerializer.Default.SerializeToElement(action.Envelope), + }, 0, 0, 0, () => { }); + } + } + deliveries[uri].Add(delivery); + Assert.True(Matches(step.GetProperty("expect"), SystemTextJsonAhpSerializer.Default.SerializeToElement(delivery.Notification))); + if (step.TryGetProperty("expectBytes", out var expectedBytes)) Assert.Equal(expectedBytes.GetInt64(), delivery.WireBytes); + break; + case "advance": + deliveries[uri][deliveries[uri].Count - 1].Release(); + break; + case "release": + deliveries[uri][step.GetProperty("index").GetInt32()].Release(); + break; + case "send": + requests.Add(client.NotifyAsync(step.GetProperty("method").GetString()!, step.GetProperty("params"))); + break; + case "unsubscribe": + await client.UnsubscribeAsync(uri).ConfigureAwait(false); + break; + case "expectStatus": + await UntilAsync(() => + { + var status = handles.TryGetValue(uri, out var owner) ? owner.FlowControl + : new FlowControlStatus { Mode = "ordinary", Writable = true }; + var json = JsonSerializer.SerializeToElement(status, StatusJson); + return Matches(step.GetProperty("match"), json); + }).ConfigureAwait(false); + break; + case "expectClosed": + await client.Completion.WaitAsync(TimeSpan.FromSeconds(1)).ConfigureAwait(false); + Assert.IsType(client.Error); + Assert.Equal("protocol", ((AhpTransportException)client.Error!).Kind); + break; + default: + throw new InvalidOperationException("Unknown corpus step: " + op); + } + } + } + catch (Exception error) + { + var state = string.Join(";", handles.Select(x => + { + lock (x.Value.Gate) + return x.Key + ":status=" + JsonSerializer.Serialize(x.Value.FlowControl, StatusJson) + + ":queue=" + string.Join(",", x.Value.Queue.Select(d => SystemTextJsonAhpSerializer.Default.Serialize(d.Notification))); + })); + throw new InvalidOperationException($"Corpus {Path.GetFileName(path)}, step {currentStep}; client={client.Error}; {state}", error); + } + finally + { + await client.ShutdownAsync().ConfigureAwait(false); + foreach (var request in requests) + try { await request.ConfigureAwait(false); } catch (AhpException) { } + } + } + + internal static async Task UntilAsync(Func condition) + { + using var timeout = new CancellationTokenSource(TimeSpan.FromSeconds(1)); + while (!condition()) await Task.Delay(1, timeout.Token).ConfigureAwait(false); + } + + internal static bool Matches(JsonElement expected, JsonElement actual) + { + if (expected.ValueKind != actual.ValueKind) return false; + if (expected.ValueKind == JsonValueKind.Object) + return expected.EnumerateObject().All(p => actual.TryGetProperty(p.Name, out var value) && Matches(p.Value, value)); + return AhpClient.JsonEqual(expected, actual); + } + + private static bool ContainsPath(JsonElement value, string path) + { + foreach (var part in path.Split('.')) + if (value.ValueKind != JsonValueKind.Object || !value.TryGetProperty(part, out value)) return false; + return true; + } +} + +internal sealed class CorpusTransport : ITransport, ITransportMessageLimits +{ + private readonly Channel _inbox = Channel.CreateUnbounded(); + private readonly List<(string Text, JsonElement Json)> _sent = new(); + public long? MaximumInboundMessageBytes { get; set; } + internal Action? OnSent { get; set; } + internal void Inject(string text) => _inbox.Writer.TryWrite(TransportMessage.FromText(text)); + internal (string Text, JsonElement Json)[] Snapshot() { lock (_sent) return _sent.ToArray(); } + public ValueTask SendAsync(TransportMessage message, CancellationToken cancellationToken = default) + { + using var document = JsonDocument.Parse(message.Text!); + var json = document.RootElement.Clone(); + lock (_sent) _sent.Add((message.Text!, json)); + OnSent?.Invoke(json); + return default; + } + public async ValueTask ReceiveAsync(CancellationToken cancellationToken = default) => + await _inbox.Reader.ReadAsync(cancellationToken).ConfigureAwait(false); + public ValueTask CloseAsync(CancellationToken cancellationToken = default) + { + _inbox.Writer.TryComplete(); + return default; + } + public ValueTask DisposeAsync() => CloseAsync(); +} diff --git a/clients/dotnet/tests/AgentHostProtocol.Tests/WindowedLifecycleTests.cs b/clients/dotnet/tests/AgentHostProtocol.Tests/WindowedLifecycleTests.cs new file mode 100644 index 000000000..ba65b77ec --- /dev/null +++ b/clients/dotnet/tests/AgentHostProtocol.Tests/WindowedLifecycleTests.cs @@ -0,0 +1,554 @@ +#nullable enable + +using System.Diagnostics.Metrics; +using System.Linq; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.AgentHostProtocol.Hosts; +using Microsoft.Extensions.Time.Testing; +using Xunit; + +namespace Microsoft.AgentHostProtocol.Tests; + +[CollectionDefinition("Window lifecycle", DisableParallelization = true)] +public sealed class WindowLifecycleIsolation { } + +[Collection("Window lifecycle")] +public sealed class WindowedLifecycleTests +{ + private const string Uri = "ahp-test:/lifecycle"; + private static readonly SystemTextJsonAhpSerializer Ser = SystemTextJsonAhpSerializer.Default; + + [Fact] + public async Task LateReconnectResponseAfterTimeoutIsCleanedUpWithoutClosing() => await CheckLateReconnectAsync(); + + private static async Task CheckLateReconnectAsync() + { + var transport = Peer(); + transport.OnSent = request => + { + if (Method(request) != "reconnect") Respond(transport, request); + }; + var time = new FakeTimeProvider(); + await using var client = AhpClient.Connect(transport, new ClientConfig + { + TimeProvider = time, + DefaultRequestTimeout = TimeSpan.FromSeconds(1), + }); + var reconnecting = client.ResumeAsync(new ReconnectParams + { + Channel = ProtocolVersion.RootResourceUri, + ClientId = "late-reconnect", + LastSeenServerSeq = 0, + Subscriptions = new List { Uri }, + Windows = new ChannelDeliveryResumeOptions + { + Items = new List + { + new() + { + Channel = Uri, + FlowControl = new SubscriptionFlowControlOptions + { + Receive = new ChannelReceiveLimits { WindowBytes = 262144, MaximumFrameBytes = 1048576 }, + }, + }, + }, + }, + }); + await WindowedDeliveryCorpusTests.UntilAsync(() => transport.Snapshot().Any(x => Method(x.Json) == "reconnect")).ConfigureAwait(false); + var request = transport.Snapshot().First(x => Method(x.Json) == "reconnect").Json; + time.Advance(TimeSpan.FromSeconds(1)); + await Assert.ThrowsAnyAsync(() => reconnecting).ConfigureAwait(false); + Assert.Null(client.GetWindowedSubscription(Uri)); + Respond(transport, request); + Notify(transport, "channel/snapshot", new ChannelSnapshotParams { Channel = Uri, Snapshot = Snapshot(Uri, 1) }); + await WindowedDeliveryCorpusTests.UntilAsync(() => transport.Snapshot().Any(x => Method(x.Json) == "unsubscribe")).ConfigureAwait(false); + await client.PingAsync().ConfigureAwait(false); + Assert.Equal(0, client.PendingWindowResponseCount); + Assert.False(client.Completion.IsCompleted); + } + + [Fact] + public async Task ClosedChannelHistoryRemainsBoundedWithoutDuplicateEntries() => await CheckHistoryAsync(); + + private static async Task CheckHistoryAsync() + { + var transport = Peer(); + await using var client = AhpClient.Connect(transport); + for (var i = 0; i < 1030; i++) + { + var (_, owner) = await client.SubscribeWindowedAsync("ahp-test:/history-" + i).ConfigureAwait(false); + await owner.UnsubscribeAsync().ConfigureAwait(false); + } + Assert.Equal(1024, client.ClosedWindowCount); + for (var i = 0; i < 5; i++) + { + var (_, owner) = await client.SubscribeWindowedAsync("ahp-test:/history-1029").ConfigureAwait(false); + await owner.UnsubscribeAsync().ConfigureAwait(false); + } + Assert.Equal(1024, client.ClosedWindowCount); + } + + [Fact] + public async Task UnsubscribeIgnoresInflightWindowTrafficAndCancelsCreditTimer() => await CheckUnsubscribeAsync(); + + private static async Task CheckUnsubscribeAsync() + { + var time = new FakeTimeProvider(); + var transport = Peer(); + await using var client = AhpClient.Connect(transport, new ClientConfig { TimeProvider = time }); + var (_, owner) = await client.SubscribeWindowedAsync(Uri, new WindowOptions { Receive = ReceiveMode.Manual }).ConfigureAwait(false); + Notify(transport, "channel/snapshot", new ChannelSnapshotParams { Channel = Uri, Snapshot = Snapshot(Uri, 1) }); + using var delivery = await owner.Deliveries.ReadAsync().ConfigureAwait(false); + delivery.Release(); + Assert.True(owner.CreditTimer); + await owner.UnsubscribeAsync().ConfigureAwait(false); + Assert.False(owner.CreditTimer); + time.Advance(TimeSpan.FromSeconds(1)); + Notify(transport, "channel/credit", new ChannelCreditParams { Channel = Uri, ConsumedBytes = 999999 }); + Notify(transport, "channel/frame", new ChannelFrameParams { Channel = Uri, Data = "{", Final = true }); + Notify(transport, "channel/snapshot", new ChannelSnapshotParams { Channel = Uri, Snapshot = Snapshot(Uri, 2) }); + await client.PingAsync().WaitAsync(TimeSpan.FromSeconds(1)).ConfigureAwait(false); + Assert.False(client.Completion.IsCompleted); + Assert.DoesNotContain(transport.Snapshot(), x => Method(x.Json) == "channel/credit"); + var (_, replacement) = await client.SubscribeWindowedAsync(Uri, new WindowOptions { Receive = ReceiveMode.Manual }).ConfigureAwait(false); + Assert.NotSame(owner, replacement); + Assert.Equal(0, replacement.Received); + Assert.Equal(0, replacement.Sent); + Assert.Equal(0, replacement.CreditSent); + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task UnsentSubscribeCancellationOrTimeoutDoesNotTrapPlainSends(bool timeout) => await CheckUnsentAsync(timeout); + + private static async Task CheckUnsentAsync(bool timeout) + { + var transport = new GatedTransport(Peer()); + var time = new FakeTimeProvider(); + await using var client = AhpClient.Connect(transport, new ClientConfig + { + TimeProvider = time, + DefaultRequestTimeout = TimeSpan.FromSeconds(1), + }); + var blocker = client.NotifyAsync("ordinary/blocked", new TcpEofParams { Channel = Uri }); + await transport.Entered.Task.WaitAsync(TimeSpan.FromSeconds(1)).ConfigureAwait(false); + using var cancellation = new CancellationTokenSource(); + var subscribing = client.SubscribeWindowedAsync(Uri, cancellationToken: timeout ? default : cancellation.Token); + await WindowedDeliveryCorpusTests.UntilAsync(() => client.PendingWindowResponseCount == 1).ConfigureAwait(false); + if (timeout) time.Advance(TimeSpan.FromSeconds(1)); + else cancellation.Cancel(); + await Assert.ThrowsAnyAsync(() => subscribing).ConfigureAwait(false); + Assert.Null(client.GetWindowedSubscription(Uri)); + Assert.Equal(0, client.PendingWindowResponseCount); + transport.Release.TrySetResult(true); + await blocker.ConfigureAwait(false); + await client.DispatchAsync(Uri, new StateAction(new SessionTitleChangedAction + { + Type = ActionType.SessionTitleChanged, + Title = "ordinary", + })).WaitAsync(TimeSpan.FromSeconds(1)).ConfigureAwait(false); + Assert.DoesNotContain(transport.Inner.Snapshot(), x => Method(x.Json) == "subscribe"); + Assert.False(client.Completion.IsCompleted); + } + + [Theory] + [InlineData("cancel")] + [InlineData("timeout")] + [InlineData("unsubscribe")] + [InlineData("rpc-error")] + public async Task AbandonedSubscribeAndLateSnapshotDoNotCloseConnection(string reason) => await CheckAbandonedAsync(reason); + + private static async Task CheckAbandonedAsync(string reason) + { + var transport = Peer(); + var time = new FakeTimeProvider(); + transport.OnSent = request => + { + if (Method(request) != "subscribe") Respond(transport, request); + }; + await using var client = AhpClient.Connect(transport, new ClientConfig + { + TimeProvider = time, + DefaultRequestTimeout = TimeSpan.FromSeconds(1), + }); + using var cancellation = new CancellationTokenSource(); + var subscribing = client.SubscribeWindowedAsync(Uri, cancellationToken: cancellation.Token); + await WindowedDeliveryCorpusTests.UntilAsync(() => transport.Snapshot().Any(x => Method(x.Json) == "subscribe")).ConfigureAwait(false); + var request = transport.Snapshot().First(x => Method(x.Json) == "subscribe").Json; + if (reason == "cancel") cancellation.Cancel(); + else if (reason == "timeout") time.Advance(TimeSpan.FromSeconds(1)); + else if (reason == "unsubscribe") await client.UnsubscribeAsync(Uri).ConfigureAwait(false); + else ReplyError(transport, request); + await Assert.ThrowsAnyAsync(() => subscribing).ConfigureAwait(false); + Assert.Null(client.GetWindowedSubscription(Uri)); + if (reason != "rpc-error") + { + Respond(transport, request); + Notify(transport, "channel/snapshot", new ChannelSnapshotParams { Channel = Uri, Snapshot = Snapshot(Uri, 1) }); + await WindowedDeliveryCorpusTests.UntilAsync(() => transport.Snapshot().Any(x => Method(x.Json) == "unsubscribe")).ConfigureAwait(false); + } + await client.NotifyAsync("ordinary/data", new TcpEofParams { Channel = Uri }).WaitAsync(TimeSpan.FromSeconds(1)).ConfigureAwait(false); + await client.PingAsync().WaitAsync(TimeSpan.FromSeconds(1)).ConfigureAwait(false); + Assert.Equal(0, client.PendingWindowResponseCount); + Assert.False(client.Completion.IsCompleted); + } + + [Fact] + public async Task LargeSingleSendAndControlsBypassTinySendBudget() => await CheckLargeAndControlsAsync(); + + private static async Task CheckLargeAndControlsAsync() + { + var transport = Peer(); + transport.OnSent = request => + { + if (Method(request) == "subscribe") + { + var parameters = Ser.Deserialize(request.GetProperty("params")); + Reply(transport, request, new SubscribeResult + { + FlowControl = new ChannelFlowControl + { + ClientReceive = parameters.FlowControl!.Receive, + HostReceive = new ChannelReceiveLimits { WindowBytes = 1, MaximumFrameBytes = 1024 }, + }, + }); + } + else Respond(transport, request); + }; + await using var client = AhpClient.Connect(transport, new ClientConfig + { + Windowing = new WindowingConfig { MaxBufferedSendBytes = 1 }, + }); + var (_, owner) = await client.SubscribeWindowedAsync(Uri, new WindowOptions { Receive = ReceiveMode.Manual }).ConfigureAwait(false); + await owner.WaitWritableAsync().ConfigureAwait(false); + var sending = owner.NotifyAsync("future/large", new TcpDataParams { Channel = Uri, Data = new string('x', 2048) }); + await WindowedDeliveryCorpusTests.UntilAsync(() => transport.Snapshot().Any(x => Method(x.Json) == "channel/frame")).ConfigureAwait(false); + Assert.False(sending.IsCompleted); + await client.PingAsync().WaitAsync(TimeSpan.FromSeconds(1)).ConfigureAwait(false); + Notify(transport, "channel/snapshot", new ChannelSnapshotParams { Channel = Uri, Snapshot = Snapshot(Uri, 1) }); + using var snapshot = await owner.Deliveries.ReadAsync().ConfigureAwait(false); + snapshot.Release(); + await WindowedDeliveryCorpusTests.UntilAsync(() => transport.Snapshot().Any(x => Method(x.Json) == "channel/credit")).ConfigureAwait(false); + await owner.NotifyAsync("tcp/reset", new TcpResetParams { Channel = Uri }).WaitAsync(TimeSpan.FromSeconds(1)).ConfigureAwait(false); + await owner.UnsubscribeAsync().WaitAsync(TimeSpan.FromSeconds(1)).ConfigureAwait(false); + await Assert.ThrowsAsync(() => sending).ConfigureAwait(false); + Assert.False(client.Completion.IsCompleted); + } + + [Theory] + [InlineData(false, false)] + [InlineData(false, true)] + [InlineData(true, false)] + [InlineData(true, true)] + public async Task ManagedReadersEndOnRemoveAndShutdown(bool remove, bool manual) => await CheckManagedEndAsync(remove, manual); + + private static async Task CheckManagedEndAsync(bool remove, bool manual) + { + await using var multi = new MultiHostClient(); + await multi.AddHostAsync(Config(_ => Peer())).ConfigureAwait(false); + var (_, owner) = await multi.SubscribeWindowedAsync("lifecycle", Uri, + new WindowOptions { Receive = manual ? ReceiveMode.Manual : ReceiveMode.Auto }).ConfigureAwait(false); + var observer = multi.EventsForHost("lifecycle", Uri); + await using var iterator = owner.ReadEventsAsync().GetAsyncEnumerator(); + var waiting = manual ? owner.Deliveries.WaitToReadAsync().AsTask() : iterator.MoveNextAsync().AsTask(); + if (remove) await multi.RemoveHostAsync("lifecycle").WaitAsync(TimeSpan.FromSeconds(2)).ConfigureAwait(false); + else await multi.ShutdownAsync().WaitAsync(TimeSpan.FromSeconds(2)).ConfigureAwait(false); + await Assert.ThrowsAsync(() => waiting).ConfigureAwait(false); + await observer.Completion.WaitAsync(TimeSpan.FromSeconds(1)).ConfigureAwait(false); + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task ClosedManagedOwnerIsNeverResumed(bool synchronous) => await CheckManagedCloseAsync(synchronous); + + private static async Task CheckManagedCloseAsync(bool synchronous) + { + var transports = new List(); + await using var multi = new MultiHostClient(); + await multi.AddHostAsync(Config(_ => + { + var transport = Peer(); + transports.Add(transport); + return transport; + })).ConfigureAwait(false); + var (_, owner) = await multi.SubscribeWindowedAsync("lifecycle", Uri).ConfigureAwait(false); + if (synchronous) owner.Dispose(); + else await owner.UnsubscribeAsync().ConfigureAwait(false); + Assert.DoesNotContain(Uri, multi.Host("lifecycle")!.Subscriptions); + var generation = multi.Host("lifecycle")!.Generation; + await multi.ReconnectAsync("lifecycle").ConfigureAwait(false); + await WindowedDeliveryCorpusTests.UntilAsync(() => multi.Host("lifecycle")!.Generation > generation).ConfigureAwait(false); + var request = transports[1].Snapshot().First(x => Method(x.Json) == "reconnect").Json; + Assert.DoesNotContain(request.GetProperty("params").GetProperty("subscriptions").EnumerateArray(), x => x.GetString() == Uri); + Assert.False(request.GetProperty("params").TryGetProperty("windows", out _)); + Assert.True(owner.Closed); + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task FreshResubscribeFailureIsPerChannelAndShutdownCanCancelIt(bool stall) => await CheckFreshFailureAsync(stall); + + private static async Task CheckFreshFailureAsync(bool stall) + { + var transports = new List(); + await using var multi = new MultiHostClient(); + await multi.AddHostAsync(Config(index => + { + var transport = Peer(); + transports.Add(transport); + if (index > 0) + transport.OnSent = request => + { + if (Method(request) == "reconnect") ReplyError(transport, request); + else if (Method(request) == "subscribe" && request.GetProperty("params").GetProperty("channel").GetString() == Uri) + { + if (!stall) ReplyError(transport, request); + } + else Respond(transport, request); + }; + return transport; + })).ConfigureAwait(false); + var (_, failed) = await multi.SubscribeWindowedAsync("lifecycle", Uri).ConfigureAwait(false); + var (_, kept) = await multi.SubscribeWindowedAsync("lifecycle", "ahp-test:/kept").ConfigureAwait(false); + var generation = multi.Host("lifecycle")!.Generation; + await multi.ReconnectAsync("lifecycle").ConfigureAwait(false); + await WindowedDeliveryCorpusTests.UntilAsync(() => transports.Count > 1 + && transports[1].Snapshot().Any(x => Method(x.Json) == "subscribe")).ConfigureAwait(false); + if (stall) + { + await multi.ShutdownAsync().WaitAsync(TimeSpan.FromSeconds(2)).ConfigureAwait(false); + Assert.True(failed.Closed); + Assert.True(kept.Closed); + } + else + { + await WindowedDeliveryCorpusTests.UntilAsync(() => multi.Host("lifecycle")!.Generation > generation).ConfigureAwait(false); + Assert.True(failed.Closed); + Assert.IsType(failed.CloseError); + Assert.False(kept.Closed); + Assert.DoesNotContain(Uri, multi.Host("lifecycle")!.Subscriptions); + Assert.Equal(HostStateKind.Connected, multi.Host("lifecycle")!.State.Kind); + } + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task OrdinaryFallbackRecoveryPrecedesBackToBackLiveEvents(bool snapshot) => await CheckOrdinaryRecoveryAsync(snapshot); + + private static async Task CheckOrdinaryRecoveryAsync(bool snapshot) + { + await using var multi = new MultiHostClient(); + await multi.AddHostAsync(Config(index => + { + var transport = Peer(); + transport.OnSent = request => + { + if (Method(request) == "subscribe") + Reply(transport, request, new SubscribeResult { Snapshot = Snapshot(Uri, 1) }); + else if (Method(request) == "reconnect" && index > 0) + { + if (snapshot) + Reply(transport, request, new ReconnectSnapshotResult + { + Type = ReconnectResultType.Snapshot, + Snapshots = new List { Snapshot(Uri, 2) }, + }); + else + Reply(transport, request, new ReconnectReplayResult + { + Type = ReconnectResultType.Replay, + Actions = new List { Action(2) }, + Missing = new List(), + }); + Notify(transport, "action", Action(3)); + } + else Respond(transport, request); + }; + return transport; + })).ConfigureAwait(false); + var (_, owner) = await multi.SubscribeWindowedAsync("lifecycle", Uri, new WindowOptions { Receive = ReceiveMode.Manual }).ConfigureAwait(false); + using var initial = await owner.Deliveries.ReadAsync().ConfigureAwait(false); + var observer = multi.EventsForHost("lifecycle", Uri); + var generation = multi.Host("lifecycle")!.Generation; + await multi.ReconnectAsync("lifecycle").ConfigureAwait(false); + await WindowedDeliveryCorpusTests.UntilAsync(() => multi.Host("lifecycle")!.Generation > generation).ConfigureAwait(false); + using var first = await owner.Deliveries.ReadAsync().ConfigureAwait(false); + using var second = await owner.Deliveries.ReadAsync().ConfigureAwait(false); + Assert.Equal(2, Sequence(first.Event)); + Assert.Equal(3, Sequence(second.Event)); + Assert.Equal(2, Sequence(await observer.ReadAsync().ConfigureAwait(false))); + Assert.Equal(3, Sequence(await observer.ReadAsync().ConfigureAwait(false))); + Assert.Equal(3, multi.Host("lifecycle")!.ServerSeq); + } + + [Fact] + public async Task RegistryIsPrunedWithoutGaugeCollectionAndReconnectMetricsRemain() => await CheckTelemetryAsync(); + + private static async Task CheckTelemetryAsync() + { + var baseline = AhpTelemetry.TrackedWindowCount; + var transport = Peer(); + await using (var client = AhpClient.Connect(transport)) + for (var i = 0; i < 40; i++) + { + var (_, owner) = await client.SubscribeWindowedAsync("ahp-test:/registry-" + i).ConfigureAwait(false); + await owner.UnsubscribeAsync().ConfigureAwait(false); + } + Assert.True(AhpTelemetry.TrackedWindowCount <= baseline); + long reconnects = 0; + using var listener = new MeterListener + { + InstrumentPublished = (instrument, meter) => + { + if (instrument.Name == "ahp.client.reconnects") meter.EnableMeasurementEvents(instrument); + }, + }; + listener.SetMeasurementEventCallback((_, value, _, _) => Interlocked.Add(ref reconnects, value)); + listener.Start(); + await using var multi = new MultiHostClient(); + await multi.AddHostAsync(Config(_ => Peer())).ConfigureAwait(false); + var generation = multi.Host("lifecycle")!.Generation; + await multi.ReconnectAsync("lifecycle").ConfigureAwait(false); + await WindowedDeliveryCorpusTests.UntilAsync(() => multi.Host("lifecycle")!.Generation > generation && Volatile.Read(ref reconnects) >= 2).ConfigureAwait(false); + Assert.Equal(2, Volatile.Read(ref reconnects)); + } + + [Theory] + [InlineData("notify")] + [InlineData("dispatchAction")] + [InlineData("unsubscribe")] + public async Task OrdinarySendsDrainBeforeShutdownAndObserversEnd(string method) => await CheckOrdinaryShutdownAsync(method); + + private static async Task CheckOrdinaryShutdownAsync(string method) + { + var transport = Peer(); + await using var client = AhpClient.Connect(transport); + using var events = client.CreateEventStream(); + var sent = method switch + { + "dispatchAction" => client.DispatchAsync(Uri, new StateAction(new SessionTitleChangedAction { Type = ActionType.SessionTitleChanged, Title = "sent" })), + "unsubscribe" => client.UnsubscribeAsync(Uri), + _ => client.NotifyAsync("notify", new TcpEofParams { Channel = Uri }), + }; + await client.ShutdownAsync().WaitAsync(TimeSpan.FromSeconds(2)).ConfigureAwait(false); + await sent.ConfigureAwait(false); + Assert.Contains(transport.Snapshot(), x => Method(x.Json) == method); + await events.Events.Completion.WaitAsync(TimeSpan.FromSeconds(1)).ConfigureAwait(false); + } + + private static HostConfig Config(Func factory) + { + var index = -1; + return new HostConfig + { + Id = "lifecycle", + ReconnectPolicy = ReconnectPolicy.Disabled, + ClientConfig = new ClientConfig { DefaultRequestTimeout = TimeSpan.Zero }, + TransportFactory = (_, _) => Task.FromResult(factory(Interlocked.Increment(ref index))), + }; + } + + private static CorpusTransport Peer() + { + var transport = new CorpusTransport(); + transport.OnSent = request => Respond(transport, request); + return transport; + } + + private static void Respond(CorpusTransport transport, JsonElement request) + { + if (!request.TryGetProperty("id", out _)) return; + object result; + if (Method(request) == "initialize") + result = new InitializeResult { ProtocolVersion = ProtocolVersion.Supported[0], Snapshots = new List(), ServerSeq = 0 }; + else if (Method(request) == "subscribe") + { + var parameters = Ser.Deserialize(request.GetProperty("params")); + result = new SubscribeResult + { + FlowControl = parameters.FlowControl is { } flow + ? new ChannelFlowControl { ClientReceive = flow.Receive, HostReceive = flow.Receive } : null, + }; + } + else if (Method(request) == "reconnect") + { + var parameters = Ser.Deserialize(request.GetProperty("params")); + result = new ReconnectReplayResult + { + Type = ReconnectResultType.Replay, + Actions = new List(), + Missing = new List(), + Windows = parameters.Windows is null ? null : new ChannelDeliveryResumeResult + { + Items = parameters.Windows.Items.Select(x => new ResumedChannelSubscription + { + Channel = x.Channel, + FlowControl = new ChannelFlowControl { ClientReceive = x.FlowControl.Receive, HostReceive = x.FlowControl.Receive }, + Recovery = ChannelRecoveryKind.Live, + }).ToList(), + }, + }; + } + else result = new Dictionary(); + Reply(transport, request, result); + } + + private static void Reply(CorpusTransport transport, JsonElement request, object result) => + transport.Inject(Ser.EncodeMessage(new JsonRpcMessage + { + SuccessResponse = new JsonRpcSuccessResponse { Id = request.GetProperty("id").GetUInt64(), Result = Ser.SerializeToElement(result) }, + }).Text!); + + private static void ReplyError(CorpusTransport transport, JsonElement request) => + transport.Inject(Ser.EncodeMessage(new JsonRpcMessage + { + ErrorResponse = new JsonRpcErrorResponse + { + Id = request.GetProperty("id").GetUInt64(), + Error = new JsonRpcErrorObject { Code = -32601, Message = "not available" }, + }, + }).Text!); + + private static void Notify(CorpusTransport transport, string method, T parameters) => + transport.Inject(Ser.EncodeMessage(new JsonRpcMessage + { + Notification = new JsonRpcNotification { Method = method, Params = Ser.SerializeToElement(parameters) }, + }).Text!); + + private static string? Method(JsonElement message) => message.GetProperty("method").GetString(); + private static Snapshot Snapshot(string uri, long sequence) => new() { Resource = uri, State = new SnapshotState { Empty = new EmptyState() }, FromSeq = sequence }; + private static ActionEnvelope Action(long sequence) => new() + { + Channel = Uri, + ServerSeq = sequence, + Action = new StateAction(new SessionTitleChangedAction { Type = ActionType.SessionTitleChanged, Title = sequence.ToString(System.Globalization.CultureInfo.InvariantCulture) }), + }; + private static long Sequence(SubscriptionEvent value) => value switch + { + SubscriptionEventAction action => action.Envelope.ServerSeq, + SubscriptionEventSnapshot snapshot => snapshot.Snapshot.FromSeq, + _ => throw new InvalidOperationException("Unexpected event."), + }; + + private sealed class GatedTransport(CorpusTransport inner) : ITransport + { + internal CorpusTransport Inner => inner; + internal readonly TaskCompletionSource Entered = WindowedSubscription.Signal(); + internal readonly TaskCompletionSource Release = WindowedSubscription.Signal(); + public async ValueTask SendAsync(TransportMessage message, CancellationToken cancellationToken = default) + { + Entered.TrySetResult(true); + await Release.Task.WaitAsync(cancellationToken).ConfigureAwait(false); + await inner.SendAsync(message, cancellationToken).ConfigureAwait(false); + } + public ValueTask ReceiveAsync(CancellationToken cancellationToken = default) => inner.ReceiveAsync(cancellationToken); + public ValueTask CloseAsync(CancellationToken cancellationToken = default) => inner.CloseAsync(cancellationToken); + public ValueTask DisposeAsync() => inner.DisposeAsync(); + } +} diff --git a/clients/go/README.md b/clients/go/README.md index 13562b870..b1c68ee97 100644 --- a/clients/go/README.md +++ b/clients/go/README.md @@ -60,6 +60,238 @@ for evt := range sub.Events() { } ``` +## Windowed delivery + +Windowing is **off by default**. `Config.Windowing` supplies connection defaults; +`SubscribeWithOptions` opts one subscription in. Sizes are UTF-16 bytes of the +actual JSON text, including envelopes and escaping, not payload/UTF-8 sizes. + +```go +cfg := ahp.DefaultConfig() +cfg.Windowing.WindowBytes = 256 << 10 +cfg.Windowing.MaxBufferedSendBytes = 1 << 20 // per subscription +// Optional: override the transport-derived maximum based on your transport. +cfg.Windowing.MaximumFrameBytes = 16 << 10 + +client, err := ahp.Connect(ctx, transport, cfg) +if err != nil { return err } +defer client.Shutdown(ctx) + +if _, err := client.Initialize(ctx, "client-id", + ahptypes.SupportedProtocolVersions(), nil); err != nil { return err } + +result, sub, err := client.SubscribeWithOptions(ctx, "ahp-terminal:/t1", + ahp.SubscribeOptions{Window: &ahp.WindowOptions{}}) +if err != nil { return err } +defer sub.Close() +_ = result +for { + delivery, err := sub.Next(ctx) + if err != nil { break } + switch event := delivery.Event().(type) { + case ahp.SubscriptionEventSnapshot: + // Replace the local state using event.Params.Snapshot. + _ = event + case ahp.SubscriptionEventAction: + // Apply/render event.Envelope, completing the work before Next. + _ = event + } +} +``` + +`ReceiveAuto` (default) releases the **previous** delivery when `Next` is called +again, even if that call begins waiting or its context is canceled. Thus work +performed between calls holds credit. Call `sub.Close()` if you leave the loop. +`Events()`, `Client.Events()` and host subscription event streams remain lossy +observers: reading them **does not release credit**. Use `Next` on an +options-based subscription for credit-owning consumption. Legacy subscription +APIs keep their existing buffered event-channel behavior. + +For manual acknowledgement, request `ReceiveManual` and call `Release` only +when application processing, or writing to a bounded downstream sink, completes: + +```go +_, sub, err := client.SubscribeWithOptions(ctx, uri, ahp.SubscribeOptions{ + Window: &ahp.WindowOptions{Receive: ahp.ReceiveManual}, +}) +if err != nil { return err } +defer sub.Close() +for { + delivery, err := sub.Next(ctx) + if err != nil { return err } + if err := process(delivery.Event()); err != nil { + delivery.Release() // Rejected/discarded events can also be released. + return err + } + delivery.Release() // Idempotent; may be called from another goroutine. +} +``` + +The first delivery is the snapshot. When a host declines windowing, its result +snapshot is injected as an uncharged delivery; `Release` is then a no-op and +ordinary buffering applies. There is no protocol-version gate and no change to +version offers. Inspect `sub.FlowControl().Mode` and `.FallbackReason` +(`declined` or `transport`) to distinguish fallback. TCP subscriptions require +acceptance, cannot be initial subscriptions, and never resume automatically. +TCP socket bridging is application-owned. + +Repeated `SubscribeWithOptions` for the same channel updates the **same owner** +handle; `sub.Update(ctx, options)` is equivalent. Counters continue for +windowed-to-windowed updates. Limits, snapshot views and delivery preferences +are independent; the SDK never adds `maxLatencyMs` automatically. + +### Examples + +Two runnable examples use an in-memory demo host that follows the window rules, +so they need no server: + +- [`examples/windowing_basic`](examples/windowing_basic/main.go) opts a terminal + subscription into windowed delivery and shows the host pausing when the + window is full and resuming as the client returns credit. +- [`examples/windowing_backpressure`](examples/windowing_backpressure/main.go) + adds a slow renderer, a custom coalescer that keeps only the latest terminal + resize, and a paste that waits for `WaitWritable` instead of queueing. + +```sh +go run ./examples/windowing_basic +go run ./examples/windowing_backpressure +``` + +The demo host and its in-memory transport live in +[`examples/internal/demohost`](examples/internal/demohost); they implement the +public `ahp.Transport` and `ahp.TransportCapabilitiesProvider` interfaces. + +### Send-side backpressure and observability + +`Notify` and `Dispatch` await the actual transport write, including credit and +fragment scheduling. `ErrBackpressure` means data queued behind another send +has reached `MaxBufferedSendBytes` (default 1 MiB). One head message is always +accepted when the queue is empty, even if larger than this budget; fragmentation +and the window pace it. The budget also bounds retained original inputs/receipts +while coalesced sends wait for credit. `WaitWritable` is advisory, +not a reservation. Once fragmentation starts, cancellation does not abandon +the message; its remaining fragments are completed. Unsent canceled work is +removed. Pending sends fail on connection loss and are never replayed. + +```go +if err := sub.WaitWritable(ctx); err != nil { return err } +status := sub.FlowControl() +fmt.Println(status.OutstandingSendBytes, status.BufferedSendBytes) +if err := client.Notify(ctx, "tcp/data", ahptypes.TcpDataParams{ + Channel: sub.URI(), + Data: base64.StdEncoding.EncodeToString(bytesToSend), +}); err != nil { + if errors.Is(err, ahp.ErrBackpressure) { + // Stop reading the source and retry after WaitWritable. + } + return err +} +``` + +`sub.FlowControlChanges()` is a conflated change stream. It closes when the +owner is explicitly closed or dropped; managed owners keep it across connection +replacement. `client.FlowControlEvents()` is a lossy connection-level stream +of negotiation, blocked/unblocked, credit, fragmentation and coalescing events. +Create it before subscribing if you want to observe negotiation. + +### Coalescing + +Adjacent unread `terminal/data` and compatible `chat/delta` actions merge by +default. Adjacent unsent `tcp/data` notifications merge by decoding the base64 +and concatenating bytes. Origins/rejections, distinct `dispatchAction` sends, +snapshots, EOF/reset, different parts or metadata, and connection epochs are +never crossed. Receive merging is local and does not require delivery +preferences. Original received sizes are summed; merged send text is measured +as written. Coalescing never delays sendable data to collect a batch. + +Disable merging per subscription with +`disabled := false; WindowOptions{Coalescing: &disabled}`. Register custom +coalescers in connection configuration, before `Connect`; rules run before +built-ins and must be nonblocking and must not call back into the client: + +```go +cfg := ahp.DefaultConfig() +cfg.Windowing.Coalescers = []ahp.Coalescer{{ + Method: "my-app/chunk", + Merge: func(a, b ahptypes.JsonRpcNotification, c ahp.CoalescingContext) ( + ahptypes.JsonRpcNotification, bool, + ) { + if c.Direction != ahp.CoalesceReceive { return a, false } + var earlier, later struct { + Channel string `json:"channel"` + Text string `json:"text"` + } + if json.Unmarshal(a.Params, &earlier) != nil || + json.Unmarshal(b.Params, &later) != nil { + return a, false + } + earlier.Text += later.Text + a.Params, _ = json.Marshal(earlier) + return a, true + }, +}} +``` + +Rules must preserve required events, ordering, method and channel. The SDK +checks channel/method identity and excludes control messages and reconciliation +barriers. Known AHP events retain typed models; custom/unknown notifications +are available through `SubscriptionEventNotification` and +`Delivery.Notification()`. + +### Transports and managed hosts + +Custom transports remain source-compatible: `Transport` has no new methods. +For windowing, optionally implement `TransportCapabilitiesProvider`, asserting +that text is retained on receive and sent without reserialization. Return text/ +binary `TransportMessage`s, or `NewParsedMessageWithRawText(parsed, original)`; +`NewParsedMessage` alone cannot supply exact received sizes and uses ordinary +fallback. `RawJSONText()` never reserializes. Boxing forwards capabilities. + +`ahpws.Connect` sets coder/websocket's read limit to 32 MiB by default (the +dependency's native default is only 32 KiB). Its default AHP maximum is +**22,369,620 UTF-16 bytes**, computed as `2 * floor(readLimit / 3)`. +Custom `DialOptions.ReadLimit` is converted the same way. `FromConn` cannot +inspect an externally configured read limit, so its default is **1 MiB**; +set `Config.Windowing.MaximumFrameBytes` when you know the actual limit. +`ahp.MaximumFrameBytesForReadLimit` provides that conversion for adapters. + +For multiple hosts, configure `HostConfig.ClientConfig.Windowing` and call +`multi.SubscribeWithOptions(ctx, hostID, uri, options)`. The managed handle +provides `Next`, `Update`, `FlowControl`, `FlowControlChanges`, `WaitWritable`, +and `Close(ctx)`. Hosts retain unread items ahead of the replacement snapshot +and keep re-running `initialize`, then re-subscribe with fresh offers. Releases +of old deliveries never credit the new connection. Missing permissions or +closed TCP channels produce a visible dropped-subscription error. +This is a deliberate Go policy: managed hosts do **not** use protocol +`reconnect.windows` or replay/live recovery; they receive fresh snapshots. +A failed re-subscribe drops only that owner, not healthy sibling subscriptions +or the host connection. Closing/dropping an owner removes it from the managed +registry, so subsequent reconnects do not retry it. +`multi.FlowControlEvents()` tags observe-only telemetry with its host ID. +State mirrors remain passive: apply/store a delivery before releasing it. + +Single-connection callers can use `ReconnectWithParams(ctx, +ahptypes.ReconnectParams{Windows: ...})` to exercise protocol window recovery; +the existing `Reconnect` convenience signature is unchanged. Obtain resumed +owners with `sub, ok := client.Subscription(uri)`, then consume them through +`Next` as usual. Raw `Request` +paths for subscribe/reconnect use the same negotiation hook, installed before +the next inbound data message. `ResumeSubscription` is the managed-host +initialize/re-subscribe helper, not a protocol `reconnect`. + +Closing or unsubscribing an options-based owner stops its credit timers and +invalidates held deliveries; their later `Release` calls send no credit. +A new subscribe on that URI gets fresh counters. Canceled/timed-out +negotiations are retired and unsubscribed without closing the connection; +late responses cannot revive them. Failed RPC updates retain the accepted +owner and its previous limits. + +Windowing intentionally adds no receive/reassembly memory cap or retention +timeout. Partial non-final frames are credited when stored, so messages larger +than the window can complete; the final frame stays charged until release. +Oversized data, invalid fragments/credits, or raised accepted limits close the +connection with `TransportError{Kind: "protocol"}`. + ## Code generation The contents of `ahptypes/*.go` (except `common.go`) are auto-generated diff --git a/clients/go/ahp/client.go b/clients/go/ahp/client.go index 1938bf369..59d39ec79 100644 --- a/clients/go/ahp/client.go +++ b/clients/go/ahp/client.go @@ -24,6 +24,9 @@ type Config struct { // channel. Slow consumers that fail to drain within this many // messages will drop frames and must resubscribe. SubscriptionBuffer int + + // Windowing supplies defaults, but does not opt any subscription in. + Windowing WindowingConfig } // DefaultConfig returns a [Config] with sensible defaults: 30s request @@ -32,6 +35,7 @@ func DefaultConfig() Config { return Config{ DefaultRequestTimeout: 30 * time.Second, SubscriptionBuffer: 256, + Windowing: WindowingConfig{WindowBytes: 256 << 10, MaxBufferedSendBytes: 1 << 20}, } } @@ -99,6 +103,7 @@ type Subscription struct { events chan SubscriptionEvent closeMu sync.Mutex closed bool + owner *deliveryQueue } // URI returns the channel URI this subscription is bound to. @@ -107,18 +112,38 @@ func (s *Subscription) URI() string { return s.uri } // Events returns a receive-only channel of events. The channel closes // when the [Client] shuts down or the caller invokes // [Subscription.Close]. -func (s *Subscription) Events() <-chan SubscriptionEvent { return s.events } +func (s *Subscription) Events() <-chan SubscriptionEvent { + s.closeMu.Lock() + defer s.closeMu.Unlock() + return s.events +} -// Close stops the subscription locally without notifying the server. +// Close stops this handle. Options-based owners also queue unsubscribe and +// retire their credit epoch; legacy observer handles only close locally. // Safe to call multiple times. func (s *Subscription) Close() { + s.closeEvents() + if s.owner != nil { + q := s.owner + q.mu.Lock() + q.finishLocked(ErrClosed) + e, f := q.engine, q.flow + q.mu.Unlock() + e.mu.Lock() + e.retireLocked(f, ErrClosed, true) + e.mu.Unlock() + } +} + +func (s *Subscription) closeEvents() { s.closeMu.Lock() - defer s.closeMu.Unlock() if s.closed { + s.closeMu.Unlock() return } s.closed = true close(s.events) + s.closeMu.Unlock() } func (s *Subscription) trySend(ev SubscriptionEvent) { @@ -147,15 +172,17 @@ type DispatchHandle struct { // pendingResult routes a single in-flight request's outcome to its // waiter. type pendingResult struct { - value json.RawMessage - err *ahptypes.JsonRpcError + value json.RawMessage + err *ahptypes.JsonRpcError + localErr error } // outboundMsg is the writer goroutine's input queue payload. type outboundMsg struct { msg ahptypes.JsonRpcMessage // done, if non-nil, receives the result of the send attempt. - done chan error + done chan error + hooks *requestHooks } // EventStream is a top-level fan-in receiver over every inbound event @@ -228,8 +255,10 @@ type Client struct { outbound chan outboundMsg // pending is the request-correlation map keyed by JSON-RPC id. - pendingMu sync.Mutex - pending map[uint64]chan pendingResult + pendingMu sync.Mutex + pending map[uint64]chan pendingResult + pendingHooks map[uint64]*requestHooks + flow *flowEngine // subscriptionsMu guards subscriptions and the all-events // fan-out registry. @@ -269,9 +298,15 @@ func Connect(_ context.Context, transport Transport, cfg Config) (*Client, error transport: transport, outbound: make(chan outboundMsg, 64), pending: make(map[uint64]chan pendingResult), + pendingHooks: make(map[uint64]*requestHooks), subscriptions: make(map[string][]*Subscription), done: make(chan struct{}), } + var err error + c.flow, err = newFlowEngine(c) + if err != nil { + return nil, err + } c.nextID.Store(1) c.nextClientSeq.Store(1) c.wg.Add(2) @@ -322,6 +357,7 @@ func (c *Client) shutdownWithError(err error) { c.closeErr.Store(err) } close(c.done) + c.flow.stop() // Close the transport so any blocked Recv unblocks. closeCtx, cancel := context.WithTimeout(context.Background(), 2*time.Second) @@ -341,6 +377,7 @@ func (c *Client) shutdownWithError(err error) { } close(ch) delete(c.pending, id) + delete(c.pendingHooks, id) } c.pendingMu.Unlock() @@ -353,7 +390,14 @@ func (c *Client) shutdownWithError(err error) { c.subscriptionsMu.Unlock() for _, list := range subs { for _, s := range list { - s.Close() + s.closeEvents() + if s.owner != nil { + reason := err + if reason == nil { + reason = ErrShutdown + } + s.owner.close(reason) + } } } for _, l := range listeners { @@ -365,36 +409,7 @@ func (c *Client) shutdownWithError(err error) { // runWriter drains outbound messages and ships them via the transport. // Exits when c.done is closed. func (c *Client) runWriter() { - defer c.wg.Done() - for { - select { - case <-c.done: - return - case msg, ok := <-c.outbound: - if !ok { - return - } - wire, err := EncodeMessage(msg.msg) - if err != nil { - if msg.done != nil { - msg.done <- err - } - continue - } - ctx, cancel := contextWithDone(context.Background(), c.done) - sendErr := c.transport.Send(ctx, wire) - cancel() - if msg.done != nil { - msg.done <- sendErr - } - if sendErr != nil { - // Transport faulted — trigger global shutdown so all - // consumers see EOF in bounded time. - c.shutdownWithError(fmt.Errorf("ahp: transport send: %w", sendErr)) - return - } - } - } + c.flow.runWriter() } // runReader pumps inbound messages from the transport and dispatches @@ -415,12 +430,31 @@ func (c *Client) runReader() { c.shutdownWithError(fmt.Errorf("ahp: transport recv: %w", err)) return } + select { + case <-c.done: + return + default: + } parsed, perr := msg.IntoParsed() if perr != nil { // Skip malformed frames; protocol resync is the server's // responsibility. + if c.flow.hasWindows() { + c.shutdownWithError(&TransportError{Kind: "protocol", Err: perr}) + return + } continue } + if parsed.Notification != nil { + handled, err := c.flow.receive(*parsed.Notification, msg) + if err != nil { + c.shutdownWithError(&TransportError{Kind: "protocol", Err: err}) + return + } + if handled { + continue + } + } c.dispatch(parsed) } } @@ -462,13 +496,27 @@ func (c *Client) dispatch(msg ahptypes.JsonRpcMessage) { func (c *Client) deliver(id uint64, r pendingResult) { c.pendingMu.Lock() ch, ok := c.pending[id] + hook := c.pendingHooks[id] if ok { delete(c.pending, id) + delete(c.pendingHooks, id) } c.pendingMu.Unlock() if !ok { return } + if r.err != nil && hook != nil { + hook.failed(&RPCError{Code: r.err.Code, Message: r.err.Message, Data: r.err.Data}) + } + if r.err == nil && hook != nil { + if err := hook.complete(r.value); err != nil { + r.localErr = &TransportError{Kind: "protocol", Err: err} + ch <- r + close(ch) + c.shutdownWithError(r.localErr) + return + } + } ch <- r close(ch) } @@ -638,6 +686,20 @@ func (c *Client) handleServerRequest(req ahptypes.JsonRpcRequest) { // Request sends a JSON-RPC request and decodes the response into out. // If out is nil, the result is discarded. func (c *Client) Request(ctx context.Context, method string, params any, out any) error { + return c.request(ctx, method, params, out, nil) +} + +func (c *Client) request(ctx context.Context, method string, params any, out any, opts *SubscribeOptions) error { + if err := ctx.Err(); err != nil { + return err + } + if c.cfg.DefaultRequestTimeout > 0 { + if _, ok := ctx.Deadline(); !ok { + var cancel context.CancelFunc + ctx, cancel = context.WithTimeout(ctx, c.cfg.DefaultRequestTimeout) + defer cancel() + } + } select { case <-c.done: return ErrShutdown @@ -647,11 +709,16 @@ func (c *Client) Request(ctx context.Context, method string, params any, out any if err != nil { return err } + rawParams, hook, err := c.flow.prepareRequest(method, rawParams, opts) + if err != nil { + return err + } id := c.nextID.Add(1) - 1 resultCh := make(chan pendingResult, 1) c.pendingMu.Lock() c.pending[id] = resultCh + c.pendingHooks[id] = hook c.pendingMu.Unlock() req := ahptypes.JsonRpcMessage{Request: &ahptypes.JsonRpcRequest{ @@ -660,24 +727,27 @@ func (c *Client) Request(ctx context.Context, method string, params any, out any Method: method, Params: rawParams, }} - if err := c.send(ctx, req); err != nil { + if err := c.sendOrdinary(ctx, req, hook); err != nil { + if hook != nil { + hook.cancel() + } c.pendingMu.Lock() delete(c.pending, id) + delete(c.pendingHooks, id) c.pendingMu.Unlock() return err } - // Apply the configured default timeout if no deadline is set. - if c.cfg.DefaultRequestTimeout > 0 { - if _, ok := ctx.Deadline(); !ok { - var cancel context.CancelFunc - ctx, cancel = context.WithTimeout(ctx, c.cfg.DefaultRequestTimeout) - defer cancel() - } + var abandoned <-chan struct{} + if hook != nil { + abandoned = hook.done } select { case r := <-resultCh: + if r.localErr != nil { + return r.localErr + } if r.err != nil { return &RPCError{Code: r.err.Code, Message: r.err.Message, Data: r.err.Data} } @@ -688,16 +758,30 @@ func (c *Client) Request(ctx context.Context, method string, params any, out any } return nil case <-ctx.Done(): + if hook != nil { + hook.cancel() + } c.pendingMu.Lock() delete(c.pending, id) + delete(c.pendingHooks, id) c.pendingMu.Unlock() return ctx.Err() + case <-abandoned: + c.pendingMu.Lock() + delete(c.pending, id) + delete(c.pendingHooks, id) + c.pendingMu.Unlock() + if ctx.Err() != nil { + return ctx.Err() + } + return hook.err case <-c.done: return ErrShutdown } } -// Notify sends a JSON-RPC notification (fire-and-forget). +// Notify waits until a JSON-RPC notification is written. Windowed data may +// wait for host credit, and a full buffered-send queue returns ErrBackpressure. func (c *Client) Notify(ctx context.Context, method string, params any) error { select { case <-c.done: @@ -717,11 +801,27 @@ func (c *Client) Notify(ctx context.Context, method string, params any) error { } func (c *Client) send(ctx context.Context, msg ahptypes.JsonRpcMessage) error { + if err := ctx.Err(); err != nil { + return err + } + if handled, err := c.flow.send(ctx, msg); handled { + return err + } + return c.sendOrdinary(ctx, msg, nil) +} + +func (c *Client) sendOrdinary(ctx context.Context, msg ahptypes.JsonRpcMessage, hooks *requestHooks) error { done := make(chan error, 1) + var abandoned <-chan struct{} + if hooks != nil { + abandoned = hooks.done + } select { - case c.outbound <- outboundMsg{msg: msg, done: done}: + case c.outbound <- outboundMsg{msg: msg, done: done, hooks: hooks}: case <-ctx.Done(): return ctx.Err() + case <-abandoned: + return hooks.err case <-c.done: return ErrShutdown } @@ -733,6 +833,8 @@ func (c *Client) send(ctx context.Context, msg ahptypes.JsonRpcMessage) error { return err case <-ctx.Done(): return ctx.Err() + case <-abandoned: + return hooks.err case <-c.done: return ErrShutdown } @@ -780,12 +882,15 @@ func (c *Client) Initialize(ctx context.Context, clientID string, protocolVersio // Reconnect re-establishes a dropped connection with the server's // `reconnect` flow. func (c *Client) Reconnect(ctx context.Context, clientID string, lastSeenServerSeq int64, subscriptions []string) (*ahptypes.ReconnectResult, error) { - params := ahptypes.ReconnectParams{ - Channel: ahptypes.RootResourceURI, - ClientId: clientID, - LastSeenServerSeq: lastSeenServerSeq, - Subscriptions: subscriptions, - } + return c.ReconnectWithParams(ctx, ahptypes.ReconnectParams{ + Channel: ahptypes.RootResourceURI, ClientId: clientID, + LastSeenServerSeq: lastSeenServerSeq, Subscriptions: subscriptions, + }) +} + +// ReconnectWithParams also accepts the generated Windows resume options. +func (c *Client) ReconnectWithParams(ctx context.Context, params ahptypes.ReconnectParams) (*ahptypes.ReconnectResult, error) { + params.Channel = ahptypes.RootResourceURI var out ahptypes.ReconnectResult if err := c.Request(ctx, "reconnect", params, &out); err != nil { return nil, err @@ -839,6 +944,20 @@ func (c *Client) AttachSubscription(uri string) *Subscription { // Unsubscribe sends an `unsubscribe` notification and drops every // local [Subscription] for uri. func (c *Client) Unsubscribe(ctx context.Context, uri string) error { + if err := ctx.Err(); err != nil { + return err + } + var cleanupDone chan error + c.flow.mu.Lock() + if f := c.flow.channels[uri]; f != nil { + c.flow.retireLocked(f, ErrClosed, false) + if !c.flow.stopped { + cleanupDone = make(chan error, 1) + c.flow.cleanup = append(c.flow.cleanup, flowCleanup{channel: uri, done: cleanupDone}) + c.flow.signal() + } + } + c.flow.mu.Unlock() c.subscriptionsMu.Lock() subs := c.subscriptions[uri] delete(c.subscriptions, uri) @@ -846,6 +965,16 @@ func (c *Client) Unsubscribe(ctx context.Context, uri string) error { for _, s := range subs { s.Close() } + if cleanupDone != nil { + select { + case err := <-cleanupDone: + return err + case <-ctx.Done(): + return ctx.Err() + case <-c.done: + return ErrShutdown + } + } return c.Notify(ctx, "unsubscribe", ahptypes.UnsubscribeParams{Channel: uri}) } diff --git a/clients/go/ahp/coalescing.go b/clients/go/ahp/coalescing.go new file mode 100644 index 000000000..eaa3659b7 --- /dev/null +++ b/clients/go/ahp/coalescing.go @@ -0,0 +1,115 @@ +package ahp + +import ( + "bytes" + "encoding/base64" + "encoding/json" + "strings" + + "github.com/microsoft/agent-host-protocol/clients/go/ahptypes" +) + +type CoalescingDirection string + +const ( + CoalesceReceive CoalescingDirection = "receive" + CoalesceSend CoalescingDirection = "send" +) + +type CoalescingContext struct { + Channel string + Direction CoalescingDirection +} + +// Coalescer runs before built-ins for Method. Merge must be nonblocking and +// cannot call back into the client. Return false to leave adjacent items intact. +type Coalescer struct { + Method string + Merge func(earlier, later ahptypes.JsonRpcNotification, context CoalescingContext) (ahptypes.JsonRpcNotification, bool) +} + +func (e *flowEngine) coalesce(a, b ahptypes.JsonRpcNotification, direction CoalescingDirection, f *channelFlow) (ahptypes.JsonRpcNotification, bool) { + if a.Method != b.Method || !ahptypes.CountsAgainstWindow(a.Method) || a.Method == "dispatchAction" || a.Method == "channel/snapshot" || a.Method == "tcp/eof" || a.Method == "channel/frame" { + return a, false + } + if a.Method == "action" { + var p, q ahptypes.ActionEnvelope + if json.Unmarshal(a.Params, &p) != nil || json.Unmarshal(b.Params, &q) != nil || p.Origin != nil || q.Origin != nil || p.RejectionReason != nil || q.RejectionReason != nil { + return a, false + } + } + for _, rule := range e.cfg.Coalescers { + if rule.Method == a.Method && rule.Merge != nil { + x, y := a, b + x.Params = append(json.RawMessage(nil), a.Params...) + y.Params = append(json.RawMessage(nil), b.Params...) + if n, ok := rule.Merge(x, y, CoalescingContext{Channel: f.channel, Direction: direction}); ok { + channel, err := notificationChannel(n) + if err == nil && channel == f.channel && n.Method == a.Method { + n.JsonRpc = ahptypes.JsonRpcV2 + return n, true + } + } + } + } + if direction == CoalesceSend && a.Method == "tcp/data" { + var p, q ahptypes.TcpDataParams + if json.Unmarshal(a.Params, &p) != nil || json.Unmarshal(b.Params, &q) != nil { + return a, false + } + x, err := decodeCoalescedTCP(p.Data) + if err != nil { + return a, false + } + y, err := decodeCoalescedTCP(q.Data) + if err != nil { + return a, false + } + + p.Data = base64.StdEncoding.EncodeToString(append(x, y...)) + a.Params, _ = json.Marshal(p) + return a, true + } + if direction != CoalesceReceive || a.Method != "action" { + return a, false + } + var p, q ahptypes.ActionEnvelope + if json.Unmarshal(a.Params, &p) != nil || json.Unmarshal(b.Params, &q) != nil { + return a, false + } + switch x := p.Action.Value.(type) { + case *ahptypes.TerminalDataAction: + y, ok := q.Action.Value.(*ahptypes.TerminalDataAction) + if !ok { + return a, false + } + x.Data += y.Data + case *ahptypes.ChatDeltaAction: + y, ok := q.Action.Value.(*ahptypes.ChatDeltaAction) + if !ok || x.TurnId != y.TurnId || x.PartId != y.PartId { + return a, false + } + left, _ := json.Marshal(x.Meta) + right, _ := json.Marshal(y.Meta) + if !bytes.Equal(left, right) { + return a, false + } + x.Content += y.Content + default: + return a, false + } + p.ServerSeq = q.ServerSeq + a.Params, _ = json.Marshal(p) + return a, true +} + +func decodeCoalescedTCP(text string) ([]byte, error) { + if text == "" || strings.ContainsAny(text, " \r\n\t") { + return nil, base64.CorruptInputError(0) + } + b, err := base64.StdEncoding.Strict().DecodeString(text) + if err == nil && base64.StdEncoding.EncodeToString(b) != text { + err = base64.CorruptInputError(0) + } + return b, err +} diff --git a/clients/go/ahp/delivery.go b/clients/go/ahp/delivery.go new file mode 100644 index 000000000..abc39097a --- /dev/null +++ b/clients/go/ahp/delivery.go @@ -0,0 +1,338 @@ +package ahp + +import ( + "context" + "encoding/json" + "io" + "sync" + "sync/atomic" + + "github.com/microsoft/agent-host-protocol/clients/go/ahptypes" +) + +// SubscriptionEventSnapshot is the initial or replacement channel snapshot. +type SubscriptionEventSnapshot struct { + Params ahptypes.ChannelSnapshotParams +} + +func (SubscriptionEventSnapshot) isSubscriptionEvent() {} + +type SubscriptionEventTCPData struct{ Params ahptypes.TcpDataParams } + +func (SubscriptionEventTCPData) isSubscriptionEvent() {} + +type SubscriptionEventTCPEOF struct{ Params ahptypes.TcpEofParams } + +func (SubscriptionEventTCPEOF) isSubscriptionEvent() {} + +type SubscriptionEventTCPReset struct{ Params ahptypes.TcpResetParams } + +func (SubscriptionEventTCPReset) isSubscriptionEvent() {} + +// SubscriptionEventNotification retains notifications without a dedicated event. +type SubscriptionEventNotification struct{ Notification ahptypes.JsonRpcNotification } + +func (SubscriptionEventNotification) isSubscriptionEvent() {} + +func eventForNotification(n ahptypes.JsonRpcNotification) (SubscriptionEvent, error) { + decode := func(out any) error { return json.Unmarshal(n.Params, out) } + switch n.Method { + case "action": + var p ahptypes.ActionEnvelope + err := decode(&p) + return SubscriptionEventAction{Envelope: p}, err + case "channel/snapshot": + var p ahptypes.ChannelSnapshotParams + err := decode(&p) + return SubscriptionEventSnapshot{Params: p}, err + case "tcp/data": + var p ahptypes.TcpDataParams + err := decode(&p) + return SubscriptionEventTCPData{Params: p}, err + case "tcp/eof": + var p ahptypes.TcpEofParams + err := decode(&p) + return SubscriptionEventTCPEOF{Params: p}, err + case "tcp/reset": + var p ahptypes.TcpResetParams + err := decode(&p) + return SubscriptionEventTCPReset{Params: p}, err + case "root/sessionAdded": + var p ahptypes.SessionAddedParams + err := decode(&p) + return SubscriptionEventSessionAdded{Params: p}, err + case "root/sessionRemoved": + var p ahptypes.SessionRemovedParams + err := decode(&p) + return SubscriptionEventSessionRemoved{Params: p}, err + case "root/sessionSummaryChanged": + var p ahptypes.SessionSummaryChangedParams + err := decode(&p) + return SubscriptionEventSessionSummaryChanged{Params: p}, err + case "auth/required": + var p ahptypes.AuthRequiredParams + err := decode(&p) + return SubscriptionEventAuthRequired{Params: p}, err + default: + return SubscriptionEventNotification{Notification: n}, nil + } +} + +// Delivery owns the credit for an event. Release is idempotent and can be +// called from another goroutine. WireBytes includes already-credited fragments. +type Delivery struct { + event SubscriptionEvent + notification ahptypes.JsonRpcNotification + bytes int64 + epoch uint64 + credit int64 + source *channelFlow + release func() + released atomic.Bool +} + +func newDelivery(n ahptypes.JsonRpcNotification, event SubscriptionEvent, bytes int64, release func()) *Delivery { + return &Delivery{event: event, notification: n, bytes: bytes, release: release} +} +func (d *Delivery) Event() SubscriptionEvent { return d.event } +func (d *Delivery) Notification() ahptypes.JsonRpcNotification { + n := d.notification + n.Params = append(json.RawMessage(nil), n.Params...) + return n +} +func (d *Delivery) WireBytes() int64 { return d.bytes } +func (d *Delivery) Release() { + if d != nil && d.released.CompareAndSwap(false, true) && d.release != nil { + d.release() + } +} + +type deliveryQueue struct { + mu sync.Mutex + nextGate chan struct{} + updateGate chan struct{} + engine *flowEngine + flow *channelFlow + items []*Delivery + previous *Delivery + mode ReceiveMode + changed chan struct{} + closed bool + final bool + err error + status FlowControlStatus + statusStreams []chan FlowControlStatus + unreleased int64 + releaseVersion uint64 +} + +func newDeliveryQueue() *deliveryQueue { + return &deliveryQueue{changed: make(chan struct{}), nextGate: make(chan struct{}, 1), updateGate: make(chan struct{}, 1), status: FlowControlStatus{Mode: "ordinary", Writable: true}} +} +func (q *deliveryQueue) signalLocked() { close(q.changed); q.changed = make(chan struct{}) } +func (q *deliveryQueue) bind(e *flowEngine, f *channelFlow) { + q.mu.Lock() + defer q.mu.Unlock() + q.engine, q.flow, q.closed, q.err = e, f, false, nil + q.signalLocked() +} +func (q *deliveryQueue) binding() (*flowEngine, *channelFlow) { + q.mu.Lock() + defer q.mu.Unlock() + return q.engine, q.flow +} +func (q *deliveryQueue) setMode(opts SubscribeOptions) { + q.mu.Lock() + defer q.mu.Unlock() + q.mode = ReceiveAuto + if opts.Window != nil { + q.mode = opts.Window.Receive + } +} +func (q *deliveryQueue) publish(status FlowControlStatus) { + q.mu.Lock() + defer q.mu.Unlock() + if q.final { + status.Writable = false + } + q.status = status + for _, ch := range q.statusStreams { + select { + case ch <- status: + default: + select { + case <-ch: + default: + } + select { + case ch <- status: + default: + } + } + } + q.signalLocked() +} +func (q *deliveryQueue) close(err error) { + q.mu.Lock() + defer q.mu.Unlock() + if q.final { + return + } + q.closed, q.err = true, err + q.signalLocked() +} + +func (q *deliveryQueue) finishLocked(err error) { + if q.final { + return + } + q.closed, q.final, q.err = true, true, err + q.items, q.previous = nil, nil + q.status.Writable = false + q.status.BufferedReceiveBytes = 0 + q.unreleased = 0 + q.releaseVersion++ + for _, ch := range q.statusStreams { + close(ch) + } + + q.statusStreams = nil + q.signalLocked() +} + +func (q *deliveryQueue) isFinal() bool { + q.mu.Lock() + defer q.mu.Unlock() + return q.final +} + +// Next retrieves the next credit-owning delivery. ReceiveAuto releases the +// previous delivery before waiting, including when ctx is canceled. +// Do not use Events as a substitute: it is an observe-only, lossy channel. +func (s *Subscription) Next(ctx context.Context) (*Delivery, error) { + if s.owner == nil { + return nil, io.EOF + } + q := s.owner + select { + case q.nextGate <- struct{}{}: + defer func() { <-q.nextGate }() + case <-ctx.Done(): + q.advance() + return nil, ctx.Err() + } + q.advance() + for { + if err := ctx.Err(); err != nil { + return nil, err + } + q.mu.Lock() + if len(q.items) > 0 { + d := q.items[0] + q.items[0] = nil + q.items = q.items[1:] + if q.mode == ReceiveAuto { + q.previous = d + } + q.mu.Unlock() + return d, nil + } + if q.closed { + err := q.err + q.mu.Unlock() + if err == nil { + err = io.EOF + } + return nil, err + } + changed := q.changed + q.mu.Unlock() + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-changed: + } + } +} + +func (q *deliveryQueue) advance() { + q.mu.Lock() + prev := q.previous + q.previous = nil + q.mu.Unlock() + if prev != nil { + prev.Release() + } +} + +// FlowControl returns current backpressure; its limits are an independent copy. +func (s *Subscription) FlowControl() FlowControlStatus { + if s.owner == nil { + return FlowControlStatus{Mode: "ordinary", Writable: true} + } + q := s.owner + q.mu.Lock() + defer q.mu.Unlock() + status := q.status + if status.Limits != nil { + l := *status.Limits + status.Limits = &l + } + return status +} + +// FlowControlChanges returns a conflated status stream, initially current status. +func (s *Subscription) FlowControlChanges() <-chan FlowControlStatus { + ch := make(chan FlowControlStatus, 1) + if s.owner == nil { + ch <- s.FlowControl() + return ch + } + q := s.owner + q.mu.Lock() + defer q.mu.Unlock() + if q.final { + ch <- q.status + close(ch) + return ch + } + ch <- q.status + q.statusStreams = append(q.statusStreams, ch) + return ch +} +func (s *Subscription) WaitWritable(ctx context.Context) error { + if err := ctx.Err(); err != nil { + return err + } + if s.owner == nil { + return nil + } + q := s.owner + for { + q.mu.Lock() + if q.closed { + err := q.err + q.mu.Unlock() + return err + } + if q.status.Writable { + q.mu.Unlock() + return nil + } + changed := q.changed + q.mu.Unlock() + select { + case <-ctx.Done(): + return ctx.Err() + case <-changed: + } + } +} +func (s *Subscription) Err() error { + if s.owner == nil { + return nil + } + s.owner.mu.Lock() + defer s.owner.mu.Unlock() + return s.owner.err +} diff --git a/clients/go/ahp/hosts/hosts.go b/clients/go/ahp/hosts/hosts.go index 65396c4c9..b5f7d0f5d 100644 --- a/clients/go/ahp/hosts/hosts.go +++ b/clients/go/ahp/hosts/hosts.go @@ -417,18 +417,21 @@ type hostState struct { generation uint64 cancel context.CancelFunc supervised sync.WaitGroup + managed map[string]*ManagedSubscription } // MultiHostClient is the public multi-host registry + reconnect // supervisor. type MultiHostClient struct { - mu sync.RWMutex - hosts map[HostID]*hostState - store ClientIDStore - eventsMu sync.Mutex - events []chan HostEvent - subMu sync.Mutex - subs []chan HostSubscriptionEvent + mu sync.RWMutex + hosts map[HostID]*hostState + store ClientIDStore + eventsMu sync.Mutex + events []chan HostEvent + subMu sync.Mutex + subs []chan HostSubscriptionEvent + flowMu sync.Mutex + flowEvents []chan HostFlowControlEvent rootCtx context.Context rootCancel context.CancelFunc @@ -529,6 +532,7 @@ func (m *MultiHostClient) AddHost(ctx context.Context, cfg HostConfig) (*HostHan state: HostState{Kind: HostStateDisconnected}, updatedAt: time.Now(), cancel: cancel, + managed: make(map[string]*ManagedSubscription), } m.mu.Lock() @@ -563,6 +567,7 @@ func (m *MultiHostClient) openHost(ctx context.Context, hs *hostState) error { client, err := ahp.Connect(ctx, boxed, hs.cfg.ClientConfig) if err != nil { + _ = boxed.Close(ctx) return fmt.Errorf("hosts: connect: %w", err) } @@ -572,6 +577,21 @@ func (m *MultiHostClient) openHost(ctx context.Context, hs *hostState) error { return fmt.Errorf("hosts: initialize: %w", err) } + hs.mu.RLock() + resuming := hs.client != nil + managed := make([]*ManagedSubscription, 0, len(hs.managed)) + for _, sub := range hs.managed { + managed = append(managed, sub) + } + hs.mu.RUnlock() + go m.pumpFlowEvents(hs, client, resuming) + for _, sub := range managed { + if err := sub.resume(ctx, client); err != nil { + _ = client.Shutdown(ctx) + return fmt.Errorf("hosts: re-subscribe: %w", err) + } + } + hs.mu.Lock() hs.client = client hs.protoVer = result.ProtocolVersion @@ -593,15 +613,14 @@ func (m *MultiHostClient) pumpEvents(hs *hostState, client *ahp.Client) { defer stream.Close() for ev := range stream.Events() { m.subMu.Lock() - subs := append([]chan HostSubscriptionEvent(nil), m.subs...) - m.subMu.Unlock() out := HostSubscriptionEvent{HostID: hs.id, Channel: ev.Channel, Event: ev.Event} - for _, ch := range subs { + for _, ch := range m.subs { select { case ch <- out: default: } } + m.subMu.Unlock() } } @@ -632,7 +651,9 @@ func (m *MultiHostClient) supervise(ctx context.Context, hs *hostState) { return default: } + m.dropManagedTCP(hs) if policy.IsDisabled() { + m.closeManaged(hs, ErrHostNotConnected) m.setHostState(hs, HostState{Kind: HostStateFailed, Err: errors.New("hosts: transport closed and reconnect disabled")}) return } @@ -661,6 +682,7 @@ func (m *MultiHostClient) supervise(ctx context.Context, hs *hostState) { } attempt++ if policy.MaxAttempts > 0 && attempt > policy.MaxAttempts { + m.closeManaged(hs, ErrHostNotConnected) m.setHostState(hs, HostState{Kind: HostStateFailed, Err: fmt.Errorf("hosts: exceeded %d reconnect attempts", policy.MaxAttempts)}) return } @@ -677,14 +699,13 @@ func (m *MultiHostClient) setHostState(hs *hostState, state HostState) { hs.mu.Unlock() m.eventsMu.Lock() - subs := append([]chan HostEvent(nil), m.events...) - m.eventsMu.Unlock() - for _, ch := range subs { + for _, ch := range m.events { select { case ch <- HostEvent{HostID: hs.id, State: state}: default: } } + m.eventsMu.Unlock() } // Host returns a fresh snapshot of the host with the given id, or @@ -782,6 +803,7 @@ func (m *MultiHostClient) RemoveHost(ctx context.Context, id HostID) error { return ErrUnknownHost } hs.cancel() + m.closeManaged(hs, ErrHostNotConnected) hs.mu.RLock() client := hs.client hs.mu.RUnlock() @@ -805,7 +827,12 @@ func (m *MultiHostClient) removeHostLocked(id HostID) error { func (m *MultiHostClient) Events() <-chan HostEvent { ch := make(chan HostEvent, 64) m.eventsMu.Lock() - m.events = append(m.events, ch) + select { + case <-m.rootCtx.Done(): + close(ch) + default: + m.events = append(m.events, ch) + } m.eventsMu.Unlock() return ch } @@ -815,7 +842,12 @@ func (m *MultiHostClient) Events() <-chan HostEvent { func (m *MultiHostClient) Subscriptions() <-chan HostSubscriptionEvent { ch := make(chan HostSubscriptionEvent, 256) m.subMu.Lock() - m.subs = append(m.subs, ch) + select { + case <-m.rootCtx.Done(): + close(ch) + default: + m.subs = append(m.subs, ch) + } m.subMu.Unlock() return ch } @@ -833,6 +865,7 @@ func (m *MultiHostClient) Shutdown(ctx context.Context) error { m.mu.Unlock() for _, hs := range hosts { hs.cancel() + m.closeManaged(hs, ErrHostNotConnected) hs.mu.RLock() client := hs.client hs.mu.RUnlock() @@ -841,5 +874,23 @@ func (m *MultiHostClient) Shutdown(ctx context.Context) error { } hs.supervised.Wait() } + m.eventsMu.Lock() + for _, ch := range m.events { + close(ch) + } + m.events = nil + m.eventsMu.Unlock() + m.subMu.Lock() + for _, ch := range m.subs { + close(ch) + } + m.subs = nil + m.subMu.Unlock() + m.flowMu.Lock() + for _, ch := range m.flowEvents { + close(ch) + } + m.flowEvents = nil + m.flowMu.Unlock() return nil } diff --git a/clients/go/ahp/hosts/managed_lifecycle_test.go b/clients/go/ahp/hosts/managed_lifecycle_test.go new file mode 100644 index 000000000..892d84c27 --- /dev/null +++ b/clients/go/ahp/hosts/managed_lifecycle_test.go @@ -0,0 +1,126 @@ +package hosts + +import ( + "context" + "encoding/json" + "errors" + "testing" + "time" + + "github.com/microsoft/agent-host-protocol/clients/go/ahp" + "github.com/microsoft/agent-host-protocol/clients/go/ahptypes" +) + +func TestManagedResubscribeTimeoutDropsOnlyThatChannel(t *testing.T) { + m, f, ctx := managedSetupOptions(t, false, "ahp-terminal:/bad", 20*time.Millisecond) + _, bad, err := m.SubscribeWithOptions(ctx, "host", "ahp-terminal:/bad", ahp.SubscribeOptions{Window: &ahp.WindowOptions{}}) + if err != nil { + t.Fatal(err) + } + _, good, err := m.SubscribeWithOptions(ctx, "host", "ahp-terminal:/good", ahp.SubscribeOptions{Window: &ahp.WindowOptions{}}) + if err != nil { + t.Fatal(err) + } + old, _ := m.ClientHandle("host") + oldClient, _ := old.Client() + first := <-f.servers + _ = first.Close(ctx) + waitManaged(t, func() bool { + h, _ := m.ClientHandle("host") + c, err := h.Client() + return err == nil && c != oldClient && m.Host("host").State.Kind == HostStateConnected && bad.Err() != nil + }) + if !errors.Is(bad.Err(), ahp.ErrSubscriptionDropped) { + t.Fatal(bad.Err()) + } + if good.Err() != nil { + t.Fatal("healthy channel dropped", good.Err()) + } + m.mu.RLock() + hs := m.hosts["host"] + m.mu.RUnlock() + hs.mu.RLock() + _, keptBad := hs.managed[bad.URI()] + _, keptGood := hs.managed[good.URI()] + hs.mu.RUnlock() + if keptBad || !keptGood { + t.Fatal("managed registry retained dropped channel") + } + second := <-f.servers + h, _ := m.ClientHandle("host") + secondClient, _ := h.Client() + _ = second.Close(ctx) + waitManaged(t, func() bool { + h, _ := m.ClientHandle("host") + c, err := h.Client() + return err == nil && c != secondClient && m.Host("host").State.Kind == HostStateConnected + }) + for { + select { + case record := <-f.records: + if record.generation >= 3 && record.message.Request != nil && record.message.Request.Method == "subscribe" { + var p ahptypes.SubscribeParams + _ = json.Unmarshal(record.message.Request.Params, &p) + if p.Channel == bad.URI() { + t.Fatal("later reconnect retried a dropped channel") + } + } + default: + return + } + } +} + +func TestManagedCloseDoesNotWaitForResubscribe(t *testing.T) { + for _, shutdown := range []bool{false, true} { + t.Run(map[bool]string{false: "close", true: "shutdown"}[shutdown], func(t *testing.T) { + m, f, ctx := managedSetupOptions(t, false, "ahp-terminal:/blocked", 0) + _, s, err := m.SubscribeWithOptions(ctx, "host", "ahp-terminal:/blocked", ahp.SubscribeOptions{Window: &ahp.WindowOptions{}}) + if err != nil { + t.Fatal(err) + } + first := <-f.servers + _ = first.Close(ctx) + waiting := false + deadline := time.After(time.Second) + for !waiting { + select { + case record := <-f.records: + waiting = record.generation == 2 && record.message.Request != nil && record.message.Request.Method == "subscribe" + case <-deadline: + t.Fatal("resubscribe never started") + } + } + op, cancel := context.WithTimeout(ctx, 500*time.Millisecond) + defer cancel() + done := make(chan error, 1) + go func() { + if shutdown { + done <- m.Shutdown(op) + } else { + done <- s.Close(op) + } + }() + select { + case err := <-done: + if err != nil { + t.Fatal(err) + } + case <-op.Done(): + t.Fatal("shutdown/close blocked behind resubscribe") + } + if !shutdown { + waitManaged(t, func() bool { return m.Host("host").State.Kind == HostStateConnected }) + m.mu.RLock() + hs := m.hosts["host"] + m.mu.RUnlock() + hs.mu.RLock() + n := len(hs.managed) + hs.mu.RUnlock() + if n != 0 { + t.Fatal("closed managed handle remained registered") + } + } + }) + } +} diff --git a/clients/go/ahp/hosts/managed_streams_test.go b/clients/go/ahp/hosts/managed_streams_test.go new file mode 100644 index 000000000..70fba7ebd --- /dev/null +++ b/clients/go/ahp/hosts/managed_streams_test.go @@ -0,0 +1,37 @@ +package hosts + +import ( + "context" + "testing" + "time" +) + +func awaitHostStreamClose[T any](t *testing.T, ctx context.Context, stream <-chan T) { + t.Helper() + for { + select { + case _, ok := <-stream: + if !ok { + return + } + case <-ctx.Done(): + t.Fatal("host observer stream remained open") + } + } +} + +func TestManagedShutdownClosesObserverStreams(t *testing.T) { + m, _, ctx := managedSetup(t, false) + events, subscriptions, flow := m.Events(), m.Subscriptions(), m.FlowControlEvents() + if err := m.Shutdown(ctx); err != nil { + t.Fatal(err) + } + timeout, cancel := context.WithTimeout(ctx, 500*time.Millisecond) + defer cancel() + awaitHostStreamClose(t, timeout, events) + awaitHostStreamClose(t, timeout, subscriptions) + awaitHostStreamClose(t, timeout, flow) + awaitHostStreamClose(t, timeout, m.Events()) + awaitHostStreamClose(t, timeout, m.Subscriptions()) + awaitHostStreamClose(t, timeout, m.FlowControlEvents()) +} diff --git a/clients/go/ahp/hosts/managed_subscription.go b/clients/go/ahp/hosts/managed_subscription.go new file mode 100644 index 000000000..0d2e6400d --- /dev/null +++ b/clients/go/ahp/hosts/managed_subscription.go @@ -0,0 +1,348 @@ +package hosts + +import ( + "context" + "errors" + "fmt" + "strings" + "sync" + + "github.com/microsoft/agent-host-protocol/clients/go/ahp" + "github.com/microsoft/agent-host-protocol/clients/go/ahptypes" +) + +// ManagedSubscription keeps its owner queue across initialize/re-subscribe +// cycles. Old deliveries precede the new snapshot; old releases send no credit. +type ManagedSubscription struct { + mu sync.Mutex + operation sync.Mutex + host *hostState + channel string + sub *ahp.Subscription + opts ahp.SubscribeOptions + changed chan struct{} + closed bool + err error + lifetime context.Context + cancel context.CancelFunc +} + +// SubscribeWithOptions subscribes on one host, retaining a windowed owner +// across connection replacement. A second call updates the same handle. +func (m *MultiHostClient) SubscribeWithOptions(ctx context.Context, id HostID, uri string, opts ahp.SubscribeOptions) (*ahptypes.SubscribeResult, *ManagedSubscription, error) { + m.mu.RLock() + hs := m.hosts[id] + m.mu.RUnlock() + if hs == nil { + return nil, nil, ErrUnknownHost + } + hs.mu.Lock() + s := hs.managed[uri] + created := s == nil + if s == nil { + life, cancel := context.WithCancel(context.Background()) + s = &ManagedSubscription{host: hs, channel: uri, changed: make(chan struct{}), lifetime: life, cancel: cancel} + hs.managed[uri] = s + } + hs.mu.Unlock() + s.operation.Lock() + defer s.operation.Unlock() + hs.mu.RLock() + client := hs.client + hs.mu.RUnlock() + if client == nil { + if created { + s.drop(ErrHostNotConnected) + } + return nil, s, ErrHostNotConnected + } + select { + case <-client.Done(): + if created { + s.drop(ErrHostNotConnected) + } + return nil, s, ErrHostNotConnected + default: + } + opCtx, finish := s.operationContext(ctx) + defer finish() + result, sub, err := client.SubscribeWithOptions(opCtx, uri, opts) + if err != nil { + if created { + s.drop(err) + } + return nil, s, err + } + s.mu.Lock() + if s.closed { + err := s.err + s.mu.Unlock() + sub.Close() + return nil, s, err + } + s.sub, s.opts, s.closed, s.err = sub, cloneOptions(opts), false, nil + s.signalLocked() + s.mu.Unlock() + return result, s, nil +} + +func cloneOptions(o ahp.SubscribeOptions) ahp.SubscribeOptions { + if o.Window != nil { + w := *o.Window + if w.Coalescing != nil { + v := *w.Coalescing + w.Coalescing = &v + } + o.Window = &w + } + if o.Delivery != nil { + v := *o.Delivery + if v.MaxLatencyMs != nil { + n := *v.MaxLatencyMs + v.MaxLatencyMs = &n + } + o.Delivery = &v + } + if o.View != nil { + v := *o.View + if v.Turns != nil { + n := *v.Turns + v.Turns = &n + } + o.View = &v + } + return o +} +func (s *ManagedSubscription) signalLocked() { close(s.changed); s.changed = make(chan struct{}) } +func (s *ManagedSubscription) URI() string { return s.channel } +func (s *ManagedSubscription) HostID() HostID { return s.host.id } + +func (s *ManagedSubscription) resume(ctx context.Context, client *ahp.Client) error { + s.operation.Lock() + defer s.operation.Unlock() + s.mu.Lock() + if s.closed || s.sub == nil { + s.mu.Unlock() + return nil + } + sub, opts := s.sub, cloneOptions(s.opts) + s.mu.Unlock() + if opts.Window == nil || strings.HasPrefix(s.channel, "ahp-tcp:") { + s.drop(ahp.ErrSubscriptionDropped) + return nil + } + opCtx, finish := s.operationContext(ctx) + defer finish() + _, err := client.ResumeSubscription(opCtx, sub, opts) + if err != nil { + select { + case <-client.Done(): + return err + default: + } + if ctx.Err() != nil { + return ctx.Err() + } + s.mu.Lock() + closed := s.closed + s.mu.Unlock() + if !closed { + s.drop(fmt.Errorf("%w: %v", ahp.ErrSubscriptionDropped, err)) + } + return nil + } + s.mu.Lock() + s.signalLocked() + s.mu.Unlock() + return nil +} + +// Next waits through managed reconnects. ReceiveAuto releases on the next call. +func (s *ManagedSubscription) Next(ctx context.Context) (*ahp.Delivery, error) { + for { + s.mu.Lock() + if s.closed { + err := s.err + s.mu.Unlock() + return nil, err + } + sub, changed := s.sub, s.changed + s.mu.Unlock() + if sub != nil { + d, err := sub.Next(ctx) + if err == nil { + return d, nil + } + if ctx.Err() != nil { + return nil, ctx.Err() + } + if sub.Err() == nil && !errors.Is(err, ahp.ErrClosed) && !errors.Is(err, ahp.ErrShutdown) { + return nil, err + } + } + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-changed: + } + } +} +func (s *ManagedSubscription) FlowControl() ahp.FlowControlStatus { + s.mu.Lock() + sub := s.sub + s.mu.Unlock() + if sub == nil { + return ahp.FlowControlStatus{Mode: "ordinary"} + } + return sub.FlowControl() +} +func (s *ManagedSubscription) FlowControlChanges() <-chan ahp.FlowControlStatus { + s.mu.Lock() + sub := s.sub + s.mu.Unlock() + if sub == nil { + ch := make(chan ahp.FlowControlStatus) + close(ch) + return ch + } + return sub.FlowControlChanges() +} +func (s *ManagedSubscription) WaitWritable(ctx context.Context) error { + for { + s.mu.Lock() + if s.closed { + err := s.err + s.mu.Unlock() + return err + } + sub, changed := s.sub, s.changed + s.mu.Unlock() + if sub != nil { + err := sub.WaitWritable(ctx) + if err == nil { + return nil + } + if ctx.Err() != nil { + return ctx.Err() + } + if sub.Err() == nil && !errors.Is(err, ahp.ErrClosed) && !errors.Is(err, ahp.ErrShutdown) { + return err + } + } + select { + case <-ctx.Done(): + return ctx.Err() + case <-changed: + } + } +} +func (s *ManagedSubscription) Update(ctx context.Context, opts ahp.SubscribeOptions) (*ahptypes.SubscribeResult, error) { + s.operation.Lock() + defer s.operation.Unlock() + s.mu.Lock() + sub := s.sub + s.mu.Unlock() + if sub == nil { + return nil, ErrHostNotConnected + } + opCtx, finish := s.operationContext(ctx) + defer finish() + result, err := sub.Update(opCtx, opts) + if err == nil { + s.mu.Lock() + s.opts = cloneOptions(opts) + s.mu.Unlock() + } + return result, err +} +func (s *ManagedSubscription) Err() error { s.mu.Lock(); defer s.mu.Unlock(); return s.err } +func (s *ManagedSubscription) drop(err error) { + s.mu.Lock() + if s.closed { + s.mu.Unlock() + return + } + s.closed, s.err = true, err + sub := s.sub + s.cancel() + s.signalLocked() + s.mu.Unlock() + s.host.mu.Lock() + if s.host.managed[s.channel] == s { + delete(s.host.managed, s.channel) + } + s.host.mu.Unlock() + if sub != nil { + sub.Close() + } +} + +// Close unsubscribes and permanently removes this managed owner. +func (s *ManagedSubscription) Close(ctx context.Context) error { + s.drop(ahp.ErrClosed) + return nil +} + +func (s *ManagedSubscription) operationContext(ctx context.Context) (context.Context, context.CancelFunc) { + op, cancel := context.WithCancel(ctx) + stop := context.AfterFunc(s.lifetime, cancel) + return op, func() { stop(); cancel() } +} +func (m *MultiHostClient) closeManaged(hs *hostState, err error) { + hs.mu.RLock() + subs := make([]*ManagedSubscription, 0, len(hs.managed)) + for _, s := range hs.managed { + subs = append(subs, s) + } + hs.mu.RUnlock() + for _, s := range subs { + s.drop(err) + } +} +func (m *MultiHostClient) dropManagedTCP(hs *hostState) { + hs.mu.RLock() + subs := make([]*ManagedSubscription, 0) + for uri, s := range hs.managed { + if strings.HasPrefix(uri, "ahp-tcp:") { + subs = append(subs, s) + } + } + hs.mu.RUnlock() + for _, s := range subs { + s.drop(ahp.ErrSubscriptionDropped) + } +} + +type HostFlowControlEvent struct { + HostID HostID + Event ahp.FlowControlEvent +} + +// FlowControlEvents observes telemetry from every host; it never releases credit. +func (m *MultiHostClient) FlowControlEvents() <-chan HostFlowControlEvent { + ch := make(chan HostFlowControlEvent, 64) + m.flowMu.Lock() + select { + case <-m.rootCtx.Done(): + close(ch) + default: + m.flowEvents = append(m.flowEvents, ch) + } + m.flowMu.Unlock() + return ch +} +func (m *MultiHostClient) pumpFlowEvents(hs *hostState, client *ahp.Client, resuming bool) { + for event := range client.FlowControlEvents() { + if resuming && event.Kind == "negotiated" { + event.Kind = "resumed" + } + m.flowMu.Lock() + for _, ch := range m.flowEvents { + select { + case ch <- HostFlowControlEvent{HostID: hs.id, Event: event}: + default: + } + } + m.flowMu.Unlock() + } +} diff --git a/clients/go/ahp/hosts/managed_subscription_test.go b/clients/go/ahp/hosts/managed_subscription_test.go new file mode 100644 index 000000000..e1906ea72 --- /dev/null +++ b/clients/go/ahp/hosts/managed_subscription_test.go @@ -0,0 +1,267 @@ +package hosts + +import ( + "context" + "encoding/json" + "errors" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/microsoft/agent-host-protocol/clients/go/ahp" + "github.com/microsoft/agent-host-protocol/clients/go/ahptypes" +) + +func (t *fakeTransport) TransportCapabilities() ahp.TransportCapabilities { + return ahp.TransportCapabilities{SupportsWindowedDelivery: true, MaximumFrameBytes: 16384} +} + +type managedWireRecord struct { + generation int + message ahptypes.JsonRpcMessage +} +type managedHostFactory struct { + servers chan *fakeTransport + records chan managedWireRecord + generation atomic.Int32 + rejectSecond bool + ignoreSecond string +} + +func (f *managedHostFactory) open(ctx context.Context, _ HostID) (ahp.Transport, error) { + generation := int(f.generation.Add(1)) + client, server := newFakePair() + f.servers <- server + go func() { + for { + frame, err := server.Recv(ctx) + if err != nil { + return + } + m, err := frame.IntoParsed() + if err != nil { + return + } + f.records <- managedWireRecord{generation, m} + if m.Request == nil { + continue + } + var result any + switch m.Request.Method { + case "initialize": + result = map[string]any{"protocolVersion": "1.0.0"} + case "subscribe": + var p ahptypes.SubscribeParams + _ = json.Unmarshal(m.Request.Params, &p) + if generation > 1 && p.Channel == f.ignoreSecond { + continue + } + if p.FlowControl == nil { + b, _ := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": m.Request.ID, "error": map[string]any{"code": -32602, "message": "missing window offer"}}) + _ = server.Send(ctx, ahp.NewTextMessage(string(b))) + continue + } + if f.rejectSecond && generation > 1 { + b, _ := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": m.Request.ID, "error": map[string]any{"code": ahptypes.ErrorCodePermissionDenied, "message": "permission changed"}}) + _ = server.Send(ctx, ahp.NewTextMessage(string(b))) + continue + } + result = map[string]any{"flowControl": map[string]any{"clientReceive": map[string]any{"windowBytes": 1024, "maximumFrameBytes": 1024}, "hostReceive": map[string]any{"windowBytes": 1, "maximumFrameBytes": 1024}}} + b, _ := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": m.Request.ID, "result": result}) + _ = server.Send(ctx, ahp.NewTextMessage(string(b))) + b, _ = json.Marshal(map[string]any{"jsonrpc": "2.0", "method": "channel/snapshot", "params": map[string]any{"channel": p.Channel, "snapshot": map[string]any{"resource": p.Channel, "state": map[string]any{}, "fromSeq": generation * 100}}}) + _ = server.Send(ctx, ahp.NewTextMessage(string(b))) + if generation == 1 && !strings.HasPrefix(p.Channel, "ahp-tcp:") { + b, _ = json.Marshal(map[string]any{"jsonrpc": "2.0", "method": "action", "params": map[string]any{"channel": p.Channel, "action": map[string]any{"type": "terminal/data", "data": "old"}, "serverSeq": 101}}) + _ = server.Send(ctx, ahp.NewTextMessage(string(b))) + } + continue + default: + result = nil + } + b, _ := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": m.Request.ID, "result": result}) + _ = server.Send(ctx, ahp.NewTextMessage(string(b))) + } + }() + return client, nil +} +func managedSetup(t *testing.T, reject bool) (*MultiHostClient, *managedHostFactory, context.Context) { + return managedSetupOptions(t, reject, "", 30*time.Second) +} + +func managedSetupOptions(t *testing.T, reject bool, ignore string, timeout time.Duration) (*MultiHostClient, *managedHostFactory, context.Context) { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + t.Cleanup(cancel) + m := NewMultiHostClient() + t.Cleanup(func() { _ = m.Shutdown(context.Background()) }) + f := &managedHostFactory{servers: make(chan *fakeTransport, 8), records: make(chan managedWireRecord, 128), rejectSecond: reject, ignoreSecond: ignore} + cfg := NewHostConfig("host", "Host", f.open) + cfg.ReconnectPolicy = ReconnectPolicy{InitialBackoff: time.Millisecond, MaxBackoff: time.Millisecond, MaxAttempts: 3} + cfg.ClientConfig.Windowing = ahp.WindowingConfig{WindowBytes: 2048, MaximumFrameBytes: 4096} + cfg.ClientConfig.DefaultRequestTimeout = timeout + if _, err := m.AddHost(ctx, cfg); err != nil { + t.Fatal(err) + } + return m, f, ctx +} +func waitManaged(t *testing.T, p func() bool) { + t.Helper() + deadline := time.Now().Add(time.Second) + for !p() { + if time.Now().After(deadline) { + t.Fatal("managed condition timed out") + } + time.Sleep(time.Millisecond) + } +} + +func TestManagedWindowQueueAndEpochsAcrossInitialize(t *testing.T) { + m, f, ctx := managedSetup(t, false) + events := m.FlowControlEvents() + _, sub, err := m.SubscribeWithOptions(ctx, "host", "ahp-terminal:/t", ahp.SubscribeOptions{Window: &ahp.WindowOptions{Receive: ahp.ReceiveManual}}) + if err != nil { + t.Fatal(err) + } + first, err := sub.Next(ctx) + if err != nil { + t.Fatal(err) + } + oldHandle, _ := m.ClientHandle("host") + oldClient, _ := oldHandle.Client() + waitManaged(t, func() bool { return sub.FlowControl().BufferedReceiveBytes > first.WireBytes() }) + server := <-f.servers + _ = server.Close(ctx) + waitManaged(t, func() bool { + h, _ := m.ClientHandle("host") + c, err := h.Client() + return err == nil && c != oldClient && m.Host("host").State.Kind == HostStateConnected + }) + if _, err := oldHandle.Client(); !errors.Is(err, ErrHostReconnected) { + t.Fatal("old client handle stayed valid") + } + old, err := sub.Next(ctx) + if err != nil { + t.Fatal(err) + } + if a, ok := old.Event().(ahp.SubscriptionEventAction); !ok || a.Envelope.ServerSeq != 101 { + t.Fatal("retained item was reordered") + } + fresh, err := sub.Next(ctx) + if err != nil { + t.Fatal(err) + } + if s, ok := fresh.Event().(ahp.SubscriptionEventSnapshot); !ok || s.Params.Snapshot.FromSeq != 200 { + t.Fatal("new snapshot not after retained queue") + } + buffered := sub.FlowControl().BufferedReceiveBytes + first.Release() + old.Release() + if sub.FlowControl().BufferedReceiveBytes != buffered-first.WireBytes()-old.WireBytes() { + t.Fatal("old retained receive bytes were not released") + } + fresh.Release() + waitManaged(t, func() bool { return sub.FlowControl().BufferedReceiveBytes == 0 }) + found := false + timer := time.NewTimer(time.Second) + defer timer.Stop() + for !found { + select { + case r := <-f.records: + if r.message.Request != nil && r.message.Request.Method == "reconnect" { + t.Fatal("Go hosts changed to protocol reconnect") + } + if r.message.Notification != nil && r.message.Notification.Method == "channel/credit" && r.generation == 2 { + var p ahptypes.ChannelCreditParams + _ = json.Unmarshal(r.message.Notification.Params, &p) + if p.ConsumedBytes != fresh.WireBytes() { + t.Fatal("wrong new-epoch cumulative credit") + } + found = true + } + case <-timer.C: + t.Fatal("no new credit") + } + } + resumed := false + waitManaged(t, func() bool { + for { + select { + case e := <-events: + if e.HostID == "host" && e.Event.Kind == "resumed" { + resumed = true + } + default: + return resumed + } + } + }) +} + +func TestManagedTCPDroppedAndPermissionChange(t *testing.T) { + for _, tcp := range []bool{true, false} { + t.Run(map[bool]string{true: "tcp", false: "permission"}[tcp], func(t *testing.T) { + m, f, ctx := managedSetup(t, !tcp) + uri := "ahp-terminal:/t" + if tcp { + uri = "ahp-tcp:/t" + } + _, sub, err := m.SubscribeWithOptions(ctx, "host", uri, ahp.SubscribeOptions{Window: &ahp.WindowOptions{}}) + if err != nil { + t.Fatal(err) + } + if _, err := sub.Next(ctx); err != nil { + t.Fatal(err) + } + server := <-f.servers + _ = server.Close(ctx) + waitManaged(t, func() bool { return errors.Is(sub.Err(), ahp.ErrSubscriptionDropped) }) + if _, err := sub.Next(ctx); !errors.Is(err, ahp.ErrSubscriptionDropped) { + t.Fatal("dropped stream did not fail") + } + waitManaged(t, func() bool { return f.generation.Load() > 1 && m.Host("host").State.Kind == HostStateConnected }) + for { + select { + case r := <-f.records: + if tcp && r.generation > 1 && r.message.Request != nil && r.message.Request.Method == "subscribe" { + t.Fatal("TCP automatically re-subscribed") + } + default: + return + } + } + }) + } +} + +func TestManagedSameHandleUpdateAndHostIsolation(t *testing.T) { + m, f, ctx := managedSetup(t, false) + cfg := NewHostConfig("other", "Other", f.open) + if _, err := m.AddHost(ctx, cfg); err != nil { + t.Fatal(err) + } + _, a, err := m.SubscribeWithOptions(ctx, "host", "ahp-terminal:/t", ahp.SubscribeOptions{Window: &ahp.WindowOptions{}}) + if err != nil { + t.Fatal(err) + } + _, again, err := m.SubscribeWithOptions(ctx, "host", a.URI(), ahp.SubscribeOptions{Window: &ahp.WindowOptions{WindowBytes: 2048}}) + if err != nil { + t.Fatal(err) + } + if again != a { + t.Fatal("re-subscribe created a second owner") + } + _, b, err := m.SubscribeWithOptions(ctx, "other", a.URI(), ahp.SubscribeOptions{Window: &ahp.WindowOptions{}}) + if err != nil { + t.Fatal(err) + } + h, _ := m.ClientHandle("host") + c, _ := h.Client() + if err := c.Notify(ctx, "app/data", map[string]any{"channel": a.URI(), "text": "blocked"}); err != nil { + t.Fatal(err) + } + if a.FlowControl().Writable || !b.FlowControl().Writable { + t.Fatal("windows shared across hosts") + } +} diff --git a/clients/go/ahp/transport.go b/clients/go/ahp/transport.go index 914586710..27533c1fe 100644 --- a/clients/go/ahp/transport.go +++ b/clients/go/ahp/transport.go @@ -31,6 +31,49 @@ type TransportMessage struct { parsed *ahptypes.JsonRpcMessage text string binary []byte + raw *string +} + +// TransportCapabilities describes optional support for exact JSON text +// accounting. MaximumFrameBytes is measured in AHP UTF-16 bytes, not UTF-8. +type TransportCapabilities struct { + SupportsWindowedDelivery bool + MaximumFrameBytes int64 +} + +// TransportCapabilitiesProvider is optional; Transport itself is unchanged. +// Providers must retain received text and send text messages without rewriting. +type TransportCapabilitiesProvider interface { + TransportCapabilities() TransportCapabilities +} + +// MaximumFrameBytesForReadLimit converts a UTF-8 message limit to a safe AHP +// UTF-16 limit. Unknown or unlimited limits use 1 MiB. +func MaximumFrameBytesForReadLimit(limit int64) int64 { + if limit <= 0 { + return 1 << 20 + } + return 2 * (limit / 3) +} + +// NewParsedMessageWithRawText preserves the original text alongside its parse. +func NewParsedMessageWithRawText(m ahptypes.JsonRpcMessage, raw string) TransportMessage { + return TransportMessage{kind: transportFrameParsed, parsed: &m, raw: &raw} +} + +// RawJSONText returns the original received JSON, never a reserialization. +func (m TransportMessage) RawJSONText() (string, bool) { + switch m.kind { + case transportFrameText: + return m.text, true + case transportFrameBinary: + return string(m.binary), true + case transportFrameParsed: + if m.raw != nil { + return *m.raw, true + } + } + return "", false } // NewParsedMessage wraps an already-decoded JSON-RPC message. @@ -174,3 +217,11 @@ func (b *BoxedTransport) Recv(ctx context.Context) (TransportMessage, error) { func (b *BoxedTransport) Close(ctx context.Context) error { return b.inner.Close(ctx) } + +// TransportCapabilities preserves the optional capabilities of the inner transport. +func (b *BoxedTransport) TransportCapabilities() TransportCapabilities { + if p, ok := b.inner.(TransportCapabilitiesProvider); ok { + return p.TransportCapabilities() + } + return TransportCapabilities{} +} diff --git a/clients/go/ahp/windowing.go b/clients/go/ahp/windowing.go new file mode 100644 index 000000000..21aada148 --- /dev/null +++ b/clients/go/ahp/windowing.go @@ -0,0 +1,965 @@ +package ahp + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "strings" + "sync" + "time" + "unicode/utf8" + + "github.com/microsoft/agent-host-protocol/clients/go/ahptypes" +) + +// WindowingConfig supplies connection defaults. Windowing remains opt-in. +// All byte amounts use UTF-16 bytes of the complete JSON text. +type WindowingConfig struct { + WindowBytes int64 + MaximumFrameBytes int64 + MaxBufferedSendBytes int64 + Coalescers []Coalescer +} + +// ReceiveMode controls when an owner returns credit. +type ReceiveMode uint8 + +const ( + // ReceiveAuto releases the previous delivery when Next is called again. + ReceiveAuto ReceiveMode = iota + // ReceiveManual requires explicit Delivery.Release. + ReceiveManual +) + +// WindowOptions opts a subscription in. Zero limits inherit the connection +// defaults. A nil Coalescing enables built-in and configured coalescers. +type WindowOptions struct { + WindowBytes int64 + MaximumFrameBytes int64 + Receive ReceiveMode + Coalescing *bool +} + +// SubscribeOptions configures one subscription without changing legacy APIs. +type SubscribeOptions struct { + Window *WindowOptions + Delivery *ahptypes.SubscriptionDeliveryOptions + View *ahptypes.SubscribeView + owner *Subscription +} + +// FlowControlStatus is an immutable snapshot of subscription backpressure. +type FlowControlStatus struct { + Mode string `json:"mode"` + FallbackReason string `json:"fallbackReason,omitempty"` + Limits *ahptypes.ChannelFlowControl `json:"limits,omitempty"` + OutstandingSendBytes int64 `json:"outstandingSendBytes"` + BufferedSendBytes int64 `json:"bufferedSendBytes"` + BufferedReceiveBytes int64 `json:"bufferedReceiveBytes"` + Writable bool `json:"writable"` +} + +// FlowControlEvent is lossy telemetry, never a credit-owning delivery. +type FlowControlEvent struct { + Channel string `json:"channel"` + Kind string `json:"kind"` + Status FlowControlStatus `json:"status"` +} + +var ( + ErrBackpressure = errors.New("ahp: buffered send limit reached") + ErrSubscriptionDropped = errors.New("ahp: subscription was not resumed") +) + +const maxWindowInteger int64 = 9007199254740991 + +func utf16Bytes(text string) int64 { + var size int64 + for _, r := range text { + if r > 0xffff { + size += 4 + } else { + size += 2 + } + } + return size +} + +type channelFlow struct { + channel string + owner *Subscription + opts SubscribeOptions + limits *ahptypes.ChannelFlowControl + reason string + epoch uint64 + sent, consumed, received, released, advertised int64 + creditReady bool + creditTimer *time.Timer + expectSnapshot bool + fragments strings.Builder + fragmentBytes int64 + fragmenting bool + queue []*flowSend + buffered int64 + retainedSendBytes int64 + lastSeq int64 + dropped bool + reset bool + statusInitialized bool + lastWritable bool + accepted bool + requests map[*requestHooks]struct{} +} + +type flowEngine struct { + c *Client + mu sync.Mutex + cfg WindowingConfig + capabilities TransportCapabilities + channels map[string]*channelFlow + order []string + cursor int + creditCursor int + wake chan struct{} + listeners []chan FlowControlEvent + stopped bool + controlTokens float64 + controlAt time.Time + cleanup []flowCleanup +} + +func newFlowEngine(c *Client) (*flowEngine, error) { + e := &flowEngine{c: c, cfg: c.cfg.Windowing, channels: make(map[string]*channelFlow), wake: make(chan struct{}, 1)} + e.controlTokens, e.controlAt = 32, time.Now() + if p, ok := c.transport.(TransportCapabilitiesProvider); ok { + e.capabilities = p.TransportCapabilities() + } + if e.cfg.WindowBytes == 0 { + e.cfg.WindowBytes = 256 << 10 + } + if e.cfg.MaximumFrameBytes == 0 { + e.cfg.MaximumFrameBytes = e.capabilities.MaximumFrameBytes + if e.cfg.MaximumFrameBytes == 0 { + e.cfg.MaximumFrameBytes = 1 << 20 + } + } + if e.cfg.MaxBufferedSendBytes == 0 { + e.cfg.MaxBufferedSendBytes = 1 << 20 + } + if err := validLimits(ahptypes.ChannelReceiveLimits{WindowBytes: e.cfg.WindowBytes, MaximumFrameBytes: e.cfg.MaximumFrameBytes}); err != nil { + return nil, err + } + if e.cfg.MaxBufferedSendBytes < 1 || e.cfg.MaxBufferedSendBytes > maxWindowInteger { + return nil, fmt.Errorf("ahp: invalid send buffer limit") + } + e.cfg.Coalescers = append([]Coalescer(nil), e.cfg.Coalescers...) + return e, nil +} + +func validLimits(l ahptypes.ChannelReceiveLimits) error { + if l.WindowBytes < 1 || l.WindowBytes > maxWindowInteger || l.MaximumFrameBytes < 1024 || l.MaximumFrameBytes > maxWindowInteger { + return fmt.Errorf("ahp: invalid flow-control limits") + } + return nil +} + +func validateAccepted(l *ahptypes.ChannelFlowControl, offer *ahptypes.SubscriptionFlowControlOptions) error { + if offer == nil { + return fmt.Errorf("ahp: unsolicited flow control") + } + if err := validLimits(l.ClientReceive); err != nil { + return err + } + if err := validLimits(l.HostReceive); err != nil { + return err + } + if l.ClientReceive.WindowBytes > offer.Receive.WindowBytes || l.ClientReceive.MaximumFrameBytes > offer.Receive.MaximumFrameBytes { + return fmt.Errorf("ahp: accepted limits exceed the offer") + } + return nil +} + +func copyOptions(o SubscribeOptions) SubscribeOptions { + if o.Window != nil { + w := *o.Window + if w.Coalescing != nil { + b := *w.Coalescing + w.Coalescing = &b + } + o.Window = &w + } + if o.Delivery != nil { + d := *o.Delivery + if d.MaxLatencyMs != nil { + n := *d.MaxLatencyMs + d.MaxLatencyMs = &n + } + o.Delivery = &d + } + if o.View != nil { + v := *o.View + if v.Turns != nil { + n := *v.Turns + v.Turns = &n + } + o.View = &v + } + return o +} + +func (e *flowEngine) signal() { + select { + case e.wake <- struct{}{}: + default: + } +} + +func (e *flowEngine) statusLocked(f *channelFlow) FlowControlStatus { + receiveBytes := f.received - f.released + closed := false + if f.owner != nil { + f.owner.owner.mu.Lock() + receiveBytes = f.owner.owner.unreleased + closed = f.owner.owner.final + f.owner.owner.mu.Unlock() + } + s := FlowControlStatus{Mode: "ordinary", FallbackReason: f.reason, BufferedSendBytes: f.buffered, BufferedReceiveBytes: receiveBytes, Writable: !closed && !e.stopped && !f.dropped && !f.reset} + if f.limits != nil { + l := *f.limits + s.Mode, s.Limits = "windowed", &l + s.OutstandingSendBytes = f.sent - f.consumed + s.Writable = !closed && !e.stopped && !f.dropped && !f.reset && s.OutstandingSendBytes < l.HostReceive.WindowBytes && f.buffered < e.cfg.MaxBufferedSendBytes && f.retainedSendBytes < e.cfg.MaxBufferedSendBytes + } + return s +} + +func (e *flowEngine) emitLocked(f *channelFlow, kind string) { + status := e.statusLocked(f) + if f.owner != nil && f.owner.owner != nil { + f.owner.owner.publish(status) + } + emit := func(kind string) { + if kind == "" { + return + } + for _, ch := range e.listeners { + select { + case ch <- FlowControlEvent{Channel: f.channel, Kind: kind, Status: status}: + default: + } + } + } + emit(kind) + if f.statusInitialized && !e.stopped && !f.dropped && f.lastWritable != status.Writable { + transition := "blocked" + if status.Writable { + transition = "unblocked" + } + if kind != transition { + emit(transition) + } + } + f.statusInitialized, f.lastWritable = true, status.Writable +} + +// FlowControlEvents returns a lossy stream of connection-level telemetry. +func (c *Client) FlowControlEvents() <-chan FlowControlEvent { + c.flow.mu.Lock() + defer c.flow.mu.Unlock() + ch := make(chan FlowControlEvent, 64) + if c.flow.stopped { + close(ch) + } else { + c.flow.listeners = append(c.flow.listeners, ch) + } + return ch +} + +func (e *flowEngine) getLocked(uri string) *channelFlow { + if f := e.channels[uri]; f != nil && !f.dropped { + return f + } + existed := e.channels[uri] != nil + f := &channelFlow{channel: uri} + e.channels[uri] = f + if !existed { + e.order = append(e.order, uri) + } + return f +} + +func (e *flowEngine) attachLocked(f *channelFlow, sub *Subscription) { + if sub == nil { + sub = &Subscription{uri: f.channel, events: make(chan SubscriptionEvent, e.c.cfg.SubscriptionBuffer), owner: newDeliveryQueue()} + } + f.owner = sub + sub.owner.bind(e, f) + e.c.subscriptionsMu.Lock() + found := false + for _, s := range e.c.subscriptions[f.channel] { + if s == sub { + found = true + } + } + if !found { + e.c.subscriptions[f.channel] = append(e.c.subscriptions[f.channel], sub) + } + e.c.subscriptionsMu.Unlock() +} + +// SubscribeWithOptions creates or updates the single owner for a channel. +// Use Next for credit-owning consumption; Events remains an observe-only stream. +func (c *Client) SubscribeWithOptions(ctx context.Context, uri string, opts SubscribeOptions) (*ahptypes.SubscribeResult, *Subscription, error) { + if err := ctx.Err(); err != nil { + return nil, nil, err + } + opts = copyOptions(opts) + c.flow.mu.Lock() + if c.flow.stopped { + c.flow.mu.Unlock() + return nil, nil, ErrShutdown + } + if opts.owner != nil { + current := c.flow.channels[uri] + if current == nil || current.owner != opts.owner || opts.owner.owner.isFinal() { + c.flow.mu.Unlock() + return nil, opts.owner, ErrClosed + } + } + f := c.flow.getLocked(uri) + created := f.owner == nil + if f.owner == nil { + c.flow.attachLocked(f, nil) + } + sub := f.owner + c.flow.mu.Unlock() + select { + case sub.owner.updateGate <- struct{}{}: + defer func() { <-sub.owner.updateGate }() + case <-ctx.Done(): + return nil, sub, ctx.Err() + } + opts.owner = sub + var result ahptypes.SubscribeResult + err := c.request(ctx, "subscribe", ahptypes.SubscribeParams{Channel: uri, Delivery: opts.Delivery, View: opts.View}, &result, &opts) + if err != nil { + if created { + sub.Close() + return nil, nil, err + } + c.flow.mu.Lock() + active := c.flow.channels[uri] == f && f.owner == sub && !sub.owner.isFinal() + c.flow.mu.Unlock() + if !active { + return nil, sub, ErrClosed + } + return nil, sub, err + } + return &result, sub, nil +} + +// Update renegotiates the same subscription. Existing deliveries remain ordered. +func (s *Subscription) Update(ctx context.Context, opts SubscribeOptions) (*ahptypes.SubscribeResult, error) { + if s.owner == nil { + return nil, fmt.Errorf("ahp: Update requires an options-based subscription") + } + e, _ := s.owner.binding() + opts.owner = s + r, _, err := e.c.SubscribeWithOptions(ctx, s.uri, opts) + return r, err +} + +// ResumeSubscription is used by managed hosts after a fresh initialize. It +// preserves the owner's unread queue and requests a new snapshot/credit epoch. +// The old Client must be shut down before calling it. +func (c *Client) ResumeSubscription(ctx context.Context, s *Subscription, opts SubscribeOptions) (*ahptypes.SubscribeResult, error) { + if s.owner == nil { + return nil, fmt.Errorf("ahp: not an options-based subscription") + } + s.owner.mu.Lock() + final := s.owner.final + s.owner.mu.Unlock() + if final { + return nil, ErrClosed + } + old, _ := s.owner.binding() + select { + case <-old.c.Done(): + default: + return nil, fmt.Errorf("ahp: old subscription connection is still open") + } + s.closeMu.Lock() + s.closed = false + s.events = make(chan SubscriptionEvent, c.cfg.SubscriptionBuffer) + s.closeMu.Unlock() + c.flow.mu.Lock() + f := c.flow.getLocked(s.uri) + if f.owner != nil && f.owner != s { + c.flow.mu.Unlock() + return nil, fmt.Errorf("ahp: channel already has an owner") + } + c.flow.attachLocked(f, s) + c.flow.mu.Unlock() + r, _, err := c.SubscribeWithOptions(ctx, s.uri, opts) + return r, err +} + +func (e *flowEngine) prepareRequest(method string, raw json.RawMessage, options *SubscribeOptions) (json.RawMessage, *requestHooks, error) { + if method == "initialize" { + var p ahptypes.InitializeParams + if err := json.Unmarshal(raw, &p); err != nil { + return raw, nil, err + } + for _, uri := range p.InitialSubscriptions { + if strings.HasPrefix(uri, "ahp-tcp:") { + return raw, nil, fmt.Errorf("ahp: TCP cannot be an initial subscription") + } + } + } + if method == "reconnect" { + return e.prepareReconnect(raw) + } + if method != "subscribe" { + return raw, nil, nil + } + var p ahptypes.SubscribeParams + if err := json.Unmarshal(raw, &p); err != nil { + return raw, nil, err + } + e.mu.Lock() + defer e.mu.Unlock() + if e.stopped { + return raw, nil, ErrShutdown + } + if options == nil && p.FlowControl == nil && e.channels[p.Channel] == nil { + return raw, nil, nil + } + if options != nil && options.owner != nil { + current := e.channels[p.Channel] + if current == nil || current.owner != options.owner || options.owner.owner.isFinal() { + return raw, nil, ErrClosed + } + } + f := e.getLocked(p.Channel) + if f.owner == nil { + e.attachLocked(f, nil) + } + validated := false + defer func() { + if !validated && !f.accepted && len(f.requests) == 0 { + e.retireLocked(f, ErrClosed, false) + } + }() + opts := copyOptions(f.opts) + if options != nil { + opts = copyOptions(*options) + } else { + opts.Delivery, opts.View = p.Delivery, p.View + if p.FlowControl != nil { + opts.Window = &WindowOptions{WindowBytes: p.FlowControl.Receive.WindowBytes, MaximumFrameBytes: p.FlowControl.Receive.MaximumFrameBytes} + } else { + opts.Window = nil + } + } + reason := "" + if opts.Window != nil { + w := opts.Window + if w.Receive != ReceiveAuto && w.Receive != ReceiveManual { + return raw, nil, fmt.Errorf("ahp: invalid receive mode") + } + window, frame := w.WindowBytes, w.MaximumFrameBytes + if window == 0 { + window = e.cfg.WindowBytes + } + if frame == 0 { + frame = e.cfg.MaximumFrameBytes + } + p.FlowControl = &ahptypes.SubscriptionFlowControlOptions{Receive: ahptypes.ChannelReceiveLimits{WindowBytes: window, MaximumFrameBytes: frame}} + if err := validLimits(p.FlowControl.Receive); err != nil { + return raw, nil, err + } + if !e.capabilities.SupportsWindowedDelivery { + p.FlowControl = nil + reason = "transport" + } + } + if options != nil && options.Window == nil { + p.FlowControl = nil + } + if strings.HasPrefix(p.Channel, "ahp-tcp:") && p.FlowControl == nil { + return raw, nil, fmt.Errorf("ahp: TCP requires windowed delivery") + } + offer := p.FlowControl + encoded, err := json.Marshal(p) + if err != nil { + return raw, nil, err + } + owner := f.owner + h := &requestHooks{done: make(chan struct{}), flows: map[*channelFlow]struct{}{f: {}}} + if f.requests == nil { + f.requests = make(map[*requestHooks]struct{}) + } + f.requests[h] = struct{}{} + h.cancel = func() { + e.mu.Lock() + defer e.mu.Unlock() + e.retireLocked(f, ErrClosed, true) + } + h.failed = func(err error) { + e.mu.Lock() + defer e.mu.Unlock() + delete(f.requests, h) + if !f.accepted { + e.retireLocked(f, err, true) + } + } + h.complete = func(result json.RawMessage) error { + e.mu.Lock() + defer e.mu.Unlock() + if e.stopped || e.channels[p.Channel] != f || f.owner != owner || owner.owner.isFinal() { + return nil + } + var r ahptypes.SubscribeResult + if err := json.Unmarshal(result, &r); err != nil { + return err + } + if r.FlowControl != nil { + if err := validateAccepted(r.FlowControl, offer); err != nil { + return err + } + if r.Snapshot != nil { + return fmt.Errorf("ahp: windowed result also contains a snapshot") + } + minimum, err := frameText(p.Channel, "A", false) + if err != nil || utf16Bytes(minimum) > r.FlowControl.ClientReceive.MaximumFrameBytes || utf16Bytes(minimum) > r.FlowControl.HostReceive.MaximumFrameBytes { + return fmt.Errorf("ahp: accepted frame cannot carry a character") + } + } else if strings.HasPrefix(p.Channel, "ahp-tcp:") { + return fmt.Errorf("ahp: host declined mandatory TCP windowing") + } + if (f.limits == nil) != (r.FlowControl == nil) { + f.epoch++ + if f.creditTimer != nil { + f.creditTimer.Stop() + f.creditTimer = nil + } + f.creditReady = false + f.sent, f.consumed, f.received, f.released, f.advertised = 0, 0, 0, 0, 0 + } + f.limits, f.opts, f.reason = r.FlowControl, opts, reason + f.accepted = true + delete(f.requests, h) + f.dropped = false + if opts.Window != nil && r.FlowControl == nil && reason == "" { + f.reason = "declined" + } + f.expectSnapshot = r.FlowControl != nil + owner.owner.setMode(opts) + kind := "negotiated" + if f.reason != "" { + kind = "declined" + } + e.emitLocked(f, kind) + if r.Snapshot != nil { + n := snapshotNotification(p.Channel, *r.Snapshot) + if err := e.enqueueLocked(f, n, 0, 0); err != nil { + return err + } + } + e.signal() + return nil + } + validated = true + return encoded, h, nil +} + +func snapshotNotification(uri string, snapshot ahptypes.Snapshot) ahptypes.JsonRpcNotification { + raw, _ := json.Marshal(ahptypes.ChannelSnapshotParams{Channel: uri, Snapshot: snapshot}) + return ahptypes.JsonRpcNotification{JsonRpc: ahptypes.JsonRpcV2, Method: "channel/snapshot", Params: raw} +} + +func (e *flowEngine) hasWindows() bool { + e.mu.Lock() + defer e.mu.Unlock() + for _, f := range e.channels { + if f.limits != nil { + return true + } + } + return false +} + +func notificationChannel(n ahptypes.JsonRpcNotification) (string, error) { + var p struct { + Channel string `json:"channel"` + } + err := json.Unmarshal(n.Params, &p) + return p.Channel, err +} + +func (e *flowEngine) receive(n ahptypes.JsonRpcNotification, wire TransportMessage) (bool, error) { + channel, err := notificationChannel(n) + if err != nil { + return false, nil + } + e.mu.Lock() + defer e.mu.Unlock() + if e.stopped { + return true, nil + } + f := e.channels[channel] + if f == nil { + return false, nil + } + if f.dropped { + return true, nil + } + if f.owner != nil && f.owner.owner.isFinal() { + e.retireLocked(f, ErrClosed, true) + return true, nil + } + if !f.accepted { + return true, nil + } + if f.reset && ahptypes.CountsAgainstWindow(n.Method) { + return true, fmt.Errorf("ahp: data after TCP reset") + } + if f.limits != nil && n.JsonRpc != ahptypes.JsonRpcV2 { + return true, fmt.Errorf("ahp: invalid JSON-RPC notification") + } + if !ahptypes.CountsAgainstWindow(n.Method) { + if raw, ok := wire.RawJSONText(); ok && utf16Bytes(raw) > 1<<20 { + return true, fmt.Errorf("ahp: control message too large") + } + } + if f.limits == nil && n.Method == "channel/credit" { + return true, nil + } + if n.Method == "channel/credit" { + var p ahptypes.ChannelCreditParams + if err := json.Unmarshal(n.Params, &p); err != nil { + return true, err + } + if p.ConsumedBytes < 0 || p.ConsumedBytes > f.sent { + return true, fmt.Errorf("ahp: credit exceeds sent bytes") + } + if p.ConsumedBytes > f.consumed { + f.consumed = p.ConsumedBytes + e.emitLocked(f, "creditReceived") + e.signal() + } + return true, nil + } + var cost int64 + if f.limits != nil && ahptypes.CountsAgainstWindow(n.Method) { + text, ok := wire.RawJSONText() + if !ok || !utf8.ValidString(text) { + return true, fmt.Errorf("ahp: raw UTF-8 JSON text required") + } + cost = utf16Bytes(text) + if cost > f.limits.ClientReceive.MaximumFrameBytes { + return true, fmt.Errorf("ahp: data exceeds maximumFrameBytes") + } + if f.received > maxWindowInteger-cost { + return true, fmt.Errorf("ahp: received counter overflow") + } + f.received += cost + } + if n.Method == "channel/frame" { + if f.limits == nil { + return true, fmt.Errorf("ahp: fragment on ordinary subscription") + } + var p ahptypes.ChannelFrameParams + if err := json.Unmarshal(n.Params, &p); err != nil { + return true, err + } + if p.Data == "" || !validFragmentString(n.Params) { + return true, fmt.Errorf("ahp: invalid fragment data") + } + f.fragmenting = true + f.fragments.WriteString(p.Data) + f.fragmentBytes += cost + if p.Final == nil || !*p.Final { + e.releaseLocked(f, cost, true) + return true, nil + } + var joined ahptypes.JsonRpcMessage + if err := json.Unmarshal([]byte(f.fragments.String()), &joined); err != nil { + return true, fmt.Errorf("ahp: invalid joined JSON: %w", err) + } + if joined.Notification == nil || joined.Notification.JsonRpc != ahptypes.JsonRpcV2 { + return true, fmt.Errorf("ahp: frame contains a request or response") + } + n = *joined.Notification + innerChannel, err := notificationChannel(n) + if err != nil || innerChannel != channel || n.Method == "channel/frame" || !ahptypes.CountsAgainstWindow(n.Method) { + return true, fmt.Errorf("ahp: invalid joined channel or method") + } + total := f.fragmentBytes + f.fragments.Reset() + f.fragmentBytes = 0 + f.fragmenting = false + e.emitLocked(f, "fragmentedMessageReceived") + return true, e.enqueueLocked(f, n, total, cost) + } + if f.fragmenting && ahptypes.CountsAgainstWindow(n.Method) { + return true, fmt.Errorf("ahp: interleaved fragmented message") + } + if (n.Method == "channel/snapshot") && f.limits == nil { + return true, fmt.Errorf("ahp: unsolicited windowed snapshot") + } + if n.Method == "tcp/reset" { + e.failQueueLocked(f, ErrClosed) + f.reset = true + f.epoch++ + if f.owner != nil { + q := f.owner.owner + q.mu.Lock() + q.items, q.previous = nil, nil + q.unreleased = 0 + q.releaseVersion++ + q.mu.Unlock() + } + e.releaseLocked(f, f.received-f.released, true) + f.fragments.Reset() + f.fragmenting = false + f.fragmentBytes = 0 + } + return true, e.enqueueLocked(f, n, cost, cost) +} + +// encoding/json accepts unpaired escaped surrogates by replacing them; fragments +// cannot split a surrogate pair, so validate their raw JSON strings first. +func validFragmentString(params []byte) bool { + var raw map[string]json.RawMessage + if json.Unmarshal(params, &raw) != nil { + return false + } + b := raw["data"] + for i := 0; i < len(b); i++ { + if b[i] != '\\' { + continue + } + i++ + if i >= len(b) { + return false + } + if b[i] != 'u' { + continue + } + read := func(pos int) (uint16, bool) { + if pos+4 >= len(b) { + return 0, false + } + var n uint16 + for _, c := range b[pos+1 : pos+5] { + n <<= 4 + switch { + case c >= '0' && c <= '9': + n += uint16(c - '0') + case c >= 'a' && c <= 'f': + n += uint16(c - 'a' + 10) + case c >= 'A' && c <= 'F': + n += uint16(c - 'A' + 10) + default: + return 0, false + } + } + return n, true + } + n, ok := read(i) + if !ok { + return false + } + i += 4 + if n >= 0xdc00 && n <= 0xdfff { + return false + } + if n >= 0xd800 && n <= 0xdbff { + if i+2 >= len(b) || b[i+1] != '\\' || b[i+2] != 'u' { + return false + } + low, ok := read(i + 2) + if !ok || low < 0xdc00 || low > 0xdfff { + return false + } + i += 6 + } + } + return true +} + +func (e *flowEngine) enqueueLocked(f *channelFlow, n ahptypes.JsonRpcNotification, total, credit int64) error { + if f.expectSnapshot && ahptypes.CountsAgainstWindow(n.Method) { + if n.Method != "channel/snapshot" { + return fmt.Errorf("ahp: first data must be channel/snapshot") + } + f.expectSnapshot = false + } + ev, err := eventForNotification(n) + if err != nil { + return err + } + if p, ok := ev.(SubscriptionEventSnapshot); ok { + if p.Params.Snapshot.Resource != f.channel { + return fmt.Errorf("ahp: snapshot resource mismatch") + } + } + epoch := f.epoch + if f.owner != nil { + q := f.owner.owner + q.mu.Lock() + if q.closed { + q.mu.Unlock() + e.releaseLocked(f, credit, false) + e.c.fanOut(f.channel, ev) + return nil + } + version := q.releaseVersion + d := newDelivery(n, ev, total, e.deliveryRelease(f, epoch, credit, q, version)) + d.credit, d.source = credit, f + q.unreleased += credit + if !q.closed { + if f.limits != nil && coalescingEnabled(f.opts) && len(q.items) > 0 { + last := q.items[len(q.items)-1] + if last.epoch == epoch && last.source == f { + if merged, ok := e.coalesce(last.notification, n, CoalesceReceive, f); ok { + event, err := eventForNotification(merged) + if err != nil { + q.mu.Unlock() + return err + } + old := last + combinedCredit := old.credit + credit + d = newDelivery(merged, event, old.bytes+total, e.deliveryRelease(f, epoch, combinedCredit, q, version)) + d.epoch, d.credit, d.source = epoch, combinedCredit, f + q.items[len(q.items)-1] = d + q.signalLocked() + q.mu.Unlock() + keepSequence(f, event) + e.emitLocked(f, "coalesced") + e.c.fanOut(f.channel, ev) + return nil + } + } + } + d.epoch = epoch + if f.limits != nil || len(q.items) < e.c.cfg.SubscriptionBuffer { + q.items = append(q.items, d) + q.signalLocked() + keepSequence(f, ev) + } + + } + q.mu.Unlock() + } else { + e.releaseLocked(f, credit, false) + } + e.emitLocked(f, "") + e.c.fanOut(f.channel, ev) + return nil +} + +func keepSequence(f *channelFlow, event SubscriptionEvent) { + var seq int64 + switch v := event.(type) { + case SubscriptionEventSnapshot: + seq = v.Params.Snapshot.FromSeq + case SubscriptionEventAction: + seq = v.Envelope.ServerSeq + } + if seq > f.lastSeq { + f.lastSeq = seq + } +} + +func (e *flowEngine) deliveryRelease(f *channelFlow, epoch uint64, credit int64, q *deliveryQueue, version uint64) func() { + return func() { + e.mu.Lock() + q.mu.Lock() + current, currentFlow := q.engine, q.flow + changed := q.releaseVersion == version + if changed { + q.unreleased -= credit + } + q.mu.Unlock() + if changed && !e.stopped && !f.dropped && f.epoch == epoch { + e.releaseLocked(f, credit, false) + } else if current == e && changed { + e.emitLocked(f, "") + } + e.mu.Unlock() + if current != e && changed { + current.mu.Lock() + current.emitLocked(currentFlow, "") + current.mu.Unlock() + } + } +} + +func coalescingEnabled(opts SubscribeOptions) bool { + return opts.Window != nil && (opts.Window.Coalescing == nil || *opts.Window.Coalescing) +} + +func (e *flowEngine) releaseLocked(f *channelFlow, bytes int64, immediate bool) { + if bytes == 0 || f.limits == nil { + return + } + f.released += bytes + threshold := f.limits.ClientReceive.WindowBytes / 4 + if threshold < 1 { + threshold = 1 + } + if threshold > 16<<10 { + threshold = 16 << 10 + } + if immediate || f.released-f.advertised >= threshold || f.received-f.advertised >= f.limits.ClientReceive.WindowBytes { + f.creditReady = true + e.signal() + } else if f.creditTimer == nil { + epoch := f.epoch + f.creditTimer = time.AfterFunc(5*time.Millisecond, func() { + e.mu.Lock() + defer e.mu.Unlock() + if f.epoch != epoch { + return + } + f.creditTimer = nil + if !e.stopped && f.epoch == epoch && f.released > f.advertised { + f.creditReady = true + e.signal() + } + }) + } + e.emitLocked(f, "") +} + +func (e *flowEngine) stop() { + e.mu.Lock() + defer e.mu.Unlock() + if e.stopped { + return + } + e.stopped = true + for _, f := range e.channels { + if f.creditTimer != nil { + f.creditTimer.Stop() + } + f.fragments.Reset() + f.fragmenting = false + f.fragmentBytes = 0 + e.failQueueLocked(f, ErrShutdown) + e.emitLocked(f, "dropped") + for h := range f.requests { + h.drop(ErrShutdown) + } + } + e.channels = make(map[string]*channelFlow) + e.order = nil + e.cleanup = nil + for _, ch := range e.listeners { + close(ch) + } + e.listeners = nil +} diff --git a/clients/go/ahp/windowing_corpus_test.go b/clients/go/ahp/windowing_corpus_test.go new file mode 100644 index 000000000..059e01870 --- /dev/null +++ b/clients/go/ahp/windowing_corpus_test.go @@ -0,0 +1,467 @@ +package ahp + +import ( + "context" + "encoding/json" + "os" + "path/filepath" + "reflect" + "strings" + "sync" + "testing" + "time" + + "github.com/microsoft/agent-host-protocol/clients/go/ahptypes" +) + +type corpusTransport struct { + in chan TransportMessage + sent chan string + done chan struct{} + once sync.Once + capable bool +} + +func newCorpusTransport() *corpusTransport { + return &corpusTransport{in: make(chan TransportMessage, 128), sent: make(chan string, 128), done: make(chan struct{}), capable: true} +} +func (t *corpusTransport) Send(ctx context.Context, m TransportMessage) error { + b, _, err := m.Bytes() + if err != nil { + return err + } + select { + case t.sent <- string(b): + return nil + case <-ctx.Done(): + return ctx.Err() + case <-t.done: + return ErrClosed + } +} +func (t *corpusTransport) Recv(ctx context.Context) (TransportMessage, error) { + select { + case m := <-t.in: + return m, nil + case <-ctx.Done(): + return TransportMessage{}, ctx.Err() + case <-t.done: + return TransportMessage{}, ErrClosed + } +} +func (t *corpusTransport) Close(context.Context) error { + t.once.Do(func() { close(t.done) }) + return nil +} +func (t *corpusTransport) TransportCapabilities() TransportCapabilities { + return TransportCapabilities{SupportsWindowedDelivery: t.capable, MaximumFrameBytes: 16384} +} + +func corpusDir(t *testing.T) string { + return filepath.Join(findFixtureDir(t), "..", "windowed-delivery") +} +func TestWindowedDeliverySizes(t *testing.T) { + data, err := os.ReadFile(filepath.Join(corpusDir(t), "sizes.json")) + if err != nil { + t.Fatal(err) + } + var corpus struct { + Cases []struct { + Name, Text string + UTF16Bytes int64 `json:"utf16Bytes"` + UTF8Bytes int `json:"utf8Bytes"` + } + } + if err := json.Unmarshal(data, &corpus); err != nil { + t.Fatal(err) + } + if len(corpus.Cases) == 0 { + t.Fatal("empty size corpus") + } + for _, c := range corpus.Cases { + t.Run(c.Name, func(t *testing.T) { + if got := utf16Bytes(c.Text); got != c.UTF16Bytes { + t.Fatalf("size %d, want %d", got, c.UTF16Bytes) + } + if len(c.Text) != c.UTF8Bytes { + t.Fatal("UTF-8 size mismatch") + } + }) + } +} + +type corpusStep struct { + Op string + Channel string + Window *WindowOptions + Receive string + Coalescing *bool + To string + Result json.RawMessage + Text string + Match any + Expect any + Absent []string + Reassemble bool + WithinMs int + ExpectBytes *int64 + Index int + Method string + Params json.RawMessage +} +type corpusRecord struct { + text string + object map[string]any + answered bool +} +type corpusSubscription struct { + ready chan struct{} + sub *Subscription + err error +} +type corpusDriver struct { + t *testing.T + ctx context.Context + client *Client + transport *corpusTransport + pending []*corpusRecord + history []*corpusRecord + subs map[string]*corpusSubscription + deliveries map[string][]*Delivery +} + +func partialMatch(actual, expected any) bool { + if object, ok := expected.(map[string]any); ok { + got, ok := actual.(map[string]any) + if !ok { + return false + } + for k, v := range object { + if !partialMatch(got[k], v) { + return false + } + } + return true + } + return reflect.DeepEqual(actual, expected) +} +func jsonObject(value any) map[string]any { + b, _ := json.Marshal(value) + var out map[string]any + _ = json.Unmarshal(b, &out) + return out +} +func (d *corpusDriver) record(text string) *corpusRecord { + var object map[string]any + if err := json.Unmarshal([]byte(text), &object); err != nil { + d.t.Fatal(err) + } + r := &corpusRecord{text: text, object: object} + d.history = append(d.history, r) + return r +} +func (d *corpusDriver) take(ctx context.Context) *corpusRecord { + if len(d.pending) > 0 { + r := d.pending[0] + d.pending = d.pending[1:] + return r + } + select { + case text := <-d.transport.sent: + return d.record(text) + case <-ctx.Done(): + d.t.Fatalf("waiting for sent message: %v (client error %v)", ctx.Err(), d.client.Err()) + return nil + } +} +func (d *corpusDriver) subscription(uri string) *Subscription { + d.t.Helper() + s := d.subs[uri] + if s == nil { + d.t.Fatalf("no subscription %s", uri) + } + select { + case <-s.ready: + if s.err != nil { + d.t.Fatal(s.err) + } + return s.sub + case <-time.After(time.Second): + d.t.Fatalf("subscribe did not complete: %v", d.client.Err()) + return nil + } +} +func (d *corpusDriver) inject(text string) { + select { + case d.transport.in <- NewTextMessage(text): + case <-d.ctx.Done(): + d.t.Fatal(d.ctx.Err()) + } +} +func (d *corpusDriver) expectSent(s corpusStep) { + ctx, cancel := context.WithTimeout(d.ctx, time.Second) + defer cancel() + var joined strings.Builder + for { + r := d.take(ctx) + method, _ := r.object["method"].(string) + if s.Reassemble && method == "channel/frame" { + params := r.object["params"].(map[string]any) + channel := params["channel"].(string) + status := d.subscription(channel).FlowControl() + if utf16Bytes(r.text) > status.Limits.HostReceive.MaximumFrameBytes { + d.t.Fatal("oversized outgoing fragment") + } + joined.WriteString(params["data"].(string)) + if params["final"] != true { + continue + } + var message any + if err := json.Unmarshal([]byte(joined.String()), &message); err != nil { + d.t.Fatal(err) + } + if !partialMatch(message, s.Match) { + d.t.Fatalf("reassembled mismatch: %v != %v", message, s.Match) + } + return + } + if partialMatch(r.object, s.Match) { + for _, path := range s.Absent { + var cur any = r.object + for _, key := range strings.Split(path, ".") { + m, ok := cur.(map[string]any) + if !ok { + cur = nil + break + } + cur = m[key] + } + if cur != nil { + d.t.Fatalf("unexpected field %s", path) + } + } + return + } + if method == "channel/credit" { + continue + } + d.t.Fatalf("unexpected sent message %s; want %v", r.text, s.Match) + } +} +func (d *corpusDriver) step(s corpusStep) { + d.t.Helper() + switch s.Op { + case "subscribe": + opts := SubscribeOptions{Window: s.Window} + if opts.Window != nil { + opts.Window.Coalescing = s.Coalescing + if s.Receive == "manual" { + opts.Window.Receive = ReceiveManual + } + } + old := d.subs[s.Channel] + state := &corpusSubscription{ready: make(chan struct{})} + d.subs[s.Channel] = state + go func() { + if old != nil { + <-old.ready + state.sub = old.sub + if old.err != nil { + state.err = old.err + } else { + _, state.err = old.sub.Update(d.ctx, opts) + } + } else { + _, state.sub, state.err = d.client.SubscribeWithOptions(d.ctx, s.Channel, opts) + } + close(state.ready) + }() + case "respond": + var request *corpusRecord + for _, r := range d.history { + if r.object["method"] == s.To && r.object["id"] != nil && !r.answered { + request = r + break + } + } + if request == nil { + d.t.Fatalf("no unanswered %s", s.To) + } + request.answered = true + b, _ := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": request.object["id"], "result": s.Result}) + d.inject(string(b)) + case "unsubscribe": + if err := d.client.Unsubscribe(d.ctx, s.Channel); err != nil { + d.t.Fatal(err) + } + delete(d.subs, s.Channel) + case "hostSends": + d.inject(s.Text) + case "expectSent": + d.expectSent(s) + case "expectNotSent": + for _, r := range d.pending { + if partialMatch(r.object, s.Match) { + d.t.Fatalf("already sent %s", r.text) + } + } + timer := time.NewTimer(time.Duration(s.WithinMs) * time.Millisecond) + defer timer.Stop() + for { + select { + case text := <-d.transport.sent: + r := d.record(text) + d.pending = append(d.pending, r) + if partialMatch(r.object, s.Match) { + d.t.Fatalf("unexpected send %s", text) + } + case <-timer.C: + return + } + } + case "next": + ctx, cancel := context.WithTimeout(d.ctx, time.Second) + defer cancel() + item, err := d.subscription(s.Channel).Next(ctx) + if err != nil { + d.t.Fatal(err) + } + if !partialMatch(jsonObject(item.Notification()), s.Expect) { + d.t.Fatalf("delivery %s, want %v", mustJSON(item.Notification()), s.Expect) + } + if s.ExpectBytes != nil && item.WireBytes() != *s.ExpectBytes { + d.t.Fatalf("delivery bytes %d, want %d", item.WireBytes(), *s.ExpectBytes) + } + d.deliveries[s.Channel] = append(d.deliveries[s.Channel], item) + case "advance": + items := d.deliveries[s.Channel] + if len(items) == 0 { + d.t.Fatal("no previous delivery") + } + items[len(items)-1].Release() + case "release": + d.deliveries[s.Channel][s.Index].Release() + case "send": + d.subscription(s.Channel) + ctx, cancel := context.WithCancel(d.ctx) + d.t.Cleanup(cancel) + done := make(chan error, 1) + go func() { done <- d.client.Notify(ctx, s.Method, s.Params) }() + // Ensure sequential scenario sends are admitted in the scripted order. + deadline := time.Now().Add(time.Second) + for { + select { + case err := <-done: + if err != nil { + d.t.Fatal(err) + } + return + default: + } + d.client.flow.mu.Lock() + admitted := false + if f := d.client.flow.channels[s.Channel]; f != nil { + for _, m := range f.queue { + for _, p := range m.parts { + if p.ctx == ctx { + admitted = true + } + } + } + } + d.client.flow.mu.Unlock() + if admitted { + return + } + if time.Now().After(deadline) { + d.t.Fatal("send not admitted") + } + time.Sleep(time.Millisecond) + } + case "hostCreditsAll": + for { + select { + case text := <-d.transport.sent: + r := d.record(text) + d.pending = append(d.pending, r) + default: + goto drained + } + } + drained: + var total int64 + for _, r := range d.history { + method, _ := r.object["method"].(string) + params, _ := r.object["params"].(map[string]any) + if r.object["id"] == nil && params["channel"] == s.Channel && ahptypes.CountsAgainstWindow(method) { + total += utf16Bytes(r.text) + } + } + b, _ := json.Marshal(map[string]any{"jsonrpc": "2.0", "method": "channel/credit", "params": map[string]any{"channel": s.Channel, "consumedBytes": total}}) + d.inject(string(b)) + case "expectStatus": + sub := d.subscription(s.Channel) + deadline := time.Now().Add(time.Second) + for { + if partialMatch(jsonObject(sub.FlowControl()), s.Match) { + return + } + if time.Now().After(deadline) { + d.t.Fatalf("status %s, want %v", mustJSON(sub.FlowControl()), s.Match) + } + time.Sleep(time.Millisecond) + } + case "expectClosed": + select { + case <-d.client.Done(): + if d.client.Err() == nil { + d.t.Fatal("missing protocol error") + } + case <-time.After(time.Second): + d.t.Fatal("connection did not close") + } + default: + d.t.Fatalf("unsupported corpus step %s", s.Op) + } +} +func mustJSON(v any) string { b, _ := json.Marshal(v); return string(b) } + +func TestWindowedDeliveryCorpus(t *testing.T) { + files, err := filepath.Glob(filepath.Join(corpusDir(t), "scenarios", "*.json")) + if err != nil { + t.Fatal(err) + } + if len(files) == 0 { + t.Fatal("empty windowed delivery scenario corpus") + } + for _, file := range files { + t.Run(filepath.Base(file), func(t *testing.T) { + data, err := os.ReadFile(file) + if err != nil { + t.Fatal(err) + } + var scenario struct { + Windowing WindowingConfig + Steps []corpusStep + } + if err := json.Unmarshal(data, &scenario); err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + transport := newCorpusTransport() + cfg := DefaultConfig() + cfg.Windowing = scenario.Windowing + client, err := Connect(ctx, transport, cfg) + if err != nil { + t.Fatal(err) + } + defer client.Shutdown(context.Background()) + d := &corpusDriver{t: t, ctx: ctx, client: client, transport: transport, subs: make(map[string]*corpusSubscription), deliveries: make(map[string][]*Delivery)} + for i, s := range scenario.Steps { + t.Logf("step %d: %s", i, s.Op) + d.step(s) + } + }) + } +} diff --git a/clients/go/ahp/windowing_lifecycle.go b/clients/go/ahp/windowing_lifecycle.go new file mode 100644 index 000000000..1a3fd4af2 --- /dev/null +++ b/clients/go/ahp/windowing_lifecycle.go @@ -0,0 +1,153 @@ +package ahp + +import ( + "encoding/json" + "errors" + "sync" + + "github.com/microsoft/agent-host-protocol/clients/go/ahptypes" +) + +// Hooks belong to one request and one set of captured flow instances. Retiring +// a flow wakes its callers; a late response cannot revive a replacement owner. +type requestHooks struct { + complete func(json.RawMessage) error + failed func(error) + cancel func() + done chan struct{} + once sync.Once + err error + written bool // guarded by flowEngine.mu + flows map[*channelFlow]struct{} +} + +type flowCleanup struct { + channel string + done chan error +} + +func (h *requestHooks) drop(err error) { + h.once.Do(func() { h.err = err; close(h.done) }) +} + +func (e *flowEngine) retireLocked(f *channelFlow, err error, notify bool) { + if e.channels[f.channel] != f { + return + } + written := f.accepted + f.dropped = true + f.epoch++ + if f.creditTimer != nil { + f.creditTimer.Stop() + f.creditTimer = nil + } + f.creditReady = false + f.fragments.Reset() + f.fragmenting = false + e.failQueueLocked(f, err) + for h := range f.requests { + written = written || h.written + delete(h.flows, f) + if len(h.flows) == 0 { + h.drop(err) + } + } + f.requests = nil + if f.owner != nil { + q := f.owner.owner + q.mu.Lock() + q.finishLocked(err) + q.mu.Unlock() + e.emitLocked(f, "dropped") + } + e.c.subscriptionsMu.Lock() + subs := e.c.subscriptions[f.channel] + delete(e.c.subscriptions, f.channel) + if !written && !f.accepted && !errors.Is(err, ErrSubscriptionDropped) { + var closing []*Subscription + for _, sub := range subs { + if sub == f.owner { + closing = append(closing, sub) + } else { + e.c.subscriptions[f.channel] = append(e.c.subscriptions[f.channel], sub) + } + } + subs = closing + } + e.c.subscriptionsMu.Unlock() + for _, s := range subs { + s.closeEvents() + } + delete(e.channels, f.channel) + for i, uri := range e.order { + if uri == f.channel { + copy(e.order[i:], e.order[i+1:]) + e.order[len(e.order)-1] = "" + e.order = e.order[:len(e.order)-1] + if e.cursor > i { + e.cursor-- + } + if e.creditCursor > i { + e.creditCursor-- + } + break + } + } + if e.cursor >= len(e.order) { + e.cursor = 0 + } + if e.creditCursor >= len(e.order) { + e.creditCursor = 0 + } + if notify && written && !e.stopped { + e.cleanup = append(e.cleanup, flowCleanup{channel: f.channel}) + e.signal() + } +} + +// Subscription returns the current owner, including handles created by raw +// subscribe or ReconnectWithParams. Retired subscriptions are not returned. +func (c *Client) Subscription(uri string) (*Subscription, bool) { + c.flow.mu.Lock() + defer c.flow.mu.Unlock() + f := c.flow.channels[uri] + if f == nil || f.dropped || f.owner == nil { + return nil, false + } + return f.owner, true +} + +func (e *flowEngine) beginRequest(h *requestHooks) bool { + e.mu.Lock() + defer e.mu.Unlock() + select { + case <-h.done: + return false + default: + h.written = true + return true + } +} + +func (e *flowEngine) nextCleanup() (*flowWrite, error) { + e.mu.Lock() + defer e.mu.Unlock() + if len(e.cleanup) == 0 || e.stopped { + return nil, nil + } + cleanup := e.cleanup[0] + e.cleanup[0] = flowCleanup{} + e.cleanup = e.cleanup[1:] + if len(e.cleanup) == 0 { + e.cleanup = nil + } + raw, err := json.Marshal(ahptypes.UnsubscribeParams{Channel: cleanup.channel}) + if err != nil { + return nil, err + } + text, err := encodeNotification(ahptypes.JsonRpcNotification{JsonRpc: ahptypes.JsonRpcV2, Method: "unsubscribe", Params: raw}) + if err != nil { + return nil, err + } + return &flowWrite{wire: NewTextMessage(text), controlDone: cleanup.done}, nil +} diff --git a/clients/go/ahp/windowing_lifecycle_test.go b/clients/go/ahp/windowing_lifecycle_test.go new file mode 100644 index 000000000..8633c33e3 --- /dev/null +++ b/clients/go/ahp/windowing_lifecycle_test.go @@ -0,0 +1,360 @@ +package ahp + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "strings" + "testing" + "time" + + "github.com/microsoft/agent-host-protocol/clients/go/ahptypes" +) + +func flowHealthy(t *testing.T, c *Client, wire *corpusTransport, ctx context.Context) { + t.Helper() + done := make(chan error, 1) + go func() { done <- c.Ping(ctx) }() + m := readFlowWire(t, wire) + for m.Notification != nil && m.Notification.Method == "channel/credit" { + m = readFlowWire(t, wire) + } + if m.Request == nil || m.Request.Method != "ping" { + t.Fatalf("unexpected wire message: %s", mustJSON(m)) + } + hostFlowReply(t, wire, m, nil) + select { + case err := <-done: + if err != nil { + t.Fatal(err) + } + case <-ctx.Done(): + t.Fatal(ctx.Err()) + } + if c.Err() != nil { + t.Fatal(c.Err()) + } +} + +func TestWindowingNonObjectParamsPreserveLegacyBehavior(t *testing.T) { + for _, windowed := range []bool{false, true} { + t.Run(fmt.Sprint(windowed), func(t *testing.T) { + c, wire, ctx := flowTestClient(t, DefaultConfig()) + uri := "ahp-terminal:/t" + var s *Subscription + if windowed { + s = subscribeFlow(t, c, wire, ctx, uri, SubscribeOptions{Window: &WindowOptions{Receive: ReceiveManual}}, acceptedFlow(1024, 1024)) + initialFlowSnapshot(wire, uri, 0) + d, err := s.Next(ctx) + if err != nil { + t.Fatal(err) + } + d.Release() + } else { + s = c.AttachSubscription(uri) + } + for _, params := range []string{"", `,"params":[]`, `,"params":null`, `,"params":17`, `,"params":"text"`} { + wire.in <- NewTextMessage(`{"jsonrpc":"2.0","method":"action"` + params + `}`) + } + hostFlowNotification(wire, "action", map[string]any{"channel": uri, "action": map[string]any{"type": "terminal/data", "data": "alive"}, "serverSeq": 1}) + if windowed { + d, err := s.Next(ctx) + if err != nil { + t.Fatal(err) + } + d.Release() + } else { + select { + case <-s.Events(): + case <-ctx.Done(): + t.Fatal(ctx.Err()) + } + } + flowHealthy(t, c, wire, ctx) + }) + } +} + +func TestWindowingLateCanceledSubscribeDoesNotCloseConnection(t *testing.T) { + for _, defaultTimeout := range []bool{false, true} { + t.Run(fmt.Sprint(defaultTimeout), func(t *testing.T) { + cfg := DefaultConfig() + if defaultTimeout { + cfg.DefaultRequestTimeout = 15 * time.Millisecond + } + c, wire, ctx := flowTestClient(t, cfg) + op := context.Background() + cancel := func() {} + if !defaultTimeout { + op, cancel = context.WithTimeout(ctx, 15*time.Millisecond) + } + defer cancel() + done := make(chan error, 1) + go func() { + _, _, err := c.SubscribeWithOptions(op, "ahp-terminal:/late", SubscribeOptions{Window: &WindowOptions{}}) + done <- err + }() + request := readFlowWire(t, wire) + if err := <-done; !errors.Is(err, context.DeadlineExceeded) { + t.Fatal(err) + } + unsubscribe := readFlowWire(t, wire) + if unsubscribe.Notification == nil || unsubscribe.Notification.Method != "unsubscribe" { + t.Fatal("abandoned offer was not unsubscribed") + } + hostFlowReply(t, wire, request, &ahptypes.SubscribeResult{FlowControl: acceptedFlow(512, 1024)}) + initialFlowSnapshot(wire, "ahp-terminal:/late", 0) + hostFlowNotification(wire, "channel/frame", ahptypes.ChannelFrameParams{Channel: "ahp-terminal:/late", Data: "ignored"}) + if _, ok := c.Subscription("ahp-terminal:/late"); ok { + t.Fatal("canceled owner remains reachable") + } + flowHealthy(t, c, wire, ctx) + if err := c.Notify(ctx, "app/data", map[string]any{"channel": "ahp-terminal:/late", "text": "ordinary"}); err != nil { + t.Fatal(err) + } + if m := readFlowWire(t, wire); m.Notification == nil || m.Notification.Method != "app/data" { + t.Fatal("stale flow held a later send") + } + c.flow.mu.Lock() + count := len(c.flow.channels) + len(c.flow.order) + c.flow.mu.Unlock() + if count != 0 { + t.Fatal("canceled flow leaked scheduler entries") + } + }) + } +} + +func TestWindowingFailedSubscribeAndUpdateRollback(t *testing.T) { + for _, existing := range []bool{false, true} { + t.Run(fmt.Sprint(existing), func(t *testing.T) { + c, wire, ctx := flowTestClient(t, DefaultConfig()) + uri := "ahp-terminal:/failure" + var old *Subscription + if existing { + old = subscribeFlow(t, c, wire, ctx, uri, SubscribeOptions{Window: &WindowOptions{}}, acceptedFlow(512, 1024)) + } + done := make(chan error, 1) + go func() { + _, _, err := c.SubscribeWithOptions(ctx, uri, SubscribeOptions{Window: &WindowOptions{WindowBytes: 1000}}) + done <- err + }() + request := readFlowWire(t, wire) + b, _ := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": request.Request.ID, "error": map[string]any{"code": ahptypes.ErrorCodePermissionDenied, "message": "denied"}}) + wire.in <- NewTextMessage(string(b)) + var rpc *RPCError + if err := <-done; !errors.As(err, &rpc) { + t.Fatal("RPC failure was not preserved", err) + } + current, ok := c.Subscription(uri) + if existing && (!ok || current != old || current.FlowControl().Limits.ClientReceive.WindowBytes != 512) { + t.Fatal("failed update changed an accepted owner") + } + if !existing && ok { + t.Fatal("failed new owner remained registered") + } + flowHealthy(t, c, wire, ctx) + }) + } +} + +func TestWindowingCloseWhileSubscribeHookIsClaimed(t *testing.T) { + for _, unsubscribe := range []bool{false, true} { + t.Run(fmt.Sprint(unsubscribe), func(t *testing.T) { + c, wire, ctx := flowTestClient(t, DefaultConfig()) + uri := "ahp-terminal:/race" + done := make(chan error, 1) + go func() { + _, _, err := c.SubscribeWithOptions(ctx, uri, SubscribeOptions{Window: &WindowOptions{}}) + done <- err + }() + request := readFlowWire(t, wire) + claimed, allow := make(chan struct{}), make(chan struct{}) + c.pendingMu.Lock() + h := c.pendingHooks[request.Request.ID] + complete := h.complete + h.complete = func(raw json.RawMessage) error { close(claimed); <-allow; return complete(raw) } + c.pendingMu.Unlock() + hostFlowReply(t, wire, request, &ahptypes.SubscribeResult{FlowControl: acceptedFlow(512, 1024)}) + <-claimed + if unsubscribe { + if err := c.Unsubscribe(ctx, uri); err != nil { + t.Fatal(err) + } + } else { + s, ok := c.Subscription(uri) + if !ok { + t.Fatal("owner missing") + } + s.Close() + } + close(allow) + if err := <-done; !errors.Is(err, ErrClosed) { + t.Fatal("abandoned subscribe did not fail locally", err) + } + m := readFlowWire(t, wire) + if m.Notification == nil || m.Notification.Method != "unsubscribe" { + t.Fatal("missing unsubscribe") + } + initialFlowSnapshot(wire, uri, 0) + flowHealthy(t, c, wire, ctx) + }) + } +} + +func TestWindowingFreshReconnectOwnersAreReachable(t *testing.T) { + c, wire, ctx := flowTestClient(t, DefaultConfig()) + done := make(chan error, 1) + go func() { + _, err := c.ReconnectWithParams(ctx, ahptypes.ReconnectParams{ClientId: "client", Subscriptions: []string{"ahp-terminal:/a", "ahp-terminal:/b"}, Windows: &ahptypes.ChannelDeliveryResumeOptions{Items: []ahptypes.ChannelDeliveryResumeParams{ + {Channel: "ahp-terminal:/a", FlowControl: ahptypes.SubscriptionFlowControlOptions{Receive: ahptypes.ChannelReceiveLimits{WindowBytes: 512, MaximumFrameBytes: 1024}}}, + {Channel: "ahp-terminal:/b", FlowControl: ahptypes.SubscriptionFlowControlOptions{Receive: ahptypes.ChannelReceiveLimits{WindowBytes: 512, MaximumFrameBytes: 1024}}}, + }}}) + done <- err + }() + request := readFlowWire(t, wire) + closed, ok := c.Subscription("ahp-terminal:/a") + if !ok { + t.Fatal("staged reconnect owner inaccessible") + } + closed.Close() + m := readFlowWire(t, wire) + if m.Notification.Method != "unsubscribe" { + t.Fatal(m) + } + hostFlowReply(t, wire, request, map[string]any{"type": "snapshot", "snapshots": []any{}, "windows": map[string]any{"items": []ahptypes.ResumedChannelSubscription{ + {Channel: "ahp-terminal:/a", FlowControl: *acceptedFlow(512, 1024), Recovery: ahptypes.ChannelRecoveryKindSnapshot}, + {Channel: "ahp-terminal:/b", FlowControl: *acceptedFlow(512, 1024), Recovery: ahptypes.ChannelRecoveryKindSnapshot}, + }}}) + initialFlowSnapshot(wire, "ahp-terminal:/a", 0) + initialFlowSnapshot(wire, "ahp-terminal:/b", 0) + if err := <-done; err != nil { + t.Fatal(err) + } + kept, ok := c.Subscription("ahp-terminal:/b") + if !ok { + t.Fatal("resumed owner inaccessible") + } + d, err := kept.Next(ctx) + if err != nil { + t.Fatal(err) + } + d.Release() + m = readFlowWire(t, wire) + if m.Notification.Method != "channel/credit" { + t.Fatal("reachable owner did not return credit") + } + flowHealthy(t, c, wire, ctx) +} + +func TestWindowingControlBypassesFullAndOversizedSendQueues(t *testing.T) { + cfg := DefaultConfig() + cfg.Windowing.MaxBufferedSendBytes = 1 + c, wire, ctx := flowTestClient(t, cfg) + s := subscribeFlow(t, c, wire, ctx, "ahp-tcp:/control", SubscribeOptions{Window: &WindowOptions{Receive: ReceiveManual}}, acceptedFlow(512, 1024)) + initialFlowSnapshot(wire, s.URI(), 0) + d, err := s.Next(ctx) + if err != nil { + t.Fatal(err) + } + done := make(chan error, 1) + go func() { + done <- c.Notify(ctx, "tcp/data", ahptypes.TcpDataParams{Channel: s.URI(), Data: strings.Repeat("eHh4", 1000)}) + }() + m := readFlowWire(t, wire) + if m.Notification.Method != "channel/frame" { + t.Fatal("oversized head did not fragment") + } + waitFlow(t, func() bool { return !s.FlowControl().Writable }) + d.Release() + m = readFlowWire(t, wire) + for m.Notification.Method == "channel/frame" { + m = readFlowWire(t, wire) + } + if m.Notification.Method != "channel/credit" { + t.Fatal("credit was blocked by cap") + } + flowHealthy(t, c, wire, ctx) + if err := c.Notify(ctx, "tcp/reset", ahptypes.TcpResetParams{Channel: s.URI()}); err != nil { + t.Fatal(err) + } + if err := <-done; !errors.Is(err, ErrClosed) { + t.Fatal("reset did not cancel pending send", err) + } + m = readFlowWire(t, wire) + if m.Notification.Method != "tcp/reset" { + t.Fatal(m) + } + if err := c.Unsubscribe(ctx, s.URI()); err != nil { + t.Fatal(err) + } + m = readFlowWire(t, wire) + if m.Notification.Method != "unsubscribe" { + t.Fatal(m) + } + flowHealthy(t, c, wire, ctx) +} + +func TestWindowingPlainSendsThenShutdownPreserveWireAndStreams(t *testing.T) { + for _, method := range []string{"notify", "dispatch", "unsubscribe"} { + t.Run(method, func(t *testing.T) { + c, wire, ctx := flowTestClient(t, DefaultConfig()) + sub := c.AttachSubscription("ahp-terminal:/plain") + events := c.Events() + var err error + switch method { + case "notify": + err = c.Notify(ctx, "app/data", map[string]any{"channel": sub.URI()}) + case "dispatch": + _, err = c.Dispatch(ctx, sub.URI(), ahptypes.StateAction{Value: &ahptypes.TerminalInputAction{Type: ahptypes.ActionTypeTerminalInput, Data: "x"}}) + case "unsubscribe": + err = c.Unsubscribe(ctx, sub.URI()) + } + if err != nil { + t.Fatal(err) + } + if err := c.Shutdown(ctx); err != nil { + t.Fatal(err) + } + if m := readFlowWire(t, wire); m.Notification == nil { + t.Fatal("plain message not on wire before shutdown") + } + select { + case _, ok := <-events.Events(): + if ok { + t.Fatal("global stream still open") + } + default: + t.Fatal("global stream not closed") + } + select { + case _, ok := <-sub.Events(): + if ok { + t.Fatal("subscription stream still open") + } + default: + t.Fatal("subscription stream not closed") + } + }) + } +} + +func TestWindowingSchedulerEntriesArePruned(t *testing.T) { + c, wire, ctx := flowTestClient(t, DefaultConfig()) + for i := 0; i < 50; i++ { + uri := fmt.Sprintf("ahp-terminal:/prune-%d", i) + s := subscribeFlow(t, c, wire, ctx, uri, SubscribeOptions{Window: &WindowOptions{}}, acceptedFlow(512, 1024)) + s.Close() + if m := readFlowWire(t, wire); m.Notification.Method != "unsubscribe" { + t.Fatal(m) + } + } + c.flow.mu.Lock() + n := len(c.flow.order) + len(c.flow.channels) + len(c.flow.cleanup) + c.flow.mu.Unlock() + if n != 0 { + t.Fatal("bookkeeping grew with closed channels", n) + } + flowHealthy(t, c, wire, ctx) +} diff --git a/clients/go/ahp/windowing_reconnect.go b/clients/go/ahp/windowing_reconnect.go new file mode 100644 index 000000000..337eef328 --- /dev/null +++ b/clients/go/ahp/windowing_reconnect.go @@ -0,0 +1,154 @@ +package ahp + +import ( + "encoding/json" + "fmt" + "strings" + + "github.com/microsoft/agent-host-protocol/clients/go/ahptypes" +) + +func (e *flowEngine) prepareReconnect(raw json.RawMessage) (json.RawMessage, *requestHooks, error) { + var p ahptypes.ReconnectParams + if err := json.Unmarshal(raw, &p); err != nil { + return raw, nil, err + } + subs := make(map[string]bool) + for _, uri := range p.Subscriptions { + if strings.HasPrefix(uri, "ahp-tcp:") { + return raw, nil, fmt.Errorf("ahp: TCP cannot resume") + } + subs[uri] = true + } + if p.Windows == nil { + return raw, nil, nil + } + requested := make(map[string]ahptypes.ChannelDeliveryResumeParams) + for _, item := range p.Windows.Items { + if _, exists := requested[item.Channel]; exists || !subs[item.Channel] || strings.HasPrefix(item.Channel, "ahp-tcp:") { + return raw, nil, fmt.Errorf("ahp: invalid reconnect window channel") + } + if err := validLimits(item.FlowControl.Receive); err != nil { + return raw, nil, err + } + if item.LastAppliedServerSeq != nil && (*item.LastAppliedServerSeq < 0 || *item.LastAppliedServerSeq > maxWindowInteger) { + return raw, nil, fmt.Errorf("ahp: invalid resume sequence") + } + requested[item.Channel] = item + } + if !e.capabilities.SupportsWindowedDelivery { + return raw, nil, fmt.Errorf("ahp: reconnect windows require raw transport") + } + e.mu.Lock() + defer e.mu.Unlock() + if e.stopped { + return raw, nil, ErrShutdown + } + h := &requestHooks{done: make(chan struct{}), flows: make(map[*channelFlow]struct{})} + flows := make(map[string]*channelFlow) + owners := make(map[string]*Subscription) + for uri := range requested { + f := e.getLocked(uri) + if f.owner == nil { + e.attachLocked(f, nil) + } + flows[uri], owners[uri] = f, f.owner + if f.requests == nil { + f.requests = make(map[*requestHooks]struct{}) + } + f.requests[h] = struct{}{} + h.flows[f] = struct{}{} + } + h.cancel = func() { + e.mu.Lock() + defer e.mu.Unlock() + for _, f := range flows { + e.retireLocked(f, ErrClosed, true) + } + } + h.failed = func(err error) { + e.mu.Lock() + defer e.mu.Unlock() + for _, f := range flows { + delete(f.requests, h) + if !f.accepted { + e.retireLocked(f, err, true) + } + } + } + h.complete = func(result json.RawMessage) error { + e.mu.Lock() + defer e.mu.Unlock() + if e.stopped { + return nil + } + active := len(flows) == 0 + for uri, f := range flows { + if e.channels[uri] == f && f.owner == owners[uri] && !owners[uri].owner.isFinal() { + active = true + break + } + } + if !active { + return nil + } + var r struct { + Windows *ahptypes.ChannelDeliveryResumeResult `json:"windows"` + } + if err := json.Unmarshal(result, &r); err != nil { + return err + } + kept := make(map[string]bool) + if r.Windows != nil { + for _, item := range r.Windows.Items { + offer, ok := requested[item.Channel] + if !ok || kept[item.Channel] { + return fmt.Errorf("ahp: unexpected resumed channel") + } + kept[item.Channel] = true + f, owner := flows[item.Channel], owners[item.Channel] + if e.channels[item.Channel] != f || f.owner != owner || owner.owner.isFinal() { + continue + } + if err := validateAccepted(&item.FlowControl, &offer.FlowControl); err != nil { + return err + } + if item.Recovery != ahptypes.ChannelRecoveryKindSnapshot && item.Recovery != ahptypes.ChannelRecoveryKindReplay && item.Recovery != ahptypes.ChannelRecoveryKindLive { + return fmt.Errorf("ahp: unknown recovery kind") + } + delete(f.requests, h) + e.failQueueLocked(f, ErrShutdown) + f.epoch++ + if f.creditTimer != nil { + f.creditTimer.Stop() + f.creditTimer = nil + } + f.creditReady = false + f.sent, f.consumed, f.received, f.released, f.advertised = 0, 0, 0, 0, 0 + f.fragments.Reset() + f.fragmenting = false + f.fragmentBytes = 0 + l := item.FlowControl + f.limits = &l + f.accepted = true + if f.opts.Window == nil { + f.opts.Window = &WindowOptions{} + } + f.opts.Window.WindowBytes = offer.FlowControl.Receive.WindowBytes + f.opts.Window.MaximumFrameBytes = offer.FlowControl.Receive.MaximumFrameBytes + f.expectSnapshot = item.Recovery == ahptypes.ChannelRecoveryKindSnapshot + owner.owner.setMode(f.opts) + e.emitLocked(f, "resumed") + } + } + for uri, f := range flows { + if !kept[uri] { + delete(f.requests, h) + e.retireLocked(f, ErrSubscriptionDropped, false) + } + } + e.signal() + return nil + } + return raw, h, nil +} diff --git a/clients/go/ahp/windowing_scheduler.go b/clients/go/ahp/windowing_scheduler.go new file mode 100644 index 000000000..ca73d0540 --- /dev/null +++ b/clients/go/ahp/windowing_scheduler.go @@ -0,0 +1,548 @@ +package ahp + +import ( + "context" + "encoding/json" + "fmt" + "time" + "unicode/utf8" + + "github.com/microsoft/agent-host-protocol/clients/go/ahptypes" +) + +type sendPart struct { + n ahptypes.JsonRpcNotification + ctx context.Context + done chan error + bytes int64 +} +type flowSend struct { + n ahptypes.JsonRpcNotification + text string + offset int + remaining int64 + parts []*sendPart + started bool + retainedBytes int64 +} +type flowWrite struct { + wire TransportMessage + message *flowSend + final bool + creditFlow *channelFlow + epoch uint64 + controlDone chan error +} + +func encodeNotification(n ahptypes.JsonRpcNotification) (string, error) { + b, err := json.Marshal(n) + return string(b), err +} + +func (e *flowEngine) send(ctx context.Context, msg ahptypes.JsonRpcMessage) (bool, error) { + if msg.Notification == nil { + return false, nil + } + n := *msg.Notification + channel, err := notificationChannel(n) + if err != nil { + return false, nil + } + e.mu.Lock() + f := e.channels[channel] + if f != nil && n.Method == "unsubscribe" { + e.retireLocked(f, ErrClosed, false) + e.mu.Unlock() + return false, nil + } + if f != nil && f.dropped && ahptypes.CountsAgainstWindow(n.Method) { + e.mu.Unlock() + return true, ErrSubscriptionDropped + } + if f != nil && f.reset && ahptypes.CountsAgainstWindow(n.Method) { + e.mu.Unlock() + return true, ErrClosed + } + if f == nil || (f.limits == nil && len(f.queue) == 0) || !ahptypes.CountsAgainstWindow(n.Method) { + if f != nil && (n.Method == "tcp/reset" || n.Method == "unsubscribe") { + e.failQueueLocked(f, ErrClosed) + if n.Method == "tcp/reset" { + f.reset = true + e.emitLocked(f, "") + } + } + e.mu.Unlock() + return false, nil + } + if e.stopped { + e.mu.Unlock() + return true, ErrShutdown + } + text, err := encodeNotification(n) + if err != nil { + e.mu.Unlock() + return true, err + } + size := utf16Bytes(text) + queued := len(f.queue) > 0 + if queued && (size > e.cfg.MaxBufferedSendBytes || f.retainedSendBytes > e.cfg.MaxBufferedSendBytes-size) { + e.mu.Unlock() + return true, ErrBackpressure + } + part := &sendPart{n: n, ctx: ctx, done: make(chan error, 1), bytes: size} + var merged *flowSend + if coalescingEnabled(f.opts) && len(f.queue) > 0 { + last := f.queue[len(f.queue)-1] + if !last.started { + if combined, ok := e.coalesce(last.n, n, CoalesceSend, f); ok { + combinedText, err := encodeNotification(combined) + if err != nil { + e.mu.Unlock() + return true, err + } + combinedSize := utf16Bytes(combinedText) + if f.buffered-last.remaining+combinedSize <= e.cfg.MaxBufferedSendBytes { + f.buffered += combinedSize - last.remaining + last.n, last.text, last.remaining = combined, combinedText, combinedSize + last.parts = append(last.parts, part) + last.retainedBytes += part.bytes + f.retainedSendBytes += part.bytes + merged = last + e.emitLocked(f, "coalesced") + } + } + } + } + if merged == nil { + if queued && (size > e.cfg.MaxBufferedSendBytes || f.buffered > e.cfg.MaxBufferedSendBytes-size) { + e.mu.Unlock() + return true, ErrBackpressure + } + f.queue = append(f.queue, &flowSend{n: n, text: text, remaining: size, parts: []*sendPart{part}, retainedBytes: part.bytes}) + f.buffered += size + f.retainedSendBytes += part.bytes + } + e.emitLocked(f, "") + e.mu.Unlock() + e.signal() + select { + case err := <-part.done: + return true, err + case <-ctx.Done(): + e.signal() + return true, ctx.Err() + case <-e.c.done: + return true, ErrShutdown + } +} + +func (e *flowEngine) failQueueLocked(f *channelFlow, err error) { + for _, m := range f.queue { + for _, p := range m.parts { + select { + case p.done <- err: + default: + } + } + } + f.queue = nil + f.buffered = 0 + f.retainedSendBytes = 0 +} + +func (e *flowEngine) pruneLocked(f *channelFlow) { + for len(f.queue) > 0 { + m := f.queue[0] + if m.started { + return + } + live := make([]*sendPart, 0, len(m.parts)) + for _, p := range m.parts { + if err := p.ctx.Err(); err != nil { + select { + case p.done <- err: + default: + } + } else { + live = append(live, p) + } + } + if len(live) == len(m.parts) { + return + } + f.retainedSendBytes -= m.retainedBytes + f.buffered -= m.remaining + if len(live) == 0 { + f.queue = f.queue[1:] + e.emitLocked(f, "") + continue + } + n := live[0].n + for _, p := range live[1:] { + combined, ok := e.coalesce(n, p.n, CoalesceSend, f) + if ok { + n = combined + } else { + // A canceled middle item can invalidate a custom merge. Restore + // separate sends rather than drop or alter the survivors. + replacement := make([]*flowSend, 0, len(live)) + for _, part := range live { + text, _ := encodeNotification(part.n) + size := utf16Bytes(text) + replacement = append(replacement, &flowSend{n: part.n, text: text, remaining: size, parts: []*sendPart{part}, retainedBytes: part.bytes}) + f.buffered += size + f.retainedSendBytes += part.bytes + } + f.queue = append(replacement, f.queue[1:]...) + e.emitLocked(f, "") + return + } + } + text, _ := encodeNotification(n) + m.n, m.text, m.parts, m.remaining = n, text, live, utf16Bytes(text) + m.retainedBytes = 0 + for _, part := range live { + m.retainedBytes += part.bytes + } + f.retainedSendBytes += m.retainedBytes + f.buffered += m.remaining + e.emitLocked(f, "") + return + } +} + +func frameText(channel, data string, final bool) (string, error) { + p := ahptypes.ChannelFrameParams{Channel: channel, Data: data} + if final { + p.Final = &final + } + raw, err := json.Marshal(p) + if err != nil { + return "", err + } + return encodeNotification(ahptypes.JsonRpcNotification{JsonRpc: ahptypes.JsonRpcV2, Method: "channel/frame", Params: raw}) +} + +func fragmentText(channel, text string, maximum, remaining int64) (string, int, error) { + if text == "" { + return "", 0, fmt.Errorf("ahp: empty fragment") + } + _, first := utf8.DecodeRuneInString(text) + smallest, err := frameText(channel, text[:first], first == len(text)) + if err != nil { + return "", 0, err + } + minimum := utf16Bytes(smallest) + if minimum > maximum { + return "", 0, fmt.Errorf("ahp: channel/frame cannot carry a character within maximumFrameBytes") + } + target := remaining + if target < minimum { + target = minimum + } + if target > maximum { + target = maximum + } + boundaries := []int{0} + var prefixBytes int64 + for i, r := range text { + if r > 0xffff { + prefixBytes += 4 + } else { + prefixBytes += 2 + } + if prefixBytes > target { + break + } + boundaries = append(boundaries, i+utf8.RuneLen(r)) + } + low, high := 1, len(boundaries)-1 + best := 1 + for low <= high { + mid := low + (high-low)/2 + frame, err := frameText(channel, text[:boundaries[mid]], boundaries[mid] == len(text)) + if err != nil { + return "", 0, err + } + if utf16Bytes(frame) <= target { + best = mid + low = mid + 1 + } else { + high = mid - 1 + } + } + frame, err := frameText(channel, text[:boundaries[best]], boundaries[best] == len(text)) + return frame, boundaries[best], err +} + +func (e *flowEngine) nextCredit() (*flowWrite, error) { + e.mu.Lock() + defer e.mu.Unlock() + for i := 0; i < len(e.order); i++ { + index := (e.creditCursor + i) % len(e.order) + uri := e.order[index] + f := e.channels[uri] + if f == nil || !f.creditReady || f.limits == nil || f.released <= f.advertised { + continue + } + if f.owner != nil && f.owner.owner.isFinal() { + continue + } + raw, err := json.Marshal(ahptypes.ChannelCreditParams{Channel: uri, ConsumedBytes: f.released}) + if err != nil { + return nil, err + } + text, err := encodeNotification(ahptypes.JsonRpcNotification{JsonRpc: ahptypes.JsonRpcV2, Method: "channel/credit", Params: raw}) + if err != nil { + return nil, err + } + f.creditReady = false + f.advertised = f.released + if f.creditTimer != nil { + f.creditTimer.Stop() + f.creditTimer = nil + } + e.emitLocked(f, "creditSent") + e.creditCursor = (index + 1) % len(e.order) + return &flowWrite{wire: NewTextMessage(text), creditFlow: f, epoch: f.epoch}, nil + } + return nil, nil +} + +func (e *flowEngine) nextData() (*flowWrite, error) { + e.mu.Lock() + defer e.mu.Unlock() + for i := 0; i < len(e.order); i++ { + index := (e.cursor + i) % len(e.order) + f := e.channels[e.order[index]] + if f == nil { + continue + } + e.pruneLocked(f) + if len(f.queue) == 0 { + continue + } + if f.limits != nil && f.sent-f.consumed >= f.limits.HostReceive.WindowBytes { + continue + } + m := f.queue[0] + text := m.text[m.offset:] + count := len(text) + final := true + if f.limits != nil && (m.offset > 0 || m.remaining > f.limits.HostReceive.MaximumFrameBytes) { + var err error + text, count, err = fragmentText(f.channel, text, f.limits.HostReceive.MaximumFrameBytes, f.limits.HostReceive.WindowBytes-(f.sent-f.consumed)) + if err != nil { + return nil, err + } + final = m.offset+count == len(m.text) + } + cost := utf16Bytes(text) + if f.sent > maxWindowInteger-cost { + return nil, fmt.Errorf("ahp: sent counter overflow") + } + m.started = true + if f.limits != nil { + f.sent += cost + } + consumed := utf16Bytes(m.text[m.offset : m.offset+count]) + m.offset += count + m.remaining -= consumed + f.buffered -= consumed + if final { + f.queue = f.queue[1:] + f.retainedSendBytes -= m.retainedBytes + } + e.cursor = (index + 1) % len(e.order) + kind := "" + if final && m.text != text { + kind = "fragmentedMessageSent" + } + e.emitLocked(f, kind) + return &flowWrite{wire: NewTextMessage(text), message: m, final: final}, nil + } + return nil, nil +} + +func (e *flowEngine) write(w *flowWrite) bool { + if w.creditFlow != nil { + e.mu.Lock() + live := e.channels[w.creditFlow.channel] == w.creditFlow && w.creditFlow.epoch == w.epoch && !w.creditFlow.dropped && (w.creditFlow.owner == nil || !w.creditFlow.owner.owner.isFinal()) + e.mu.Unlock() + if !live { + return true + } + } + ctx, cancel := contextWithDone(context.Background(), e.c.done) + if w.message == nil && e.waitControl(ctx) != nil { + cancel() + return false + } + err := e.c.transport.Send(ctx, w.wire) + cancel() + if w.controlDone != nil { + w.controlDone <- err + } + if w.message != nil && (w.final || err != nil) { + for _, p := range w.message.parts { + select { + case p.done <- err: + default: + } + } + } + if err != nil { + e.c.shutdownWithError(fmt.Errorf("ahp: transport send: %w", err)) + return false + } + return true +} + +func (e *flowEngine) writeOrdinary(msg outboundMsg) bool { + for { + cleanup, err := e.nextCleanup() + if err != nil { + e.c.shutdownWithError(err) + return false + } + if cleanup == nil { + break + } + if !e.write(cleanup) { + return false + } + } + if msg.hooks != nil && !e.beginRequest(msg.hooks) { + if msg.done != nil { + msg.done <- msg.hooks.err + } + return true + } + wire, err := EncodeMessage(msg.msg) + if err != nil { + if msg.done != nil { + msg.done <- err + } + return true + } + control := msg.msg.Notification != nil && !ahptypes.CountsAgainstWindow(msg.msg.Notification.Method) && e.hasWindows() + if control && utf16Bytes(wire.text) > 1<<20 { + if msg.done != nil { + msg.done <- ErrBackpressure + } + return true + } + ctx, cancel := contextWithDone(context.Background(), e.c.done) + if control { + if err := e.waitControl(ctx); err != nil { + cancel() + if msg.done != nil { + msg.done <- err + } + return true + } + } + err = e.c.transport.Send(ctx, wire) + cancel() + if msg.done != nil { + msg.done <- err + } + + if err != nil { + e.c.shutdownWithError(fmt.Errorf("ahp: transport send: %w", err)) + return false + } + return true +} + +// Only the writer touches this token bucket. Control traffic has a bounded +// burst and rate independently of data windows; credits stay cumulative. +func (e *flowEngine) waitControl(ctx context.Context) error { + for { + now := time.Now() + e.controlTokens += now.Sub(e.controlAt).Seconds() * 1000 + e.controlAt = now + if e.controlTokens > 32 { + e.controlTokens = 32 + } + if e.controlTokens >= 1 { + e.controlTokens-- + return nil + } + timer := time.NewTimer(time.Millisecond) + select { + case <-ctx.Done(): + timer.Stop() + return ctx.Err() + case <-timer.C: + } + } +} + +func (e *flowEngine) runWriter() { + defer e.c.wg.Done() + controlBurst := 0 + for { + select { + case <-e.c.done: + return + default: + } + if controlBurst < 8 { + cleanup, err := e.nextCleanup() + if err != nil { + e.c.shutdownWithError(err) + return + } + if cleanup != nil { + if !e.write(cleanup) { + return + } + controlBurst++ + continue + } + w, err := e.nextCredit() + if err != nil { + e.c.shutdownWithError(&TransportError{Kind: "protocol", Err: err}) + return + } + if w != nil { + if !e.write(w) { + return + } + controlBurst++ + continue + } + select { + case msg := <-e.c.outbound: + if !e.writeOrdinary(msg) { + return + } + controlBurst++ + continue + default: + } + } + w, err := e.nextData() + if err != nil { + e.c.shutdownWithError(&TransportError{Kind: "protocol", Err: err}) + return + } + controlBurst = 0 + if w != nil { + if !e.write(w) { + return + } + continue + } + select { + case <-e.c.done: + return + case msg := <-e.c.outbound: + if !e.writeOrdinary(msg) { + return + } + case <-e.wake: + } + } +} diff --git a/clients/go/ahp/windowing_subscription_lifecycle_test.go b/clients/go/ahp/windowing_subscription_lifecycle_test.go new file mode 100644 index 000000000..4aa7a61c4 --- /dev/null +++ b/clients/go/ahp/windowing_subscription_lifecycle_test.go @@ -0,0 +1,129 @@ +package ahp + +import ( + "context" + "encoding/json" + "errors" + "testing" + "time" + + "github.com/microsoft/agent-host-protocol/clients/go/ahptypes" +) + +func TestWindowingLateReconnectTimeoutDoesNotReviveOwners(t *testing.T) { + c, wire, ctx := flowTestClient(t, DefaultConfig()) + op, cancel := context.WithTimeout(ctx, 15*time.Millisecond) + defer cancel() + done := make(chan error, 1) + uri := "ahp-terminal:/late-reconnect" + go func() { + _, err := c.ReconnectWithParams(op, ahptypes.ReconnectParams{ClientId: "client", Subscriptions: []string{uri}, Windows: &ahptypes.ChannelDeliveryResumeOptions{Items: []ahptypes.ChannelDeliveryResumeParams{ + {Channel: uri, FlowControl: ahptypes.SubscriptionFlowControlOptions{Receive: ahptypes.ChannelReceiveLimits{WindowBytes: 512, MaximumFrameBytes: 1024}}}, + }}}) + done <- err + }() + request := readFlowWire(t, wire) + if err := <-done; !errors.Is(err, context.DeadlineExceeded) { + t.Fatal(err) + } + if m := readFlowWire(t, wire); m.Notification == nil || m.Notification.Method != "unsubscribe" { + t.Fatal("timed out reconnect was not unsubscribed") + } + hostFlowReply(t, wire, request, map[string]any{"type": "snapshot", "snapshots": []any{}, "windows": map[string]any{"items": []ahptypes.ResumedChannelSubscription{ + {Channel: uri, FlowControl: *acceptedFlow(512, 1024), Recovery: ahptypes.ChannelRecoveryKindSnapshot}, + }}}) + initialFlowSnapshot(wire, uri, 0) + if _, ok := c.Subscription(uri); ok { + t.Fatal("timed out reconnect left a hidden owner") + } + flowHealthy(t, c, wire, ctx) +} + +func TestWindowingFailedOptInRetainsLegacyObserver(t *testing.T) { + c, wire, ctx := flowTestClient(t, DefaultConfig()) + uri := "ahp-terminal:/legacy" + legacy := c.AttachSubscription(uri) + done := make(chan error, 1) + go func() { + _, _, err := c.SubscribeWithOptions(ctx, uri, SubscribeOptions{Window: &WindowOptions{}}) + done <- err + }() + request := readFlowWire(t, wire) + b, _ := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": request.Request.ID, "error": map[string]any{"code": ahptypes.ErrorCodePermissionDenied, "message": "denied"}}) + wire.in <- NewTextMessage(string(b)) + var rpc *RPCError + if err := <-done; !errors.As(err, &rpc) { + t.Fatal(err) + } + hostFlowNotification(wire, "action", map[string]any{"channel": uri, "action": map[string]any{"type": "terminal/data", "data": "still observed"}, "serverSeq": 1}) + select { + case _, ok := <-legacy.Events(): + if !ok { + t.Fatal("failed opt-in closed an existing observer") + } + case <-ctx.Done(): + t.Fatal(ctx.Err()) + } + flowHealthy(t, c, wire, ctx) +} + +func TestWindowingClosedOwnerUpdateDoesNotCreateHiddenOwner(t *testing.T) { + c, wire, ctx := flowTestClient(t, DefaultConfig()) + sub := subscribeFlow(t, c, wire, ctx, "ahp-terminal:/closed", SubscribeOptions{Window: &WindowOptions{}}, acceptedFlow(512, 1024)) + changes := sub.FlowControlChanges() + sub.Close() + if m := readFlowWire(t, wire); m.Notification.Method != "unsubscribe" { + t.Fatal(m) + } + if _, err := sub.Update(ctx, SubscribeOptions{Window: &WindowOptions{}}); !errors.Is(err, ErrClosed) { + t.Fatal("closed owner was updated", err) + } + if _, ok := c.Subscription(sub.URI()); ok { + t.Fatal("Update created a replacement owner") + } + for range changes { + } + flowHealthy(t, c, wire, ctx) +} + +func TestWindowingUnsubscribeStopsPendingCreditTimer(t *testing.T) { + c, wire, ctx := flowTestClient(t, DefaultConfig()) + uri := "ahp-terminal:/timer" + sub := subscribeFlow(t, c, wire, ctx, uri, SubscribeOptions{Window: &WindowOptions{Receive: ReceiveManual}}, acceptedFlow(4096, 1024)) + initialFlowSnapshot(wire, uri, 0) + old, err := sub.Next(ctx) + if err != nil { + t.Fatal(err) + } + old.Release() + if err := c.Unsubscribe(ctx, uri); err != nil { + t.Fatal(err) + } + for { + m := readFlowWire(t, wire) + if m.Notification.Method == "unsubscribe" { + break + } + if m.Notification.Method != "channel/credit" { + t.Fatal(m) + } + } + select { + case text := <-wire.sent: + t.Fatal("credit timer survived unsubscribe", text) + case <-time.After(15 * time.Millisecond): + } + fresh := subscribeFlow(t, c, wire, ctx, uri, SubscribeOptions{Window: &WindowOptions{Receive: ReceiveManual}}, acceptedFlow(4096, 1024)) + initialFlowSnapshot(wire, uri, 0) + d, err := fresh.Next(ctx) + if err != nil { + t.Fatal(err) + } + d.Release() + m := readFlowWire(t, wire) + var credit ahptypes.ChannelCreditParams + _ = json.Unmarshal(m.Notification.Params, &credit) + if m.Notification.Method != "channel/credit" || credit.ConsumedBytes != d.WireBytes() { + t.Fatal("new counters were not reset", m) + } +} diff --git a/clients/go/ahp/windowing_test.go b/clients/go/ahp/windowing_test.go new file mode 100644 index 000000000..8a40d0fef --- /dev/null +++ b/clients/go/ahp/windowing_test.go @@ -0,0 +1,517 @@ +package ahp + +import ( + "context" + "encoding/json" + "errors" + "strings" + "testing" + "time" + "unicode/utf8" + + "github.com/microsoft/agent-host-protocol/clients/go/ahptypes" +) + +func flowTestClient(t *testing.T, cfg Config) (*Client, *corpusTransport, context.Context) { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + t.Cleanup(cancel) + wire := newCorpusTransport() + c, err := Connect(ctx, wire, cfg) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = c.Shutdown(context.Background()) }) + return c, wire, ctx +} +func readFlowWire(t *testing.T, wire *corpusTransport) ahptypes.JsonRpcMessage { + t.Helper() + select { + case text := <-wire.sent: + var m ahptypes.JsonRpcMessage + if err := json.Unmarshal([]byte(text), &m); err != nil { + t.Fatal(err) + } + return m + case <-time.After(time.Second): + t.Fatal("no outgoing message") + return ahptypes.JsonRpcMessage{} + } +} +func hostFlowReply(t *testing.T, wire *corpusTransport, request ahptypes.JsonRpcMessage, result any) { + t.Helper() + if request.Request == nil { + t.Fatal("expected request") + } + b, _ := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": request.Request.ID, "result": result}) + wire.in <- NewTextMessage(string(b)) +} +func hostFlowNotification(wire *corpusTransport, method string, params any) { + b, _ := json.Marshal(map[string]any{"jsonrpc": "2.0", "method": method, "params": params}) + wire.in <- NewTextMessage(string(b)) +} +func acceptedFlow(window, frame int64) *ahptypes.ChannelFlowControl { + return &ahptypes.ChannelFlowControl{ + ClientReceive: ahptypes.ChannelReceiveLimits{WindowBytes: window, MaximumFrameBytes: frame}, + HostReceive: ahptypes.ChannelReceiveLimits{WindowBytes: window, MaximumFrameBytes: frame}, + } +} +func subscribeFlow(t *testing.T, c *Client, wire *corpusTransport, ctx context.Context, uri string, opts SubscribeOptions, limits *ahptypes.ChannelFlowControl) *Subscription { + t.Helper() + type result struct { + sub *Subscription + err error + } + done := make(chan result, 1) + go func() { _, s, err := c.SubscribeWithOptions(ctx, uri, opts); done <- result{s, err} }() + request := readFlowWire(t, wire) + hostFlowReply(t, wire, request, &ahptypes.SubscribeResult{FlowControl: limits}) + select { + case r := <-done: + if r.err != nil { + t.Fatal(r.err) + } + return r.sub + case <-ctx.Done(): + t.Fatal(ctx.Err()) + return nil + } +} +func initialFlowSnapshot(wire *corpusTransport, uri string, seq int64) { + hostFlowNotification(wire, "channel/snapshot", map[string]any{"channel": uri, "snapshot": map[string]any{"resource": uri, "state": map[string]any{}, "fromSeq": seq}}) +} +func waitFlow(t *testing.T, predicate func() bool) { + t.Helper() + deadline := time.Now().Add(time.Second) + for !predicate() { + if time.Now().After(deadline) { + t.Fatal("flow condition did not become true") + } + time.Sleep(time.Millisecond) + } +} + +func TestFlowTransportFallbackAndSnapshot(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + wire := newCorpusTransport() + wire.capable = false + c, err := Connect(ctx, wire, DefaultConfig()) + if err != nil { + t.Fatal(err) + } + defer c.Shutdown(context.Background()) + done := make(chan *Subscription, 1) + go func() { + _, s, _ := c.SubscribeWithOptions(ctx, "ahp-terminal:/t", SubscribeOptions{Window: &WindowOptions{Receive: ReceiveManual}}) + done <- s + }() + request := readFlowWire(t, wire) + var p ahptypes.SubscribeParams + _ = json.Unmarshal(request.Request.Params, &p) + if p.FlowControl != nil { + t.Fatal("parsed transport offered windowing") + } + hostFlowReply(t, wire, request, map[string]any{"snapshot": map[string]any{"resource": p.Channel, "state": map[string]any{}, "fromSeq": 3}}) + sub := <-done + if sub.FlowControl().FallbackReason != "transport" { + t.Fatal(sub.FlowControl()) + } + d, err := sub.Next(ctx) + if err != nil { + t.Fatal(err) + } + if _, ok := d.Event().(SubscriptionEventSnapshot); !ok || d.WireBytes() != 0 { + t.Fatal("fallback did not inject an uncharged snapshot") + } + d.Release() + select { + case text := <-wire.sent: + t.Fatalf("ordinary fallback sent %s", text) + case <-time.After(15 * time.Millisecond): + } +} + +func TestFlowRawRequestHooksAndBinaryAccounting(t *testing.T) { + c, wire, ctx := flowTestClient(t, DefaultConfig()) + uri := "ahp-terminal:/generic" + done := make(chan error, 1) + go func() { + var out ahptypes.SubscribeResult + done <- c.Request(ctx, "subscribe", ahptypes.SubscribeParams{Channel: uri, FlowControl: &ahptypes.SubscriptionFlowControlOptions{Receive: ahptypes.ChannelReceiveLimits{WindowBytes: 1024, MaximumFrameBytes: 1024}}}, &out) + }() + request := readFlowWire(t, wire) + hostFlowReply(t, wire, request, &ahptypes.SubscribeResult{FlowControl: acceptedFlow(1024, 1024)}) + raw := `{ "jsonrpc":"2.0", "method":"channel/snapshot", "params":{"channel":"ahp-terminal:/generic","snapshot":{"resource":"ahp-terminal:/generic","state":{},"fromSeq":0}} }` + wire.in <- NewBinaryMessage([]byte(raw)) + if err := <-done; err != nil { + t.Fatal(err) + } + c.flow.mu.Lock() + sub := c.flow.channels[uri].owner + c.flow.mu.Unlock() + d, err := sub.Next(ctx) + if err != nil { + t.Fatal(err) + } + if d.WireBytes() != utf16Bytes(raw) { + t.Fatal("binary input was reserialized") + } + d.Release() + m := readFlowWire(t, wire) + if m.Notification.Method != "channel/credit" { + t.Fatal(m) + } +} + +func TestFlowCustomCoalescerAndValidation(t *testing.T) { + cfg := DefaultConfig() + cfg.Windowing.Coalescers = []Coalescer{{Method: "app/chunk", Merge: func(a, b ahptypes.JsonRpcNotification, ctx CoalescingContext) (ahptypes.JsonRpcNotification, bool) { + if ctx.Direction != CoalesceReceive { + return a, false + } + var p, q map[string]any + _ = json.Unmarshal(a.Params, &p) + _ = json.Unmarshal(b.Params, &q) + p["text"] = p["text"].(string) + q["text"].(string) + a.Params, _ = json.Marshal(p) + return a, true + }}} + c, wire, ctx := flowTestClient(t, cfg) + uri := "ahp-terminal:/custom" + sub := subscribeFlow(t, c, wire, ctx, uri, SubscribeOptions{Window: &WindowOptions{Receive: ReceiveManual}}, acceptedFlow(1024, 1024)) + initialFlowSnapshot(wire, uri, 0) + snapshot, err := sub.Next(ctx) + if err != nil { + t.Fatal(err) + } + snapshot.Release() + _ = readFlowWire(t, wire) + hostFlowNotification(wire, "app/chunk", map[string]any{"channel": uri, "text": "a"}) + hostFlowNotification(wire, "app/chunk", map[string]any{"channel": uri, "text": "b"}) + waitFlow(t, func() bool { + sub.owner.mu.Lock() + defer sub.owner.mu.Unlock() + return len(sub.owner.items) == 1 && strings.Contains(string(sub.owner.items[0].notification.Params), `"ab"`) + }) + d, err := sub.Next(ctx) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(string(d.Notification().Params), `"ab"`) { + t.Fatal(d.Notification()) + } + d.Release() + d.Release() + waitFlow(t, func() bool { return sub.FlowControl().BufferedReceiveBytes == 0 }) + c.flow.mu.Lock() + c.flow.cfg.Coalescers[0].Merge = func(a, b ahptypes.JsonRpcNotification, _ CoalescingContext) (ahptypes.JsonRpcNotification, bool) { + a.Params = json.RawMessage(`{"channel":"ahp-terminal:/wrong"}`) + return a, true + } + c.flow.mu.Unlock() + hostFlowNotification(wire, "app/chunk", map[string]any{"channel": uri, "text": "a"}) + hostFlowNotification(wire, "app/chunk", map[string]any{"channel": uri, "text": "b"}) + waitFlow(t, func() bool { sub.owner.mu.Lock(); defer sub.owner.mu.Unlock(); return len(sub.owner.items) == 2 }) +} + +func TestFlowFairnessControlAndCancellation(t *testing.T) { + c, wire, ctx := flowTestClient(t, DefaultConfig()) + a := subscribeFlow(t, c, wire, ctx, "ahp-terminal:/a", SubscribeOptions{Window: &WindowOptions{}}, acceptedFlow(512, 1024)) + b := subscribeFlow(t, c, wire, ctx, "ahp-terminal:/b", SubscribeOptions{Window: &WindowOptions{}}, acceptedFlow(512, 1024)) + aCtx, cancel := context.WithCancel(ctx) + firstDone := make(chan error, 1) + go func() { + firstDone <- c.Notify(aCtx, "app/data", map[string]any{"channel": a.URI(), "data": strings.Repeat("漢😀\\\"", 500)}) + }() + first := readFlowWire(t, wire) + if first.Notification.Method != "channel/frame" { + t.Fatal("large message not fragmented") + } + cancel() + if !errors.Is(<-firstDone, context.Canceled) { + t.Fatal("send cancellation not reported") + } + waitFlow(t, func() bool { return !a.FlowControl().Writable }) + blockCtx, blockCancel := context.WithTimeout(ctx, 20*time.Millisecond) + defer blockCancel() + if !errors.Is(a.WaitWritable(blockCtx), context.DeadlineExceeded) { + t.Fatal("full window writable") + } + if err := c.Notify(ctx, "app/data", map[string]any{"channel": b.URI(), "data": "other"}); err != nil { + t.Fatal(err) + } + for i := 0; ; i++ { + other := readFlowWire(t, wire) + channel, _ := notificationChannel(*other.Notification) + if channel == b.URI() { + break + } + if i >= 2 || channel != a.URI() { + t.Fatal("another channel was starved") + } + } + pingCtx, pingCancel := context.WithTimeout(ctx, 20*time.Millisecond) + defer pingCancel() + if err := c.Ping(pingCtx); err == nil { + t.Fatal("unanswered ping should time out") + } + ping := readFlowWire(t, wire) + if ping.Request == nil || ping.Request.Method != "ping" { + t.Fatal("RPC did not bypass full window") + } + // The canceled, already-fragmenting message must still finish. + for { + c.flow.mu.Lock() + total := c.flow.channels[a.URI()].sent + c.flow.mu.Unlock() + hostFlowNotification(wire, "channel/credit", ahptypes.ChannelCreditParams{Channel: a.URI(), ConsumedBytes: total}) + m := readFlowWire(t, wire) + if m.Notification == nil || m.Notification.Method != "channel/frame" { + t.Fatal(m) + } + var p ahptypes.ChannelFrameParams + _ = json.Unmarshal(m.Notification.Params, &p) + if p.Final != nil && *p.Final { + break + } + } + waitFlow(t, func() bool { return a.FlowControl().BufferedSendBytes == 0 }) +} + +func TestFlowQueueCapAndCanceledUnsentData(t *testing.T) { + cfg := DefaultConfig() + cfg.Windowing.MaxBufferedSendBytes = 600 + c, wire, ctx := flowTestClient(t, cfg) + s := subscribeFlow(t, c, wire, ctx, "ahp-terminal:/t", SubscribeOptions{Window: &WindowOptions{}}, acceptedFlow(1, 1024)) + if err := c.Notify(ctx, "app/data", map[string]any{"channel": s.URI(), "text": "first"}); err != nil { + t.Fatal(err) + } + _ = readFlowWire(t, wire) + sendCtx, cancel := context.WithCancel(ctx) + done := make(chan error, 1) + go func() { + done <- c.Notify(sendCtx, "app/data", map[string]any{"channel": s.URI(), "text": strings.Repeat("x", 120)}) + }() + waitFlow(t, func() bool { return s.FlowControl().BufferedSendBytes > 0 }) + if err := c.Notify(ctx, "app/data", map[string]any{"channel": s.URI(), "text": strings.Repeat("x", 120)}); !errors.Is(err, ErrBackpressure) { + t.Fatalf("cap returned %v", err) + } + cancel() + if !errors.Is(<-done, context.Canceled) { + t.Fatal("canceled unsent data remained pending") + } + waitFlow(t, func() bool { return s.FlowControl().BufferedSendBytes == 0 }) + c.flow.mu.Lock() + total := c.flow.channels[s.URI()].sent + c.flow.mu.Unlock() + hostFlowNotification(wire, "channel/credit", ahptypes.ChannelCreditParams{Channel: s.URI(), ConsumedBytes: total}) + waitFlow(t, func() bool { return s.FlowControl().Writable }) + select { + case text := <-wire.sent: + t.Fatalf("canceled data sent: %s", text) + case <-time.After(15 * time.Millisecond): + } +} + +func TestFlowFragmentSizingSurrogatesAndLongURI(t *testing.T) { + text := strings.Repeat("漢😀\\\"<\n", 300) + for _, remaining := range []int64{1, 200, 512, 1024} { + frame, count, err := fragmentText("ahp-terminal:/t", text, 1024, remaining) + if err != nil { + t.Fatal(err) + } + if utf16Bytes(frame) > 1024 || !utf8.ValidString(text[:count]) { + t.Fatal("invalid fragment prefix or size") + } + var n ahptypes.JsonRpcNotification + _ = json.Unmarshal([]byte(frame), &n) + var p ahptypes.ChannelFrameParams + _ = json.Unmarshal(n.Params, &p) + if p.Data != text[:count] { + t.Fatal("fragment data changed") + } + } + if _, _, err := fragmentText(strings.Repeat("x", 1000), "a", 1024, 512); err == nil { + t.Fatal("impossible frame accepted") + } + for _, p := range []string{`{"data":"\ud800"}`, `{"data":"\udc00"}`} { + if validFragmentString([]byte(p)) { + t.Fatal("split surrogate accepted") + } + } + for _, p := range []string{`{"data":"\ud83d\ude00"}`, `{"data":"\\ud800"}`, `{"data":"😀"}`} { + if !validFragmentString([]byte(p)) { + t.Fatal("valid fragment rejected", p) + } + } +} + +func TestFlowResubscribeAndOwnerClose(t *testing.T) { + c, wire, ctx := flowTestClient(t, DefaultConfig()) + uri := "ahp-terminal:/t" + sub := subscribeFlow(t, c, wire, ctx, uri, SubscribeOptions{Window: &WindowOptions{Receive: ReceiveManual}}, acceptedFlow(512, 1024)) + initialFlowSnapshot(wire, uri, 0) + d, err := sub.Next(ctx) + if err != nil { + t.Fatal(err) + } + sub.Close() + m := readFlowWire(t, wire) + if m.Notification.Method != "unsubscribe" { + t.Fatal("close must unsubscribe, not credit unread data") + } + d.Release() + c.flow.mu.Lock() + if c.flow.channels[uri] != nil || len(c.flow.order) != 0 { + t.Fatal("closed subscription remained in scheduler") + } + c.flow.mu.Unlock() + select { + case text := <-wire.sent: + t.Fatalf("closed owner sent %s", text) + case <-time.After(15 * time.Millisecond): + } +} + +func TestFlowReconnectRecoveryKindsAndDropped(t *testing.T) { + for _, kind := range []ahptypes.ChannelRecoveryKind{ahptypes.ChannelRecoveryKindSnapshot, ahptypes.ChannelRecoveryKindReplay, ahptypes.ChannelRecoveryKindLive} { + t.Run(string(kind), func(t *testing.T) { + c, wire, ctx := flowTestClient(t, DefaultConfig()) + uri := "ahp-terminal:/kept" + sub := subscribeFlow(t, c, wire, ctx, uri, SubscribeOptions{Window: &WindowOptions{Receive: ReceiveManual}}, acceptedFlow(512, 1024)) + dropped := subscribeFlow(t, c, wire, ctx, "ahp-terminal:/dropped", SubscribeOptions{Window: &WindowOptions{}}, acceptedFlow(512, 1024)) + initialFlowSnapshot(wire, uri, 7) + old, err := sub.Next(ctx) + if err != nil { + t.Fatal(err) + } + done := make(chan error, 1) + seq := int64(7) + go func() { + _, err := c.ReconnectWithParams(ctx, ahptypes.ReconnectParams{ClientId: "client", Subscriptions: []string{uri, dropped.URI()}, Windows: &ahptypes.ChannelDeliveryResumeOptions{Items: []ahptypes.ChannelDeliveryResumeParams{ + {Channel: uri, FlowControl: ahptypes.SubscriptionFlowControlOptions{Receive: ahptypes.ChannelReceiveLimits{WindowBytes: 512, MaximumFrameBytes: 1024}}, LastAppliedServerSeq: &seq}, + {Channel: dropped.URI(), FlowControl: ahptypes.SubscriptionFlowControlOptions{Receive: ahptypes.ChannelReceiveLimits{WindowBytes: 512, MaximumFrameBytes: 1024}}}, + }}}) + done <- err + }() + request := readFlowWire(t, wire) + hostFlowReply(t, wire, request, map[string]any{"type": "snapshot", "snapshots": []any{}, "windows": map[string]any{"items": []ahptypes.ResumedChannelSubscription{{Channel: uri, FlowControl: *acceptedFlow(512, 1024), Recovery: kind}}}}) + if kind == ahptypes.ChannelRecoveryKindSnapshot { + initialFlowSnapshot(wire, uri, 8) + } else { + hostFlowNotification(wire, "app/data", map[string]any{"channel": uri, "text": "new"}) + } + if err := <-done; err != nil { + t.Fatal(err) + } + old.Release() + fresh, err := sub.Next(ctx) + if err != nil { + t.Fatal(err) + } + fresh.Release() + m := readFlowWire(t, wire) + var credit ahptypes.ChannelCreditParams + _ = json.Unmarshal(m.Notification.Params, &credit) + if credit.ConsumedBytes != fresh.WireBytes() { + t.Fatal("old epoch credited on new connection") + } + if !errors.Is(dropped.Err(), ErrSubscriptionDropped) { + t.Fatal("missing reconnect window not dropped") + } + }) + } +} + +func TestFlowTCPNeverInitialOrResumed(t *testing.T) { + c, _, ctx := flowTestClient(t, DefaultConfig()) + if _, err := c.Initialize(ctx, "client", ahptypes.SupportedProtocolVersions(), []string{"ahp-tcp:/t"}); err == nil { + t.Fatal("TCP initial subscription accepted") + } + if _, err := c.ReconnectWithParams(ctx, ahptypes.ReconnectParams{ClientId: "client", Subscriptions: []string{"ahp-tcp:/t"}, Windows: &ahptypes.ChannelDeliveryResumeOptions{}}); err == nil { + t.Fatal("TCP resume accepted") + } +} + +func TestFlowConcurrentNextCancellationAndAutoAdvance(t *testing.T) { + c, wire, ctx := flowTestClient(t, DefaultConfig()) + s := subscribeFlow(t, c, wire, ctx, "ahp-terminal:/t", SubscribeOptions{Window: &WindowOptions{}}, acceptedFlow(512, 1024)) + initialFlowSnapshot(wire, s.URI(), 0) + d, err := s.Next(ctx) + if err != nil { + t.Fatal(err) + } + canceled, cancel := context.WithCancel(ctx) + cancel() + if _, err := s.Next(canceled); !errors.Is(err, context.Canceled) { + t.Fatal(err) + } + if !d.released.Load() { + t.Fatal("canceled next did not auto-release") + } + waiting, stop := context.WithCancel(ctx) + defer stop() + done := make(chan error, 1) + go func() { _, err := s.Next(waiting); done <- err }() + deadline, finish := context.WithTimeout(ctx, 15*time.Millisecond) + defer finish() + if _, err := s.Next(deadline); !errors.Is(err, context.DeadlineExceeded) { + t.Fatal("concurrent next ignored cancellation", err) + } + stop() + if err := <-done; !errors.Is(err, context.Canceled) { + t.Fatal(err) + } +} + +func TestFlowTCPResetDiscardsUnreadDataAndStopsSends(t *testing.T) { + c, wire, ctx := flowTestClient(t, DefaultConfig()) + s := subscribeFlow(t, c, wire, ctx, "ahp-tcp:/t", SubscribeOptions{Window: &WindowOptions{Receive: ReceiveManual}}, acceptedFlow(1024, 1024)) + initialFlowSnapshot(wire, s.URI(), 0) + snapshot, err := s.Next(ctx) + if err != nil { + t.Fatal(err) + } + snapshot.Release() + _ = readFlowWire(t, wire) + hostFlowNotification(wire, "tcp/data", ahptypes.TcpDataParams{Channel: s.URI(), Data: "aGVsbG8="}) + hostFlowNotification(wire, "tcp/reset", ahptypes.TcpResetParams{Channel: s.URI()}) + waitFlow(t, func() bool { c.flow.mu.Lock(); defer c.flow.mu.Unlock(); return c.flow.channels[s.URI()].reset }) + d, err := s.Next(ctx) + if err != nil { + t.Fatal(err) + } + if _, ok := d.Event().(SubscriptionEventTCPReset); !ok { + t.Fatal("reset left unread TCP data ahead of it") + } + if err := c.Notify(ctx, "tcp/data", ahptypes.TcpDataParams{Channel: s.URI(), Data: "eA=="}); !errors.Is(err, ErrClosed) { + t.Fatal("send after reset accepted", err) + } +} + +func TestFlowCreditSchedulingIsFair(t *testing.T) { + e := &flowEngine{cfg: WindowingConfig{MaxBufferedSendBytes: 1 << 20}, channels: make(map[string]*channelFlow), order: []string{"a", "b"}} + for _, uri := range e.order { + e.channels[uri] = &channelFlow{channel: uri, limits: acceptedFlow(512, 1024), released: 200, creditReady: true} + } + first, err := e.nextCredit() + if err != nil { + t.Fatal(err) + } + var n ahptypes.JsonRpcNotification + _ = json.Unmarshal([]byte(first.wire.text), &n) + if channel, _ := notificationChannel(n); channel != "a" { + t.Fatal(channel) + } + e.channels["a"].released += 200 + e.channels["a"].creditReady = true + second, err := e.nextCredit() + if err != nil { + t.Fatal(err) + } + _ = json.Unmarshal([]byte(second.wire.text), &n) + if channel, _ := notificationChannel(n); channel != "b" { + t.Fatal("busy channel starved another channel's credit") + } +} diff --git a/clients/go/ahptypes/state.generated.go b/clients/go/ahptypes/state.generated.go index ed37ca904..75025bc45 100644 --- a/clients/go/ahptypes/state.generated.go +++ b/clients/go/ahptypes/state.generated.go @@ -6530,7 +6530,7 @@ func (s *SnapshotState) UnmarshalJSON(data []byte) error { return err } s.Automations = &v - case containsAll(probe, "lifecycle"): + case containsAll(probe, "lifecycle", "provider"): var v SessionState if err := json.Unmarshal(data, &v); err != nil { return err diff --git a/clients/go/ahpws/transport.go b/clients/go/ahpws/transport.go index c9b0dc4b5..f8ee8df9c 100644 --- a/clients/go/ahpws/transport.go +++ b/clients/go/ahpws/transport.go @@ -53,7 +53,8 @@ const defaultReadLimit int64 = 32 * 1024 * 1024 // [github.com/coder/websocket.Conn]. Use [Connect] to dial or // [FromConn] to wrap an existing connection. type Transport struct { - conn *websocket.Conn + conn *websocket.Conn + readLimit int64 } // Connect dials the given URL and returns a ready-to-use @@ -80,7 +81,7 @@ func Connect(ctx context.Context, url string, opts ...DialOptions) (*Transport, limit = defaultReadLimit } conn.SetReadLimit(limit) - return &Transport{conn: conn}, nil + return &Transport{conn: conn, readLimit: limit}, nil } // FromConn wraps an already-connected @@ -91,6 +92,15 @@ func FromConn(conn *websocket.Conn) *Transport { return &Transport{conn: conn} } +// TransportCapabilities reports a safe UTF-16 frame limit derived from the +// configured UTF-8 message read limit. FromConn uses the unknown-limit fallback. +func (t *Transport) TransportCapabilities() ahp.TransportCapabilities { + return ahp.TransportCapabilities{ + SupportsWindowedDelivery: true, + MaximumFrameBytes: ahp.MaximumFrameBytesForReadLimit(t.readLimit), + } +} + // Send writes a single message. Text frames are produced for // [ahp.NewTextMessage] / [ahp.NewParsedMessage] payloads, binary // frames for [ahp.NewBinaryMessage] payloads. diff --git a/clients/go/ahpws/transport_test.go b/clients/go/ahpws/transport_test.go new file mode 100644 index 000000000..cc0004ceb --- /dev/null +++ b/clients/go/ahpws/transport_test.go @@ -0,0 +1,66 @@ +package ahpws + +import ( + "context" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/coder/websocket" + "github.com/microsoft/agent-host-protocol/clients/go/ahp" +) + +func TestWindowingReadLimitDefaults(t *testing.T) { + if got := (&Transport{readLimit: defaultReadLimit}).TransportCapabilities(); !got.SupportsWindowedDelivery || got.MaximumFrameBytes != 22369620 { + t.Fatal(got) + } + if got := FromConn(nil).TransportCapabilities().MaximumFrameBytes; got != 1048576 { + t.Fatal(got) + } + if got := (&Transport{readLimit: 1536}).TransportCapabilities().MaximumFrameBytes; got != 1024 { + t.Fatal(got) + } + for _, limit := range []int64{-1, 0} { + if ahp.MaximumFrameBytesForReadLimit(limit) != 1048576 { + t.Fatal("unknown limit default") + } + } + boxed := ahp.NewBoxedTransport(&Transport{readLimit: defaultReadLimit}) + if boxed.TransportCapabilities().MaximumFrameBytes != 22369620 { + t.Fatal("boxing hid capabilities") + } +} + +func TestConnectOverridesCoderNativeReadLimitAndPreservesText(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + text := `{"jsonrpc":"2.0","method":"app/data","params":{"channel":"ahp-terminal:/t","text":"` + strings.Repeat("漢", 12000) + `"}}` + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := websocket.Accept(w, r, nil) + if err != nil { + return + } + defer conn.CloseNow() + _ = conn.Write(ctx, websocket.MessageText, []byte(text)) + _, _, _ = conn.Read(ctx) + })) + defer server.Close() + transport, err := Connect(ctx, "ws"+strings.TrimPrefix(server.URL, "http")) + if err != nil { + t.Fatal(err) + } + defer transport.Close(ctx) + if transport.TransportCapabilities().MaximumFrameBytes != 22369620 { + t.Fatal("effective default not reported") + } + msg, err := transport.Recv(ctx) + if err != nil { + t.Fatal(err) + } + raw, ok := msg.RawJSONText() + if !ok || raw != text { + t.Fatal("received text changed") + } +} diff --git a/clients/go/examples/internal/demohost/demohost.go b/clients/go/examples/internal/demohost/demohost.go new file mode 100644 index 000000000..7b08c06da --- /dev/null +++ b/clients/go/examples/internal/demohost/demohost.go @@ -0,0 +1,348 @@ +// Package demohost is a small in-memory AHP host for the windowing examples. +// +// It implements only what the examples need: `initialize`, a windowed +// `subscribe` to one terminal, and the host side of the window rules: +// +// - It counts every data message it sends, in UTF-16 bytes, and stops while +// `sent - consumed` is at or above the client's window. +// - It counts every data message the client sends, and returns +// `channel/credit` after it has processed each one. +// +// A real host applies the same rules to every windowed subscription. +package demohost + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "strings" + "sync" + "time" + + "github.com/microsoft/agent-host-protocol/clients/go/ahp" + "github.com/microsoft/agent-host-protocol/clients/go/ahptypes" +) + +// Options configures a [Host]. +type Options struct { + // ClientWindowBytes is the largest window the host accepts for data that it sends. + ClientWindowBytes int64 + // HostWindowBytes is the window the host offers for data that the client sends. + HostWindowBytes int64 + // InputDelay is how long the host takes to process one message from the client. + InputDelay time.Duration +} + +const maximumFrameBytes = 16_384 + +// ErrClosed is returned by [Host.SendAction] after [Host.Close]. +var ErrClosed = errors.New("demohost: closed") + +// Host is an in-memory AHP host that follows the window rules. +type Host struct { + opts Options + client *pipe + transport *pipe + subscribed chan struct{} + done chan struct{} + closeOnce sync.Once + + mu sync.Mutex + limits *ahptypes.ChannelFlowControl + channel string + serverSeq int64 + + // Data that the host sends to the client. + sentBytes int64 + consumedBytes int64 + credit chan struct{} // closed and replaced when consumedBytes grows + + // Data that the client sends to the host. + receivedBytes int64 + processedBytes int64 + inputTail chan struct{} // closed when the previous input has been processed + receivedInput strings.Builder +} + +// New starts a host. Give [Host.ClientTransport] to [ahp.Connect]. +func New(opts Options) *Host { + client, transport := newPipe() + idle := make(chan struct{}) + close(idle) + h := &Host{ + opts: opts, + client: client, + transport: transport, + subscribed: make(chan struct{}), + done: make(chan struct{}), + credit: make(chan struct{}), + inputTail: idle, + } + go h.run() + return h +} + +// ClientTransport returns the half of the transport to give to the client. +func (h *Host) ClientTransport() ahp.Transport { return h.client } + +// Subscribed is closed after the host has accepted the subscription and sent the snapshot. +func (h *Host) Subscribed() <-chan struct{} { return h.subscribed } + +// ReceivedInput returns the terminal input that the host has processed, in order. +func (h *Host) ReceivedInput() string { + h.mu.Lock() + defer h.mu.Unlock() + return h.receivedInput.String() +} + +// SendAction sends one terminal action, first waiting for credit if the +// client's window is full. Call it after [Host.Subscribed] is closed. +func (h *Host) SendAction(ctx context.Context, action ahptypes.StateAction) error { + h.mu.Lock() + defer h.mu.Unlock() + if h.limits == nil { + return errors.New("demohost: not subscribed") + } + window := h.limits.ClientReceive.WindowBytes + if h.sentBytes-h.consumedBytes >= window { + logf("window full (%d of %d bytes outstanding); pausing output", h.sentBytes-h.consumedBytes, window) + for h.sentBytes-h.consumedBytes >= window { + credit := h.credit + h.mu.Unlock() + select { + case <-credit: + h.mu.Lock() + case <-ctx.Done(): + h.mu.Lock() + return ctx.Err() + case <-h.done: + h.mu.Lock() + return ErrClosed + } + } + logf("credit received (consumedBytes=%d); resuming output", h.consumedBytes) + } + h.serverSeq++ + h.sendDataLocked(notification{ + JsonRpc: ahptypes.JsonRpcV2, + Method: "action", + Params: actionEnvelope{ + Channel: h.channel, + // Encode the variant directly: the union's MarshalJSON would escape + // characters such as '>' that a JavaScript host sends as-is. + Action: json.RawMessage(encode(action.Value)), + ServerSeq: h.serverSeq, + }, + }) + return nil +} + +// Close closes the connection to the client. +func (h *Host) Close() { + h.closeOnce.Do(func() { + close(h.done) + _ = h.transport.Close(context.Background()) + }) +} + +type notification struct { + JsonRpc ahptypes.JsonRpcVersion `json:"jsonrpc"` + Method string `json:"method"` + Params any `json:"params"` +} + +// actionEnvelope has the same fields, in the same order, as [ahptypes.ActionEnvelope]. +type actionEnvelope struct { + Channel string `json:"channel"` + Action json.RawMessage `json:"action"` + ServerSeq int64 `json:"serverSeq"` +} + +type response struct { + JsonRpc ahptypes.JsonRpcVersion `json:"jsonrpc"` + ID uint64 `json:"id"` + Result any `json:"result"` +} + +func (h *Host) sendDataLocked(message any) { + text := encode(message) + h.sentBytes += utf16Bytes(text) + _ = h.transport.Send(context.Background(), ahp.NewTextMessage(text)) +} + +func (h *Host) sendControl(message any) { + _ = h.transport.Send(context.Background(), ahp.NewTextMessage(encode(message))) +} + +func (h *Host) run() { + for { + frame, err := h.transport.Recv(context.Background()) + if err != nil { + return + } + text, ok := frame.RawJSONText() + if !ok { + continue + } + var message ahptypes.JsonRpcMessage + if err := json.Unmarshal([]byte(text), &message); err != nil { + continue + } + switch { + case message.Request != nil: + switch message.Request.Method { + case "initialize": + h.sendControl(response{ + JsonRpc: ahptypes.JsonRpcV2, + ID: message.Request.ID, + Result: ahptypes.InitializeResult{ + ProtocolVersion: ahptypes.SupportedProtocolVersions()[0], + Snapshots: []ahptypes.Snapshot{}, + }, + }) + case "subscribe": + h.acceptSubscription(message.Request) + } + case message.Notification != nil: + switch message.Notification.Method { + case "channel/credit": + var params ahptypes.ChannelCreditParams + if json.Unmarshal(message.Notification.Params, ¶ms) == nil { + h.receiveCredit(params.ConsumedBytes) + } + case "dispatchAction": + var params ahptypes.DispatchActionParams + if json.Unmarshal(message.Notification.Params, ¶ms) == nil { + h.receiveInput(text, params.Action) + } + } + } + } +} + +func (h *Host) acceptSubscription(request *ahptypes.JsonRpcRequest) { + var params ahptypes.SubscribeParams + if err := json.Unmarshal(request.Params, ¶ms); err != nil || params.FlowControl == nil { + panic("demohost: the examples expect a windowed subscribe") + } + offered := params.FlowControl.Receive + + h.mu.Lock() + defer h.mu.Unlock() + h.channel = params.Channel + // A host can lower the limits that the client offered, but never raise them. + h.limits = &ahptypes.ChannelFlowControl{ + ClientReceive: ahptypes.ChannelReceiveLimits{ + WindowBytes: min(offered.WindowBytes, h.opts.ClientWindowBytes), + MaximumFrameBytes: min(offered.MaximumFrameBytes, maximumFrameBytes), + }, + HostReceive: ahptypes.ChannelReceiveLimits{WindowBytes: h.opts.HostWindowBytes, MaximumFrameBytes: maximumFrameBytes}, + } + logf("accepted windowed delivery: client window %d bytes (offered %d), host window %d bytes", + h.limits.ClientReceive.WindowBytes, offered.WindowBytes, h.limits.HostReceive.WindowBytes) + h.sendControl(response{ + JsonRpc: ahptypes.JsonRpcV2, + ID: request.ID, + Result: ahptypes.SubscribeResult{FlowControl: h.limits}, + }) + + // The first data message on a windowed subscription is always the snapshot. + h.sendDataLocked(notification{ + JsonRpc: ahptypes.JsonRpcV2, + Method: "channel/snapshot", + Params: ahptypes.ChannelSnapshotParams{ + Channel: h.channel, + Snapshot: ahptypes.Snapshot{ + Resource: h.channel, + State: ahptypes.SnapshotState{Terminal: &ahptypes.TerminalState{ + Title: "demo shell", + Content: []ahptypes.TerminalContentPart{}, + Lifecycle: ahptypes.TerminalLifecycleState{Value: &ahptypes.TerminalRunningLifecycleState{Status: ahptypes.TerminalLifecycleStatusRunning}}, + Claim: ahptypes.TerminalClaim{Value: &ahptypes.TerminalClientClaim{Kind: ahptypes.TerminalClaimKindClient, ClientId: "windowing-example"}}, + }}, + FromSeq: h.serverSeq, + }, + }, + }) + close(h.subscribed) +} + +func (h *Host) receiveCredit(consumedBytes int64) { + h.mu.Lock() + defer h.mu.Unlock() + // Credit is a running total, so ignore values that are not larger. + if consumedBytes > h.consumedBytes { + h.consumedBytes = consumedBytes + close(h.credit) + h.credit = make(chan struct{}) + } +} + +func (h *Host) receiveInput(text string, action ahptypes.StateAction) { + size := utf16Bytes(text) + h.mu.Lock() + h.receivedBytes += size + limits := h.limits.HostReceive + // The client can go past the window only with the message that fills it. + if h.receivedBytes-h.processedBytes > limits.WindowBytes+limits.MaximumFrameBytes { + logf("client went past the host window") + } + // Process input one message at a time, then return credit for it. + previous, done := h.inputTail, make(chan struct{}) + h.inputTail = done + h.mu.Unlock() + + go func() { + defer close(done) + select { + case <-previous: + case <-h.done: + return + } + select { + case <-time.After(h.opts.InputDelay): + case <-h.done: + return + } + h.mu.Lock() + if input, ok := action.Value.(*ahptypes.TerminalInputAction); ok { + h.receivedInput.WriteString(input.Data) + } + h.processedBytes += size + credit := ahptypes.ChannelCreditParams{Channel: h.channel, ConsumedBytes: h.processedBytes} + h.mu.Unlock() + h.sendControl(notification{JsonRpc: ahptypes.JsonRpcV2, Method: "channel/credit", Params: credit}) + }() +} + +// encode serializes a message the way a JavaScript host would, without +// escaping HTML characters. +func encode(message any) string { + var b bytes.Buffer + enc := json.NewEncoder(&b) + enc.SetEscapeHTML(false) + if err := enc.Encode(message); err != nil { + panic(err) + } + return strings.TrimSuffix(b.String(), "\n") +} + +// utf16Bytes returns the AHP size of text: two bytes per UTF-16 code unit. +// Go strings are UTF-8, so count code units per rune rather than using len. +func utf16Bytes(text string) int64 { + var size int64 + for _, r := range text { + if r > 0xffff { + size += 4 // a surrogate pair + } else { + size += 2 + } + } + return size +} + +func logf(format string, args ...any) { + fmt.Printf(" host | "+format+"\n", args...) +} diff --git a/clients/go/examples/internal/demohost/pipe.go b/clients/go/examples/internal/demohost/pipe.go new file mode 100644 index 000000000..e33bb6506 --- /dev/null +++ b/clients/go/examples/internal/demohost/pipe.go @@ -0,0 +1,107 @@ +package demohost + +import ( + "context" + "sync" + + "github.com/microsoft/agent-host-protocol/clients/go/ahp" +) + +// pipe is one half of an in-memory transport pair. Text sent on one half is +// received, unchanged, by the other half. Because the text is never +// reserialized, it reports support for windowed delivery. +type pipe struct { + mu sync.Mutex + inbox []string + changed chan struct{} // closed and replaced whenever inbox or closed changes + closed bool + peer *pipe +} + +var _ ahp.TransportCapabilitiesProvider = (*pipe)(nil) + +// newPipe returns a connected pair of in-memory transports. +func newPipe() (*pipe, *pipe) { + a := &pipe{changed: make(chan struct{})} + b := &pipe{changed: make(chan struct{})} + a.peer, b.peer = b, a + return a, b +} + +func (p *pipe) Send(ctx context.Context, m ahp.TransportMessage) error { + if err := ctx.Err(); err != nil { + return err + } + text, _, err := m.Bytes() + if err != nil { + return err + } + return p.peer.deliver(string(text)) +} + +func (p *pipe) Recv(ctx context.Context) (ahp.TransportMessage, error) { + for { + p.mu.Lock() + if len(p.inbox) > 0 { + text := p.inbox[0] + p.inbox = p.inbox[1:] + p.mu.Unlock() + return ahp.NewTextMessage(text), nil + } + if p.closed { + p.mu.Unlock() + return ahp.TransportMessage{}, closedError() + } + changed := p.changed + p.mu.Unlock() + select { + case <-changed: + case <-ctx.Done(): + return ahp.TransportMessage{}, ctx.Err() + } + } +} + +// Close closes both halves. Messages that are already queued can still be read. +func (p *pipe) Close(context.Context) error { + p.markClosed() + p.peer.markClosed() + return nil +} + +// TransportCapabilities reports that text is kept exactly as sent. The pipe +// has no frame size limit, so the client uses its default maximum. +func (p *pipe) TransportCapabilities() ahp.TransportCapabilities { + return ahp.TransportCapabilities{SupportsWindowedDelivery: true} +} + +// deliver queues text for this half to receive. Both halves close together, +// so a closed receiver also means a closed sender. +func (p *pipe) deliver(text string) error { + p.mu.Lock() + defer p.mu.Unlock() + if p.closed { + return closedError() + } + p.inbox = append(p.inbox, text) + p.signalLocked() + return nil +} + +func (p *pipe) markClosed() { + p.mu.Lock() + defer p.mu.Unlock() + if !p.closed { + p.closed = true + p.signalLocked() + } +} + +func (p *pipe) signalLocked() { + close(p.changed) + p.changed = make(chan struct{}) +} + +func closedError() error { + return &ahp.TransportError{Kind: "closed", Err: ahp.ErrClosed} +} diff --git a/clients/go/examples/windowing_backpressure/main.go b/clients/go/examples/windowing_backpressure/main.go new file mode 100644 index 000000000..fb36ae46b --- /dev/null +++ b/clients/go/examples/windowing_backpressure/main.go @@ -0,0 +1,230 @@ +// Command windowing_backpressure demonstrates windowed delivery with +// backpressure and a custom coalescer. +// +// Part 1, receiving: the client renders slowly. While it is busy, the host +// fills the window and the client queues what arrives. Adjacent queued +// messages are combined before the renderer sees them: +// +// - The built-in coalescer joins adjacent `terminal/data` output. +// - A custom coalescer keeps only the last of adjacent `terminal/resized` +// actions, because only the final size matters. +// +// Part 2, sending: the client pastes a large block of text into the terminal. +// The host processes input slowly, so the client's send window fills. The +// client waits until the subscription is writable before it sends each chunk, +// instead of queueing the whole paste in memory. +// +// Run: go run ./examples/windowing_backpressure +package main + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "log" + "sort" + "strings" + "sync" + "time" + + "github.com/microsoft/agent-host-protocol/clients/go/ahp" + "github.com/microsoft/agent-host-protocol/clients/go/ahptypes" + "github.com/microsoft/agent-host-protocol/clients/go/examples/internal/demohost" +) + +const terminal = "ahp-terminal:/demo" + +func logf(format string, args ...any) { + fmt.Printf(" client | "+format+"\n", args...) +} + +// latestResize runs on two adjacent messages that are still queued. It +// returns the combined message and true, or false to leave them separate. The +// client never passes it messages that the protocol requires to stay separate, +// such as actions with an origin. +var latestResize = ahp.Coalescer{ + Method: "action", + Merge: func(earlier, later ahptypes.JsonRpcNotification, c ahp.CoalescingContext) (ahptypes.JsonRpcNotification, bool) { + if c.Direction != ahp.CoalesceReceive { + return earlier, false + } + var first, second ahptypes.ActionEnvelope + if json.Unmarshal(earlier.Params, &first) != nil || json.Unmarshal(later.Params, &second) != nil { + return earlier, false + } + _, firstIsResize := first.Action.Value.(*ahptypes.TerminalResizedAction) + _, secondIsResize := second.Action.Value.(*ahptypes.TerminalResizedAction) + if firstIsResize && secondIsResize { + return later, true + } + return earlier, false + }, +} + +func terminalData(data string) ahptypes.StateAction { + return ahptypes.StateAction{Value: &ahptypes.TerminalDataAction{Type: ahptypes.ActionTypeTerminalData, Data: data}} +} + +func terminalResized(cols, rows int64) ahptypes.StateAction { + return ahptypes.StateAction{Value: &ahptypes.TerminalResizedAction{Type: ahptypes.ActionTypeTerminalResized, Cols: cols, Rows: rows}} +} + +func main() { + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + host := demohost.New(demohost.Options{ClientWindowBytes: 2048, HostWindowBytes: 1024, InputDelay: 15 * time.Millisecond}) + cfg := ahp.DefaultConfig() + cfg.Windowing.Coalescers = []ahp.Coalescer{latestResize} + // Allow at most 2048 bytes of unsent data to wait behind the message being sent. + cfg.Windowing.MaxBufferedSendBytes = 2048 + client, err := ahp.Connect(ctx, host.ClientTransport(), cfg) + if err != nil { + log.Fatalf("connect: %v", err) + } + if _, err := client.Initialize(ctx, "windowing-example", ahptypes.SupportedProtocolVersions(), nil); err != nil { + log.Fatalf("initialize: %v", err) + } + + // Count what the client did. This stream is for monitoring only and can + // drop events; it never affects credit. + var countsMu sync.Mutex + counts := map[string]int{} + go func() { + for event := range client.FlowControlEvents() { + countsMu.Lock() + counts[event.Kind]++ + countsMu.Unlock() + } + }() + + _, sub, err := client.SubscribeWithOptions(ctx, terminal, ahp.SubscribeOptions{ + Window: &ahp.WindowOptions{WindowBytes: 4096}, + }) + if err != nil { + log.Fatalf("subscribe: %v", err) + } + status := sub.FlowControl() + if status.Limits == nil { + log.Fatalf("the host did not accept windowed delivery (%s)", status.FallbackReason) + } + logf("subscribed: accepted window=%d bytes", status.Limits.ClientReceive.WindowBytes) + + // ─── Part 1: a slow renderer ──────────────────────────────────────────── + + fmt.Println("\nPart 1: slow renderer") + hostActions := []ahptypes.StateAction{ + terminalData("Building 120 modules\r\n"), + // The user drags the corner of the terminal window. + terminalResized(100, 30), + terminalResized(110, 32), + terminalResized(120, 36), + terminalResized(132, 40), + } + for percent := 10; percent <= 100; percent += 10 { + hostActions = append(hostActions, terminalData(fmt.Sprintf("[%3d%%] compiled %d modules\r\n", percent, percent*120/100))) + } + hostActions = append(hostActions, terminalData("Build finished\r\n")) + go func() { + <-host.Subscribed() + for _, action := range hostActions { + if err := host.SendAction(ctx, action); err != nil { + return + } + } + }() + + renders := 0 +render: + for { + delivery, err := sub.Next(ctx) + if err != nil { + log.Fatalf("next: %v", err) + } + renders++ + switch event := delivery.Event().(type) { + case ahp.SubscriptionEventSnapshot: + logf("render snapshot (%d bytes)", delivery.WireBytes()) + case ahp.SubscriptionEventAction: + switch action := event.Envelope.Action.Value.(type) { + case *ahptypes.TerminalResizedAction: + logf("render resize to %dx%d (%d bytes)", action.Cols, action.Rows, delivery.WireBytes()) + case *ahptypes.TerminalDataAction: + var lines []string + for _, line := range strings.Split(action.Data, "\r\n") { + if line != "" { + lines = append(lines, line) + } + } + logf("render %d line(s) of output (%d bytes): %q", len(lines), delivery.WireBytes(), lines[len(lines)-1]) + if strings.Contains(action.Data, "Build finished") { + break render + } + } + } + time.Sleep(60 * time.Millisecond) // The renderer is slow; the host fills the window meanwhile. + } + logf("rendered %d deliveries for %d host messages", renders, len(hostActions)+1) + + // ─── Part 2: pasting into a slow terminal ─────────────────────────────── + + fmt.Println("\nPart 2: pasting into a slow terminal") + paste := make([]string, 8) + for i := range paste { + paste[i] = fmt.Sprintf("echo \"line %d of the pasted script\"\n", i+1) + } + waits := 0 + for i, chunk := range paste { + // Wait for room in the host window and the local send queue, rather + // than queueing the whole paste. With one producer, this keeps sends + // within the limit. + if status := sub.FlowControl(); !status.Writable { + logf("not writable (%d bytes awaiting credit, %d bytes queued); waiting", status.OutstandingSendBytes, status.BufferedSendBytes) + waits++ + if err := sub.WaitWritable(ctx); err != nil { + log.Fatalf("wait writable: %v", err) + } + } + logf("send chunk %d of %d", i+1, len(paste)) + input := ahptypes.StateAction{Value: &ahptypes.TerminalInputAction{Type: ahptypes.ActionTypeTerminalInput, Data: chunk}} + // Dispatch returns when the message is written to the transport. + if _, err := client.Dispatch(ctx, terminal, input); err != nil { + // If another producer filled the queue first, the send fails + // instead of growing the queue. Wait, then send the chunk again. + if !errors.Is(err, ahp.ErrBackpressure) { + log.Fatalf("dispatch: %v", err) + } + if err := sub.WaitWritable(ctx); err != nil { + log.Fatalf("wait writable: %v", err) + } + if _, err := client.Dispatch(ctx, terminal, input); err != nil { + log.Fatalf("dispatch: %v", err) + } + } + } + // Wait until the host has processed, and credited, every chunk. + for status := range sub.FlowControlChanges() { + if status.OutstandingSendBytes == 0 && status.BufferedSendBytes == 0 { + break + } + } + intact := host.ReceivedInput() == strings.Join(paste, "") + logf("sent %d chunks, waiting %d times; host received the paste intact: %t", len(paste), waits, intact) + + fmt.Println("\nFlow-control events:") + countsMu.Lock() + kinds := make([]string, 0, len(counts)) + for kind := range counts { + kinds = append(kinds, kind) + } + sort.Strings(kinds) + for _, kind := range kinds { + fmt.Printf(" %s: %d\n", kind, counts[kind]) + } + countsMu.Unlock() + + sub.Close() + host.Close() + _ = client.Shutdown(ctx) +} diff --git a/clients/go/examples/windowing_basic/main.go b/clients/go/examples/windowing_basic/main.go new file mode 100644 index 000000000..06499e0c0 --- /dev/null +++ b/clients/go/examples/windowing_basic/main.go @@ -0,0 +1,110 @@ +// Command windowing_basic demonstrates basic windowed delivery. +// +// The client offers windowed delivery on a terminal subscription. The host +// accepts a smaller window, then sends a snapshot and terminal output. When +// the window is full, the host waits; each time the client takes the next +// delivery, it releases the previous one and returns credit, so the host can +// continue. +// +// Run: go run ./examples/windowing_basic +package main + +import ( + "context" + "fmt" + "log" + "strings" + "time" + "unicode" + + "github.com/microsoft/agent-host-protocol/clients/go/ahp" + "github.com/microsoft/agent-host-protocol/clients/go/ahptypes" + "github.com/microsoft/agent-host-protocol/clients/go/examples/internal/demohost" +) + +const terminal = "ahp-terminal:/demo" + +var output = []string{ + "$ npm test", + "> demo@1.0.0 test", + "> node --test", + "ok 1 - parses the configuration file", + "ok 2 - rejects an unknown option", + "ok 3 - reads input from stdin", + "ok 4 - writes the report", + "# pass 4, fail 0", +} + +func logf(format string, args ...any) { + fmt.Printf(" client | "+format+"\n", args...) +} + +func main() { + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + // The host accepts at most a 1024-byte window, so it fills after a few messages. + host := demohost.New(demohost.Options{ClientWindowBytes: 1024, HostWindowBytes: 4096}) + client, err := ahp.Connect(ctx, host.ClientTransport(), ahp.DefaultConfig()) + if err != nil { + log.Fatalf("connect: %v", err) + } + if _, err := client.Initialize(ctx, "windowing-example", ahptypes.SupportedProtocolVersions(), nil); err != nil { + log.Fatalf("initialize: %v", err) + } + + // Windowing is off by default. Opt in for this subscription and offer a + // 4096-byte window. Coalescing is turned off so that every host message is + // a separate delivery; see windowing_backpressure for coalescing. + coalescing := false + _, sub, err := client.SubscribeWithOptions(ctx, terminal, ahp.SubscribeOptions{ + Window: &ahp.WindowOptions{WindowBytes: 4096, Coalescing: &coalescing}, + }) + if err != nil { + log.Fatalf("subscribe: %v", err) + } + status := sub.FlowControl() + if status.Limits == nil { + log.Fatalf("the host did not accept windowed delivery (%s)", status.FallbackReason) + } + logf("subscribed: mode=%s, accepted window=%d bytes", status.Mode, status.Limits.ClientReceive.WindowBytes) + + // The host writes terminal output as fast as the window allows. + go func() { + <-host.Subscribed() + for _, line := range output { + action := ahptypes.StateAction{Value: &ahptypes.TerminalDataAction{Type: ahptypes.ActionTypeTerminalData, Data: line + "\r\n"}} + if err := host.SendAction(ctx, action); err != nil { + return + } + } + }() + + // With automatic release (the default), calling Next again releases the + // delivery that it returned before. The client then returns credit for it. + lines := 0 + for lines < len(output) { + delivery, err := sub.Next(ctx) + if err != nil { + log.Fatalf("next: %v", err) + } + buffered := sub.FlowControl().BufferedReceiveBytes + switch event := delivery.Event().(type) { + case ahp.SubscriptionEventSnapshot: + logf("snapshot (%d bytes; %d bytes not yet released)", delivery.WireBytes(), buffered) + case ahp.SubscriptionEventAction: + if data, ok := event.Envelope.Action.Value.(*ahptypes.TerminalDataAction); ok { + text := strings.TrimRightFunc(data.Data, unicode.IsSpace) + logf("output %q (%d bytes; %d bytes not yet released)", text, delivery.WireBytes(), buffered) + lines++ + } + } + time.Sleep(20 * time.Millisecond) // Simulate rendering the delivery. + } + // Leaving the loop does not unsubscribe; closing the subscription does. + sub.Close() + + logf("received all %d lines", lines) + host.Close() + _ = client.Shutdown(ctx) +} diff --git a/clients/kotlin/src/main/kotlin/com/microsoft/agenthostprotocol/generated/State.generated.kt b/clients/kotlin/src/main/kotlin/com/microsoft/agenthostprotocol/generated/State.generated.kt index 84eca910f..cf8227d7e 100644 --- a/clients/kotlin/src/main/kotlin/com/microsoft/agenthostprotocol/generated/State.generated.kt +++ b/clients/kotlin/src/main/kotlin/com/microsoft/agenthostprotocol/generated/State.generated.kt @@ -7670,7 +7670,8 @@ internal object SnapshotStateSerializer : KSerializer { // Try the most distinctive shape first. AutomationRunState has required // `automation`, `origin`, and `sessions`; AutomationState has // required `entries`; SessionState has required - // `lifecycle`; ChatState has required `turns`; ChangesetState has + // `lifecycle` + `provider` (TerminalState also has `lifecycle`); + // ChatState has required `turns`; ChangesetState has // required `status` + `files`; ResourceWatchState has required // `root` + `recursive`; AnnotationsState has required `annotations` // (checked after session, whose optional annotations summary reuses the @@ -7683,7 +7684,7 @@ internal object SnapshotStateSerializer : KSerializer { SnapshotState.AutomationRun(input.json.decodeFromJsonElement(AutomationRunState.serializer(), element)) obj.containsKey("entries") -> SnapshotState.Automations(input.json.decodeFromJsonElement(AutomationState.serializer(), element)) - obj.containsKey("lifecycle") -> SnapshotState.Session(input.json.decodeFromJsonElement(SessionState.serializer(), element)) + obj.containsKey("lifecycle") && obj.containsKey("provider") -> SnapshotState.Session(input.json.decodeFromJsonElement(SessionState.serializer(), element)) obj.containsKey("turns") -> SnapshotState.Chat(input.json.decodeFromJsonElement(ChatState.serializer(), element)) obj.containsKey("instanceId") && obj.containsKey("extensionId") && obj.containsKey("canvasId") -> SnapshotState.Canvas(input.json.decodeFromJsonElement(CanvasState.serializer(), element)) diff --git a/clients/rust/Cargo.lock b/clients/rust/Cargo.lock index 676a75eb3..12c4db4b0 100644 --- a/clients/rust/Cargo.lock +++ b/clients/rust/Cargo.lock @@ -8,6 +8,7 @@ version = "1.0.0" dependencies = [ "ahp-types", "ahp-ws", + "base64", "jiff", "serde", "serde_json", @@ -46,6 +47,12 @@ version = "1.0.102" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" +[[package]] +name = "base64" +version = "0.22.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" + [[package]] name = "bitflags" version = "1.3.2" diff --git a/clients/rust/README.md b/clients/rust/README.md index 56ef52c56..fe32b6129 100644 --- a/clients/rust/README.md +++ b/clients/rust/README.md @@ -85,3 +85,18 @@ cargo test --workspace ## Multi-host clients See [MULTI_HOST.md](MULTI_HOST.md) for the Rust SDK's multi-host registry, reconnect supervision, fan-in events, aggregated views, and single-host convenience API. +## Consumer-driven windowing + +The `ahp` client now supports opt-in windowed subscriptions. Existing +`ClientConfig`, `Client::connect`, ordinary subscription and observer APIs remain +source-compatible. Use `Client::builder().windowing(WindowingConfig::default())` +for connection defaults and `subscribe_windowed(uri, WindowOptions::default())` +for opt-in. Keep a `Delivery` alive while processing; Drop or `release()` returns +credit. The returned sender exposes `wait_writable()`, awaited `send`/`dispatch`, +and status watch streams. Default and custom connection-wide coalescers operate +only on adjacent queued data. + +See [the ahp examples](crates/ahp/README.md#windowed-subscriptions) for automatic +and manual release, send pressure, custom coalescers and managed reconnects, and +[the websocket limits](crates/ahp-ws/README.md#windowing-defaults) for transport-safe +UTF-16 frame defaults. diff --git a/clients/rust/crates/ahp-ws/README.md b/clients/rust/crates/ahp-ws/README.md index 2040fc339..d072ad976 100644 --- a/clients/rust/crates/ahp-ws/README.md +++ b/clients/rust/crates/ahp-ws/README.md @@ -59,3 +59,18 @@ With no TLS feature enabled, only `ws://` works; `wss://` fails at connect time. - [`ahp`](https://crates.io/crates/ahp) — the main client crate - [`ahp-types`](https://crates.io/crates/ahp-types) — wire types only - [Protocol documentation](https://microsoft.github.io/agent-host-protocol/) +## Windowing defaults + +This transport preserves exact incoming text/binary JSON and sends the client's +serialized text unchanged, so `Client::subscribe_windowed` can negotiate windowing. +Configure the client with `Client::builder().windowing(WindowingConfig::default())`; +configuration alone does not opt subscriptions in. + +Tungstenite 0.29 defaults to a 16 MiB incoming frame payload limit and a 64 MiB +incoming message limit. AHP uses UTF-16 bytes, not these UTF-8 sizes. The offered +default is `2 * floor(min(finite limits) / 3)`: **11,184,810** UTF-16 bytes. +When both limits are unlimited, the default is **1,048,576**. +`from_stream` inspects the stream's effective configuration, not library defaults. +Use `WebSocketTransport::connect_with_config(url, WebSocketConfig)` to customize +the actual websocket limits. Large accepted AHP data is fragmented automatically; +RPC results and ordinary initialization snapshots are not windowed. diff --git a/clients/rust/crates/ahp-ws/src/transport.rs b/clients/rust/crates/ahp-ws/src/transport.rs index dae77b0e0..8892f37d1 100644 --- a/clients/rust/crates/ahp-ws/src/transport.rs +++ b/clients/rust/crates/ahp-ws/src/transport.rs @@ -44,6 +44,16 @@ pub struct WebSocketTransport { } impl WebSocketTransport { + /// Connect with explicit websocket limits/TLS-independent protocol settings. + pub async fn connect_with_config( + url: &str, + config: tokio_tungstenite::tungstenite::protocol::WebSocketConfig, + ) -> Result { + let _parsed = Url::parse(url)?; + let (inner, _) = + tokio_tungstenite::connect_async_with_config(url, Some(config), false).await?; + Ok(Self { inner }) + } /// Open a new WebSocket connection to `url`. /// /// Both `ws://` and `wss://` schemes are accepted; the `wss://` TLS @@ -68,6 +78,18 @@ impl WebSocketTransport { } impl Transport for WebSocketTransport { + fn window_capabilities(&self) -> ahp::WindowTransportCapabilities { + let config = self.inner.get_config(); + let maximum = [config.max_frame_size, config.max_message_size] + .into_iter() + .flatten() + .min() + .map(|limit| 2 * (limit as u64 / 3)) + .unwrap_or(1_048_576); + ahp::WindowTransportCapabilities::default() + .raw_text(true) + .maximum_frame_bytes(maximum) + } async fn send(&mut self, msg: TransportMessage) -> Result<(), TransportError> { let frame = match msg { TransportMessage::Parsed(m) => { diff --git a/clients/rust/crates/ahp-ws/tests/windowing.rs b/clients/rust/crates/ahp-ws/tests/windowing.rs new file mode 100644 index 000000000..05d28995d --- /dev/null +++ b/clients/rust/crates/ahp-ws/tests/windowing.rs @@ -0,0 +1,84 @@ +#![allow(clippy::unwrap_used)] +use ahp::Transport; +use ahp_ws::WebSocketTransport; +use tokio::net::{TcpListener, TcpStream}; +use tokio_tungstenite::tungstenite::protocol::{Role, WebSocketConfig}; +use tokio_tungstenite::{MaybeTlsStream, WebSocketStream}; + +#[tokio::test] +async fn effective_websocket_limits_are_converted_conservatively() { + for (config, expected) in [ + (WebSocketConfig::default(), 11_184_810), + ( + WebSocketConfig::default() + .max_frame_size(Some(6144)) + .max_message_size(Some(3072)), + 2048, + ), + ( + WebSocketConfig::default() + .max_frame_size(None) + .max_message_size(Some(1536)), + 1024, + ), + ( + WebSocketConfig::default() + .max_frame_size(None) + .max_message_size(None), + 1_048_576, + ), + ] { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let (client, server) = tokio::join!( + TcpStream::connect(listener.local_addr().unwrap()), + listener.accept() + ); + let stream = WebSocketStream::from_raw_socket( + MaybeTlsStream::Plain(client.unwrap()), + Role::Client, + Some(config), + ) + .await; + let transport = WebSocketTransport::from_stream(stream); + assert!(transport.window_capabilities().raw_text); + assert_eq!( + transport.window_capabilities().maximum_frame_bytes, + Some(expected) + ); + drop(server); + } +} + +#[tokio::test] +async fn configured_transport_accepts_non_ascii_safe_bound() { + use futures_util::SinkExt; + use tokio_tungstenite::tungstenite::Message; + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let host = tokio::spawn(async move { + let (socket, _) = listener.accept().await.unwrap(); + let mut stream = tokio_tungstenite::accept_async(socket).await.unwrap(); + stream + .send(Message::Text( + format!("{{\"text\":\"{}\"}}", "漢".repeat(400)).into(), + )) + .await + .unwrap(); + }); + let config = WebSocketConfig::default() + .max_frame_size(Some(1536)) + .max_message_size(Some(1536)); + let mut transport = WebSocketTransport::connect_with_config(&format!("ws://{address}"), config) + .await + .unwrap(); + assert_eq!( + transport.window_capabilities().maximum_frame_bytes, + Some(1024) + ); + let ahp::TransportMessage::Text(text) = transport.recv().await.unwrap().unwrap() else { + unreachable!() + }; + assert!(ahp::windowing::utf16_bytes(&text) <= 1024); + assert!(text.len() > 1024); + host.await.unwrap(); +} diff --git a/clients/rust/crates/ahp/Cargo.toml b/clients/rust/crates/ahp/Cargo.toml index c4f310848..44ca95a85 100644 --- a/clients/rust/crates/ahp/Cargo.toml +++ b/clients/rust/crates/ahp/Cargo.toml @@ -27,9 +27,10 @@ thiserror = { workspace = true } tokio = { workspace = true } tracing = { workspace = true } jiff = { workspace = true } +base64 = "0.22" [dev-dependencies] -tokio = { workspace = true, features = ["full"] } +tokio = { workspace = true, features = ["full", "test-util"] } serde_json = { workspace = true } ahp-ws = { path = "../ahp-ws" } tracing-subscriber = "0.3" diff --git a/clients/rust/crates/ahp/README.md b/clients/rust/crates/ahp/README.md index 343027901..71e895193 100644 --- a/clients/rust/crates/ahp/README.md +++ b/clients/rust/crates/ahp/README.md @@ -71,3 +71,166 @@ See `tests/client_roundtrip.rs` for a complete in-memory example. - [`ahp-ws`](https://crates.io/crates/ahp-ws) — WebSocket transport - [Connecting to multiple hosts](https://github.com/microsoft/agent-host-protocol/blob/main/clients/rust/MULTI_HOST.md) — the [`hosts`](https://docs.rs/ahp/latest/ahp/hosts/) module wraps multi-host registry, reconnect, fan-in, and aggregated views; single-host consumers use `MultiHostClient::single` - [Protocol documentation](https://microsoft.github.io/agent-host-protocol/) +## Windowed subscriptions + +Windowing is **off by default**. Configure defaults through the non-breaking +builder, then opt in separately for each subscription: + +```rust +use ahp::{Client, SubscriptionEvent, WindowOptions, WindowingConfig}; +use ahp_ws::WebSocketTransport; + +async fn example() -> Result<(), Box> { +let client = Client::builder() + .default_request_timeout(Some(std::time::Duration::from_secs(30))) + .subscription_buffer(256) + .windowing(WindowingConfig::default() + .window_bytes(262_144) + .max_buffered_send_bytes(1_048_576)) + .connect(WebSocketTransport::connect("ws://localhost:12345").await?) + .await?; +client.initialize("example".into(), vec!["1.1.0".into()], vec![]).await?; + +let (_, mut terminal) = client + .subscribe_windowed("ahp-terminal:/t1", WindowOptions::default()) + .await?; +while let Some(delivery) = terminal.next().await? { + match delivery.event() { + SubscriptionEvent::Snapshot(snapshot) => println!("reset: {:?}", snapshot.state), + SubscriptionEvent::Action(action) => println!("apply: {:?}", action.action), + _ => {} + } + // Await your slow rendering here. Credit is withheld until delivery drops. + tokio::time::sleep(std::time::Duration::from_millis(10)).await; +} +Ok(()) +} +``` + +There is no auto/manual flag in Rust. Dropping `Delivery` at the end of the loop +body is automatic release. Moving it into another task or retaining it is manual +release; `delivery.release()` is idempotent and Drop releases any remaining credit. +The snapshot is always the first delivery when supplied, including ordinary +fallback (where it carries no credit). The host may lower the offered limits or +decline them; inspect `subscription.status()` for mode, fallback reason, limits, +queued amounts and writability. Offers do not change protocol version negotiation. + +```rust +async fn example(mut tcp: ahp::WindowedSubscription) -> Result<(), ahp::WindowError> { +let (tx, mut rx) = tokio::sync::mpsc::channel::(16); +let writer = tokio::spawn(async move { + while let Some(mut delivery) = rx.recv().await { + // Write tcp/data's decoded bytes to your bounded sink, awaiting completion. + // Retaining this delivery during the write slows the sender. + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + delivery.release(); + } +}); +while let Some(delivery) = tcp.next().await? { + if tx.send(delivery).await.is_err() { break; } +} +drop(tx); +let _ = writer.await; +Ok(()) +} +``` + +The sender's awaited methods complete after transport writing, not host processing. +`wait_writable()` is advisory, not a reservation. Queue admission waits when the +configured send budget is full. A single message larger than that budget is +admitted when the subscription's queue is empty, and fragments/window credit pace +it; the budget limits data waiting behind it. The send budget/status includes original queued inputs, including the +individual receipts retained by a coalesced send. Window accounting charges only +the actual merged wire text. Cancelling a queued send before it starts removes it; once fragmenting +starts, the logical message completes. + +```rust +async fn example(tcp: ahp::WindowedSubscription) -> Result<(), ahp::WindowError> { +// Create this TCP connection through createTcpConnection before subscribing. +let sender = tcp.sender(); +sender.wait_writable().await?; +sender.send("tcp/data", ahp::ahp_types::notifications::TcpDataParams { + channel: tcp.uri().into(), + data: "SGk=".into(), +}).await?; +let mut changes = tcp.status_changes(); +println!("pressure: {:?}", changes.borrow_and_update()); +Ok(()) +} +``` + +Coalescing is enabled by default and merges only adjacent queued data: +incoming `terminal/data`, compatible `chat/delta`, and outgoing TCP byte chunks. +It never merges snapshots, controls, echoes carrying origin, or individually +sequenced dispatches. Disable it with `WindowOptions::default().coalescing(false)`. +Register connection-wide custom coalescers before the built-ins: + +```rust +use ahp::{CoalescingContext, Direction, WindowingConfig}; +use ahp::ahp_types::messages::JsonRpcNotification; + +let config = WindowingConfig::default().coalescer( + |earlier: &JsonRpcNotification, later: &JsonRpcNotification, + context: &CoalescingContext| { + // An application-defined notification whose counts can be added losslessly. + if context.direction != Direction::Receive || later.method != "example/count" { + return None; + } + let total = earlier.params.as_ref()?["count"].as_u64()? + .checked_add(later.params.as_ref()?["count"].as_u64()?)?; + let mut merged = later.clone(); + merged.params.as_mut()?["count"] = total.into(); + Some(merged) + }); +``` + +Custom rules must preserve channel/method and mandatory event semantics. Receive +coalescing does not inject or require delivery latency preferences. Credits retain +original wire sizes, not the size of the local projection. Fragments are credited +as stored, except the final fragment held by the consumer. Consequently windowing +does not bound reassembled-message memory. + +Use `Client::flow_control_events()` for lossy connection telemetry; the same event +points emit `tracing` debug records. Existing global/broadcast event streams are +observers, remain lossy, and never release credit. + +Multi-host consumers use `HostConfig::with_windowing(config)` and +`MultiHostClient::subscribe_windowed(&host_id, uri, options)`. The owner queue +survives reconnect, including unread old deliveries; their later release never +credits a new connection. Snapshot/replay/live recovery is per channel, dropped +subscriptions report an error, pending sends fail instead of replaying, and TCP +channels never resume. `MultiHostStateMirror::apply_delivery` accepts this reliable +feed. Multi-host flow events include the host id. + +Dropping a managed handle also removes its runtime subscription intent. Failed or +cancelled subscribe calls do not leave inactive send queues; late replies are +installed for the same subscription generation or cleaned up without closing the +connection. Graceful shutdown drains already-enqueued ordinary messages and +always closes observer streams, including when a stalled transport must be aborted. + +Custom transports retain source compatibility and default to ordinary fallback. +For exact-text transports, implement `Transport::window_capabilities()` or wrap +with `ahp::transport::RawTextTransport::new(transport, Some(safe_utf16_limit))`. +Never return parsed-only data after asserting this capability: exact received text +is required for accounting. Advanced generated subscribe/reconnect requests use +the same negotiation hook; obtain their owning queue with `windowed_subscription`. + +### Examples + +Two runnable examples use an in-memory demo host that follows the window rules, +so they need no server: + +- [`examples/windowing_basic.rs`](https://github.com/microsoft/agent-host-protocol/blob/main/clients/rust/crates/ahp/examples/windowing_basic.rs) opts a terminal + subscription into windowed delivery and shows the host pausing when the + window is full and resuming as the client returns credit. +- [`examples/windowing_backpressure.rs`](https://github.com/microsoft/agent-host-protocol/blob/main/clients/rust/crates/ahp/examples/windowing_backpressure.rs) + adds a slow renderer, a custom coalescer that keeps only the latest terminal + resize, and a paste that waits for `wait_writable()` instead of queueing. + +```sh +cargo run -p ahp --example windowing_basic +cargo run -p ahp --example windowing_backpressure +``` + +The shared host is in +[`examples/support/windowing_demo_host.rs`](https://github.com/microsoft/agent-host-protocol/blob/main/clients/rust/crates/ahp/examples/support/windowing_demo_host.rs). diff --git a/clients/rust/crates/ahp/examples/support/windowing_demo_host.rs b/clients/rust/crates/ahp/examples/support/windowing_demo_host.rs new file mode 100644 index 000000000..9ec6a7a4e --- /dev/null +++ b/clients/rust/crates/ahp/examples/support/windowing_demo_host.rs @@ -0,0 +1,351 @@ +//! A small in-memory AHP host for the windowing examples. +//! +//! It implements only what the examples need: `initialize`, a windowed +//! `subscribe` to one terminal, and the host side of the window rules: +//! +//! - It counts every data message it sends, in UTF-16 bytes, and stops while +//! `sent - consumed` is at or above the client's window. +//! - It counts every data message the client sends, and returns +//! `channel/credit` after it has processed each one. +//! +//! A real host applies the same rules to every windowed subscription. + +// Each example includes this file with `#[path]`, and neither uses every item. +#![allow(dead_code)] + +use std::sync::{Arc, Mutex, MutexGuard}; +use std::time::Duration; + +use ahp::ahp_types::actions::StateAction; +use ahp::ahp_types::state::{ChannelFlowControl, ChannelReceiveLimits}; +use ahp::ahp_types::SUPPORTED_PROTOCOL_VERSIONS; +use ahp::windowing::utf16_bytes; +use ahp::{Transport, TransportError, TransportMessage, WindowTransportCapabilities}; +use serde_json::{json, Value}; +use tokio::sync::{mpsc, watch}; +use tokio::task::JoinHandle; + +/// The error type that the examples return. +pub type Error = Box; + +const MAXIMUM_FRAME_BYTES: i64 = 16_384; + +fn log(message: impl std::fmt::Display) { + println!(" host | {message}"); +} + +/// Options for [`DemoHost::start`]. +pub struct DemoHostOptions { + /// The largest window the host accepts for data that it sends. + pub client_window_bytes: i64, + /// The window the host offers for data that the client sends. + pub host_window_bytes: i64, + /// How long the host takes to process one message from the client. + pub input_delay: Duration, +} + +/// The client's half of an in-memory connection to a [`DemoHost`]. +/// +/// It passes JSON text through unchanged, so it reports raw-text support and +/// the client can offer windowed delivery over it. +pub struct DemoTransport { + inbound: mpsc::UnboundedReceiver, + outbound: mpsc::UnboundedSender, +} + +impl Transport for DemoTransport { + fn window_capabilities(&self) -> WindowTransportCapabilities { + WindowTransportCapabilities::default().raw_text(true) + } + + async fn send(&mut self, message: TransportMessage) -> Result<(), TransportError> { + let text = match message { + TransportMessage::Text(text) => text, + TransportMessage::Binary(bytes) => { + String::from_utf8(bytes).map_err(|e| TransportError::Protocol(e.to_string()))? + } + TransportMessage::Parsed(message) => serde_json::to_string(&message) + .map_err(|e| TransportError::Protocol(e.to_string()))?, + }; + self.outbound.send(text).map_err(|_| TransportError::Closed) + } + + async fn recv(&mut self) -> Result, TransportError> { + Ok(self.inbound.recv().await.map(TransportMessage::Text)) + } +} + +/// An in-memory host. Clones share the same connection. +#[derive(Clone)] +pub struct DemoHost { + shared: Arc, +} + +struct Shared { + options: DemoHostOptions, + // Dropping the sender closes the connection to the client. + outbound: Mutex>>, + state: Mutex, + subscribed: watch::Sender, + // The client's running credit total for data that the host sends. + consumed_bytes: watch::Sender, + // Data messages from the client, waiting to be processed one at a time. + input: mpsc::UnboundedSender<(u64, Value)>, + tasks: Mutex>>, +} + +#[derive(Default)] +struct State { + limits: Option, + channel: String, + server_seq: i64, + sent_bytes: u64, + received_bytes: u64, + processed_bytes: u64, + received_input: String, +} + +impl DemoHost { + /// Start the host. Give the returned transport to the client. + pub fn start(options: DemoHostOptions) -> (Self, DemoTransport) { + let (to_client, client_inbound) = mpsc::unbounded_channel(); + let (client_outbound, from_client) = mpsc::unbounded_channel(); + let (input, inputs) = mpsc::unbounded_channel(); + let host = Self { + shared: Arc::new(Shared { + options, + outbound: Mutex::new(Some(to_client)), + state: Mutex::default(), + subscribed: watch::Sender::new(false), + consumed_bytes: watch::Sender::new(0), + input, + tasks: Mutex::default(), + }), + }; + let tasks = vec![ + tokio::spawn(host.clone().run(from_client)), + tokio::spawn(host.clone().process_input(inputs)), + ]; + *lock(&host.shared.tasks) = tasks; + let transport = DemoTransport { + inbound: client_inbound, + outbound: client_outbound, + }; + (host, transport) + } + + /// Wait until the host has accepted the subscription and sent the snapshot. + pub async fn when_subscribed(&self) { + let _ = self + .shared + .subscribed + .subscribe() + .wait_for(|subscribed| *subscribed) + .await; + } + + /// Terminal input that the host has processed, in order. + pub fn received_input(&self) -> String { + lock(&self.shared.state).received_input.clone() + } + + /// Send one terminal action, first waiting for credit if the client's window is full. + pub async fn send_action(&self, action: StateAction) -> Result<(), Error> { + let (window, sent) = { + let state = lock(&self.shared.state); + let limits = state.limits.as_ref().ok_or("no windowed subscription")?; + (limits.client_receive.window_bytes as u64, state.sent_bytes) + }; + let mut consumed = self.shared.consumed_bytes.subscribe(); + let outstanding = sent - *consumed.borrow_and_update(); + if outstanding >= window { + log(format!( + "window full ({outstanding} of {window} bytes outstanding); pausing output" + )); + let consumed = *consumed + .wait_for(|consumed| sent - consumed < window) + .await?; + log(format!( + "credit received (consumedBytes={consumed}); resuming output" + )); + } + let message = { + let mut state = lock(&self.shared.state); + state.server_seq += 1; + json!({ + "jsonrpc": "2.0", + "method": "action", + "params": { + "channel": state.channel, + "action": action, + "serverSeq": state.server_seq, + }, + }) + }; + self.send_data(&message); + Ok(()) + } + + /// Close the connection and stop the host. + pub fn close(&self) { + lock(&self.shared.outbound).take(); + for task in lock(&self.shared.tasks).drain(..) { + task.abort(); + } + } + + fn send_data(&self, message: &Value) { + let text = message.to_string(); + lock(&self.shared.state).sent_bytes += utf16_bytes(&text); + self.send_text(text); + } + + fn send_control(&self, message: &Value) { + self.send_text(message.to_string()); + } + + fn send_text(&self, text: String) { + if let Some(outbound) = lock(&self.shared.outbound).as_ref() { + let _ = outbound.send(text); + } + } + + async fn run(self, mut from_client: mpsc::UnboundedReceiver) { + while let Some(text) = from_client.recv().await { + let Ok(message) = serde_json::from_str::(&text) else { + continue; + }; + match message["method"].as_str() { + Some("initialize") => self.send_control(&json!({ + "jsonrpc": "2.0", + "id": message["id"], + "result": { + "protocolVersion": SUPPORTED_PROTOCOL_VERSIONS[0], + "serverSeq": 0, + "snapshots": [], + }, + })), + Some("subscribe") => self.accept_subscription(&message), + Some("channel/credit") => { + // Credit is a running total, so ignore values that are not larger. + if let Some(total) = message["params"]["consumedBytes"].as_u64() { + self.shared.consumed_bytes.send_if_modified(|consumed| { + let larger = total > *consumed; + if larger { + *consumed = total; + } + larger + }); + } + } + Some("dispatchAction") => { + self.receive_input(&text, message["params"]["action"].clone()); + } + _ => {} + } + } + } + + fn accept_subscription(&self, message: &Value) { + let params = &message["params"]; + let offered = ¶ms["flowControl"]["receive"]; + let (Some(offered_window), Some(offered_frame)) = ( + offered["windowBytes"].as_i64(), + offered["maximumFrameBytes"].as_i64(), + ) else { + self.send_control(&json!({ + "jsonrpc": "2.0", + "id": message["id"], + "error": { "code": -32602, "message": "the examples expect a windowed subscribe" }, + })); + return; + }; + let channel = params["channel"].as_str().unwrap_or_default().to_owned(); + // A host can lower the limits that the client offered, but never raise them. + let limits = ChannelFlowControl { + client_receive: ChannelReceiveLimits { + window_bytes: offered_window.min(self.shared.options.client_window_bytes), + maximum_frame_bytes: offered_frame.min(MAXIMUM_FRAME_BYTES), + }, + host_receive: ChannelReceiveLimits { + window_bytes: self.shared.options.host_window_bytes, + maximum_frame_bytes: MAXIMUM_FRAME_BYTES, + }, + }; + log(format!( + "accepted windowed delivery: client window {} bytes (offered {offered_window}), host window {} bytes", + limits.client_receive.window_bytes, limits.host_receive.window_bytes, + )); + self.send_control(&json!({ + "jsonrpc": "2.0", + "id": message["id"], + "result": { "flowControl": limits }, + })); + let server_seq = { + let mut state = lock(&self.shared.state); + state.channel.clone_from(&channel); + state.limits = Some(limits); + state.server_seq + }; + + // The first data message on a windowed subscription is always the snapshot. + self.send_data(&json!({ + "jsonrpc": "2.0", + "method": "channel/snapshot", + "params": { + "channel": channel, + "snapshot": { + "resource": channel, + "state": { + "title": "demo shell", + "content": [], + "lifecycle": { "status": "running" }, + "claim": { "kind": "client", "clientId": "windowing-example" }, + }, + "fromSeq": server_seq, + }, + }, + })); + self.shared.subscribed.send_replace(true); + } + + fn receive_input(&self, text: &str, action: Value) { + let bytes = utf16_bytes(text); + let mut state = lock(&self.shared.state); + state.received_bytes += bytes; + if let Some(limits) = &state.limits { + let limits = &limits.host_receive; + // The client can go past the window only with the message that fills it. + if state.received_bytes - state.processed_bytes + > (limits.window_bytes + limits.maximum_frame_bytes) as u64 + { + log("client went past the host window"); + } + } + let _ = self.shared.input.send((bytes, action)); + } + + // Process input one message at a time, then return credit for it. + async fn process_input(self, mut inputs: mpsc::UnboundedReceiver<(u64, Value)>) { + while let Some((bytes, action)) = inputs.recv().await { + tokio::time::sleep(self.shared.options.input_delay).await; + let message = { + let mut state = lock(&self.shared.state); + if let Ok(StateAction::TerminalInput(input)) = serde_json::from_value(action) { + state.received_input.push_str(&input.data); + } + state.processed_bytes += bytes; + json!({ + "jsonrpc": "2.0", + "method": "channel/credit", + "params": { "channel": state.channel, "consumedBytes": state.processed_bytes }, + }) + }; + self.send_control(&message); + } + } +} + +fn lock(mutex: &Mutex) -> MutexGuard<'_, T> { + mutex.lock().unwrap_or_else(|e| e.into_inner()) +} diff --git a/clients/rust/crates/ahp/examples/windowing_backpressure.rs b/clients/rust/crates/ahp/examples/windowing_backpressure.rs new file mode 100644 index 000000000..34e4a434a --- /dev/null +++ b/clients/rust/crates/ahp/examples/windowing_backpressure.rs @@ -0,0 +1,257 @@ +//! Windowed delivery with backpressure and a custom coalescer. +//! +//! Part 1, receiving: the client renders slowly. While it is busy, the host +//! fills the window and the client queues what arrives. Adjacent queued +//! messages are combined before the renderer sees them: +//! +//! - The built-in coalescer joins adjacent `terminal/data` output. +//! - A custom coalescer keeps only the last of adjacent `terminal/resized` +//! actions, because only the final size matters. +//! +//! Part 2, sending: the client pastes a large block of text into the terminal. +//! The host processes input slowly, so the client's send window fills. The +//! client waits until the subscription is writable before it sends each chunk, +//! instead of queueing the whole paste in memory. +//! +//! Run with: +//! +//! ```sh +//! cargo run -p ahp --example windowing_backpressure +//! ``` + +#[path = "support/windowing_demo_host.rs"] +mod windowing_demo_host; + +use std::collections::BTreeMap; +use std::time::Duration; + +use ahp::ahp_types::actions::{ + StateAction, TerminalDataAction, TerminalInputAction, TerminalResizedAction, +}; +use ahp::ahp_types::messages::JsonRpcNotification; +use ahp::ahp_types::SUPPORTED_PROTOCOL_VERSIONS; +use ahp::{ + Client, CoalescingContext, Direction, FlowControlEvent, SubscriptionEvent, WindowOptions, + WindowingConfig, +}; +use tokio::sync::{broadcast, oneshot}; +use windowing_demo_host::{DemoHost, DemoHostOptions, Error}; + +const TERMINAL: &str = "ahp-terminal:/demo"; + +fn log(message: impl std::fmt::Display) { + println!(" client | {message}"); +} + +// A custom coalescer runs on two adjacent messages that are still queued. It +// returns the combined message, or `None` to leave them separate. The client +// never passes it messages that the protocol requires to stay separate, such +// as actions with an origin. +fn latest_resize( + earlier: &JsonRpcNotification, + later: &JsonRpcNotification, + context: &CoalescingContext, +) -> Option { + if context.direction != Direction::Receive || later.method != "action" { + return None; + } + let is_resize = |notification: &JsonRpcNotification| { + notification + .params + .as_ref() + .is_some_and(|params| params["action"]["type"] == "terminal/resized") + }; + (is_resize(earlier) && is_resize(later)).then(|| later.clone()) +} + +// Count flow-control events by kind until `stop` fires, then count the events +// that are still buffered and return the totals. +async fn count_events( + mut events: broadcast::Receiver, + mut stop: oneshot::Receiver<()>, +) -> BTreeMap { + use broadcast::error::{RecvError, TryRecvError}; + let mut counts = BTreeMap::new(); + loop { + let event = tokio::select! { + event = events.recv() => event, + _ = &mut stop => break, + }; + match event { + Ok(event) => *counts.entry(event.kind).or_default() += 1, + Err(RecvError::Lagged(_)) => {} + Err(RecvError::Closed) => return counts, + } + } + loop { + match events.try_recv() { + Ok(event) => *counts.entry(event.kind).or_default() += 1, + Err(TryRecvError::Lagged(_)) => {} + Err(_) => return counts, + } + } +} + +// A single-threaded runtime keeps the host and client log lines in a stable order. +#[tokio::main(flavor = "current_thread")] +async fn main() -> Result<(), Error> { + let (host, transport) = DemoHost::start(DemoHostOptions { + client_window_bytes: 2048, + host_window_bytes: 1024, + input_delay: Duration::from_millis(15), + }); + let client = Client::builder() + .windowing( + WindowingConfig::default() + .coalescer(latest_resize) + // Allow at most 2048 bytes of unsent data to wait behind the message being sent. + .max_buffered_send_bytes(2048), + ) + .connect(transport) + .await?; + let versions = SUPPORTED_PROTOCOL_VERSIONS.iter().map(|v| v.to_string()); + client + .initialize("windowing-example".into(), versions.collect(), vec![]) + .await?; + + // Count what the client did. This stream is for monitoring only and can drop + // events; it never affects credit. + let (stop_counting, stop) = oneshot::channel(); + let counter = tokio::spawn(count_events(client.flow_control_events(), stop)); + + let (_, mut subscription) = client + .subscribe_windowed(TERMINAL, WindowOptions::default().window_bytes(4096)) + .await?; + let accepted = subscription + .status() + .limits + .map_or("none".into(), |limits| { + limits.client_receive.window_bytes.to_string() + }); + log(format!("subscribed: accepted window={accepted} bytes")); + + // ─── Part 1: a slow renderer ──────────────────────────────────────────── + + println!("\nPart 1: slow renderer"); + let data = |data: String| StateAction::TerminalData(TerminalDataAction { data }); + let resize = |cols, rows| StateAction::TerminalResized(TerminalResizedAction { cols, rows }); + let mut host_actions = vec![ + data("Building 120 modules\r\n".into()), + // The user drags the corner of the terminal window. + resize(100, 30), + resize(110, 32), + resize(120, 36), + resize(132, 40), + ]; + host_actions.extend((1..=10).map(|step| { + let percent = step * 10; + data(format!( + "[{percent:>3}%] compiled {} modules\r\n", + percent * 120 / 100 + )) + })); + host_actions.push(data("Build finished\r\n".into())); + let host_messages = host_actions.len() + 1; // The snapshot, then each action. + let writer = tokio::spawn({ + let host = host.clone(); + async move { + host.when_subscribed().await; + for action in host_actions { + host.send_action(action).await?; + } + Ok::<_, Error>(()) + } + }); + + let mut renders = 0; + while let Some(delivery) = subscription.next().await? { + renders += 1; + let bytes = delivery.wire_bytes(); + match delivery.event() { + SubscriptionEvent::Snapshot(_) => log(format!("render snapshot ({bytes} bytes)")), + SubscriptionEvent::Action(envelope) => match &envelope.action { + StateAction::TerminalResized(size) => log(format!( + "render resize to {}x{} ({bytes} bytes)", + size.cols, size.rows + )), + StateAction::TerminalData(output) => { + let lines: Vec<_> = output + .data + .split("\r\n") + .filter(|l| !l.is_empty()) + .collect(); + let last = lines.last().copied().unwrap_or_default(); + log(format!( + "render {} line(s) of output ({bytes} bytes): {last:?}", + lines.len() + )); + if output.data.contains("Build finished") { + break; + } + } + _ => {} + }, + _ => {} + } + // The renderer is slow; the host fills the window meanwhile. + tokio::time::sleep(Duration::from_millis(60)).await; + } + log(format!( + "rendered {renders} deliveries for {host_messages} host messages" + )); + writer.await??; + + // ─── Part 2: pasting into a slow terminal ─────────────────────────────── + + println!("\nPart 2: pasting into a slow terminal"); + let paste: Vec = (1..=8) + .map(|i| format!("echo \"line {i} of the pasted script\"\n")) + .collect(); + let sender = subscription.sender(); + let mut waits = 0; + for (index, chunk) in paste.iter().enumerate() { + // Wait for room in the host window and the local send queue, rather than + // queueing the whole paste. With one producer, this keeps sends within the limit. + let status = subscription.status(); + if !status.writable { + log(format!( + "not writable ({} bytes awaiting credit, {} bytes queued); waiting", + status.outstanding_send_bytes, status.buffered_send_bytes + )); + waits += 1; + subscription.wait_writable().await?; + } + log(format!("send chunk {} of {}", index + 1, paste.len())); + // Completes when the message is written to the transport. If another + // producer has filled the send queue, this waits for room instead of + // failing, so there is no backpressure error to retry. An error here + // means the subscription or connection has closed. + sender + .dispatch(StateAction::TerminalInput(TerminalInputAction { + data: chunk.clone(), + })) + .await?; + } + // Wait until the host has processed, and credited, every chunk. + subscription + .status_changes() + .wait_for(|status| status.outstanding_send_bytes == 0 && status.buffered_send_bytes == 0) + .await?; + log(format!( + "sent {} chunks, waiting {waits} times; host received the paste intact: {}", + paste.len(), + host.received_input() == paste.concat() + )); + + let _ = stop_counting.send(()); + println!("\nFlow-control events:"); + for (kind, count) in counter.await? { + println!(" {kind}: {count}"); + } + + client.unsubscribe(TERMINAL.into()).await?; + drop(subscription); + host.close(); + client.shutdown().await; + Ok(()) +} diff --git a/clients/rust/crates/ahp/examples/windowing_basic.rs b/clients/rust/crates/ahp/examples/windowing_basic.rs new file mode 100644 index 000000000..99a1a66ed --- /dev/null +++ b/clients/rust/crates/ahp/examples/windowing_basic.rs @@ -0,0 +1,132 @@ +//! Basic windowed delivery. +//! +//! The client offers windowed delivery on a terminal subscription. The host +//! accepts a smaller window, then sends a snapshot and terminal output. When +//! the window is full, the host waits; each time the client finishes with a +//! delivery, it releases it and returns credit, so the host can continue. +//! +//! Run with: +//! +//! ```sh +//! cargo run -p ahp --example windowing_basic +//! ``` + +#[path = "support/windowing_demo_host.rs"] +mod windowing_demo_host; + +use std::time::Duration; + +use ahp::ahp_types::actions::{StateAction, TerminalDataAction}; +use ahp::ahp_types::SUPPORTED_PROTOCOL_VERSIONS; +use ahp::{Client, FlowControlStatus, SubscriptionEvent, WindowOptions}; +use windowing_demo_host::{DemoHost, DemoHostOptions, Error}; + +const TERMINAL: &str = "ahp-terminal:/demo"; +const OUTPUT: [&str; 8] = [ + "$ npm test", + "> demo@1.0.0 test", + "> node --test", + "ok 1 - parses the configuration file", + "ok 2 - rejects an unknown option", + "ok 3 - reads input from stdin", + "ok 4 - writes the report", + "# pass 4, fail 0", +]; + +fn log(message: impl std::fmt::Display) { + println!(" client | {message}"); +} + +fn accepted_window(status: &FlowControlStatus) -> String { + status.limits.as_ref().map_or("none".into(), |limits| { + limits.client_receive.window_bytes.to_string() + }) +} + +// A single-threaded runtime keeps the host and client log lines in a stable order. +#[tokio::main(flavor = "current_thread")] +async fn main() -> Result<(), Error> { + // The host accepts at most a 1024-byte window, so it fills after a few messages. + let (host, transport) = DemoHost::start(DemoHostOptions { + client_window_bytes: 1024, + host_window_bytes: 4096, + input_delay: Duration::ZERO, + }); + let client = Client::builder().connect(transport).await?; + let versions = SUPPORTED_PROTOCOL_VERSIONS.iter().map(|v| v.to_string()); + client + .initialize("windowing-example".into(), versions.collect(), vec![]) + .await?; + + // Windowing is off by default. Opt in for this subscription and offer a + // 4096-byte window. Coalescing is turned off so that every host message is a + // separate delivery; see windowing_backpressure.rs for coalescing. + let (_, mut subscription) = client + .subscribe_windowed( + TERMINAL, + WindowOptions::default() + .window_bytes(4096) + .coalescing(false), + ) + .await?; + let status = subscription.status(); + log(format!( + "subscribed: mode={}, accepted window={} bytes", + status.mode, + accepted_window(&status) + )); + + // The host writes terminal output as fast as the window allows. + let writer = tokio::spawn({ + let host = host.clone(); + async move { + host.when_subscribed().await; + for line in OUTPUT { + let data = format!("{line}\r\n"); + host.send_action(StateAction::TerminalData(TerminalDataAction { data })) + .await?; + } + Ok::<_, Error>(()) + } + }); + + // Each delivery holds its credit until it is dropped. Here it drops at the + // end of the loop body, before the next one is taken, which is automatic + // release. The client then returns credit for it. To release later, move + // the delivery elsewhere and call `release()` or drop it there. + let mut lines = 0; + while let Some(delivery) = subscription.next().await? { + let buffered = subscription.status().buffered_receive_bytes; + let bytes = delivery.wire_bytes(); + match delivery.event() { + SubscriptionEvent::Snapshot(_) => { + log(format!( + "snapshot ({bytes} bytes; {buffered} bytes not yet released)" + )); + } + SubscriptionEvent::Action(envelope) => { + if let StateAction::TerminalData(output) = &envelope.action { + let text = output.data.trim_end(); + log(format!( + "output {text:?} ({bytes} bytes; {buffered} bytes not yet released)" + )); + lines += 1; + } + } + _ => {} + } + tokio::time::sleep(Duration::from_millis(20)).await; // Simulate rendering the delivery. + if lines == OUTPUT.len() { + break; + } + } + + log(format!("received all {lines} lines")); + writer.await??; + // Dropping the subscription also unsubscribes, in the background. + client.unsubscribe(TERMINAL.into()).await?; + drop(subscription); + host.close(); + client.shutdown().await; + Ok(()) +} diff --git a/clients/rust/crates/ahp/src/client.rs b/clients/rust/crates/ahp/src/client.rs index 5e3ac3538..bc7cbd955 100644 --- a/clients/rust/crates/ahp/src/client.rs +++ b/clients/rust/crates/ahp/src/client.rs @@ -26,7 +26,7 @@ use std::future::Future; use std::pin::Pin; use std::sync::{ atomic::{AtomicU64, Ordering}, - Arc, + Arc, Weak, }; use std::time::Duration; @@ -45,8 +45,8 @@ use ahp_types::commands::{ use ahp_types::common::{Uri, ROOT_RESOURCE_URI}; use ahp_types::errors::json_rpc_error_codes; use ahp_types::messages::{ - ActionNotificationParams, JsonRpcError, JsonRpcErrorResponse, JsonRpcMessage, - JsonRpcNotification, JsonRpcRequest, JsonRpcSuccessResponse, JsonRpcVersion, + JsonRpcError, JsonRpcErrorResponse, JsonRpcMessage, JsonRpcNotification, JsonRpcRequest, + JsonRpcSuccessResponse, JsonRpcVersion, }; use ahp_types::notifications::{ AuthRequiredParams, SessionAddedParams, SessionRemovedParams, SessionSummaryChangedParams, @@ -91,6 +91,16 @@ impl Default for ClientConfig { #[derive(Debug, Clone)] #[non_exhaustive] pub enum SubscriptionEvent { + /// Initial or replacement channel state. + Snapshot(ahp_types::state::Snapshot), + /// TCP bytes received. + TcpData(ahp_types::notifications::TcpDataParams), + /// TCP half-close received. + TcpEof(ahp_types::notifications::TcpEofParams), + /// TCP abort received. + TcpReset(ahp_types::notifications::TcpResetParams), + /// Future notification unknown to this SDK. + Other(JsonRpcNotification), /// A write-ahead action envelope for this subscription's channel. Action(ActionEnvelope), /// `root/sessionAdded`: a new session was created on the root channel. @@ -189,7 +199,34 @@ pub struct DispatchHandle { type PendingMap = HashMap>>; +struct ControlRate { + since: tokio::time::Instant, + count: usize, +} +impl ControlRate { + fn new() -> Self { + Self { + since: tokio::time::Instant::now(), + count: 0, + } + } + async fn wait(&mut self) { + let period = Duration::from_millis(5); + if self.since.elapsed() >= period { + self.since = tokio::time::Instant::now(); + self.count = 0; + } + if self.count >= 64 { + tokio::time::sleep(period.saturating_sub(self.since.elapsed())).await; + self.since = tokio::time::Instant::now(); + self.count = 0; + } + self.count += 1; + } +} + struct Shared { + windowing: crate::windowing::Engine, pending: Mutex, subscriptions: Mutex>>, /// Top-level all-events broadcast. @@ -211,9 +248,47 @@ struct Shared { enum Outbound { Message(JsonRpcMessage), + Request(JsonRpcMessage, bool), + AwaitedMessage( + JsonRpcMessage, + oneshot::Sender>, + ), Shutdown, } +struct RequestGuard { + shared: Arc, + id: u64, + armed: bool, +} +impl Drop for RequestGuard { + fn drop(&mut self) { + if !self.armed { + return; + } + self.shared.windowing.expire_request(self.id); + if let Ok(mut pending) = self.shared.pending.try_lock() { + pending.remove(&self.id); + } else if let Ok(runtime) = tokio::runtime::Handle::try_current() { + let shared = self.shared.clone(); + let id = self.id; + runtime.spawn(async move { + shared.pending.lock().await.remove(&id); + }); + } + } +} + +struct AcquiringSubscription { + channel: Arc, + client: Client, +} +impl Drop for AcquiringSubscription { + fn drop(&mut self) { + self.channel.end_acquisition(self.client.clone()); + } +} + // ─── Server-initiated request handling ─────────────────────────────────────── /// Future returned by a [`ServerRequestHandler`]. @@ -356,6 +431,19 @@ struct DriveHandle { handle: Mutex>>, } +pub(crate) struct ClientBinding { + shared: Weak, + reader: Weak, +} +impl ClientBinding { + pub(crate) fn upgrade(&self) -> Option { + Some(Client { + shared: self.shared.upgrade()?, + _reader: self.reader.upgrade()?, + }) + } +} + impl Drop for DriveHandle { fn drop(&mut self) { if let Ok(mut guard) = self.handle.try_lock() { @@ -367,15 +455,33 @@ impl Drop for DriveHandle { } impl Client { + /// Configure a client without changing public ClientConfig struct literals. + pub fn builder() -> crate::ClientBuilder { + crate::ClientBuilder::default() + } /// Connect a client to the given transport. Spawns a background task /// that drives the transport until shutdown. pub async fn connect( transport: T, config: ClientConfig, ) -> Result { + Self::connect_configured(transport, config, crate::WindowingConfig::default()).await + } + + pub(crate) async fn connect_configured( + transport: T, + config: ClientConfig, + windowing: crate::WindowingConfig, + ) -> Result { + let capabilities = transport.window_capabilities(); let (outbound_tx, outbound_rx) = mpsc::channel::(64); let (all_events_tx, _) = broadcast::channel::(config.subscription_buffer); let shared = Arc::new(Shared { + windowing: crate::windowing::Engine::new( + windowing, + capabilities, + config.subscription_buffer, + ), pending: Mutex::new(HashMap::new()), subscriptions: Mutex::new(HashMap::new()), all_events: std::sync::Mutex::new(Some(all_events_tx)), @@ -398,16 +504,26 @@ impl Client { /// Gracefully shut down the client, aborting any in-flight requests /// with [`ClientError::Shutdown`]. pub async fn shutdown(&self) { - let _ = self.shared.outbound.send(Outbound::Shutdown).await; - // Fail any pending in-flight requests. - let mut pending = self.shared.pending.lock().await; - for (_, tx) in pending.drain() { - let _ = tx.send(Err(JsonRpcError { - code: -32000, - message: "client shut down".into(), - data: None, - })); + self.shared.windowing.close(None); + let must_abort = tokio::time::timeout( + Duration::from_secs(10), + self.shared.outbound.send(Outbound::Shutdown), + ) + .await + .is_err(); + if let Some(mut reader) = self._reader.handle.lock().await.take() { + if must_abort { + reader.abort(); + } + if tokio::time::timeout(Duration::from_secs(10), &mut reader) + .await + .is_err() + { + reader.abort(); + let _ = reader.await; + } } + teardown(&self.shared, None).await; } /// Send a JSON-RPC request and await its result. @@ -416,8 +532,48 @@ impl Client { P: Serialize, R: DeserializeOwned, { + self.request_inner(method, params, None).await + } + + async fn request_inner( + &self, + method: &str, + params: P, + options: Option<(crate::WindowOptions, Arc)>, + ) -> Result { + if self.shared.outbound.is_closed() { + return Err(ClientError::Shutdown); + } let id = self.shared.next_id.fetch_add(1, Ordering::Relaxed); - let params_val = serde_json::to_value(¶ms)?; + let mut params_val = serde_json::to_value(¶ms)?; + let expected = options.as_ref().map(|(_, channel)| channel.clone()); + let registration = if let Some((options, channel)) = options { + self.shared + .windowing + .register_owned(id, &channel, options, ¶ms_val) + } else { + self.shared + .windowing + .register_request(id, method, &mut params_val) + }; + if let Err(error) = registration { + if let Some(channel) = expected.or_else(|| { + params_val["channel"] + .as_str() + .and_then(|uri| self.shared.windowing.subscription(uri)) + }) { + self.shared.windowing.cleanup_unaccepted(&channel); + } + return Err(ClientError::Transport(crate::TransportError::Protocol( + error.to_string(), + ))); + } + let mut guard = RequestGuard { + shared: self.shared.clone(), + id, + armed: true, + }; + let tracked = self.shared.windowing.has_request(id); let params_any = if params_val.is_null() { None } else { @@ -439,7 +595,7 @@ impl Client { if self .shared .outbound - .send(Outbound::Message(req)) + .send(Outbound::Request(req, tracked)) .await .is_err() { @@ -451,13 +607,13 @@ impl Client { Some(dur) => match tokio::time::timeout(dur, rx).await { Ok(r) => r, Err(_) => { - self.shared.pending.lock().await.remove(&id); return Err(ClientError::Cancelled); } }, None => rx.await, }; + guard.armed = false; match result { Ok(Ok(value)) => Ok(serde_json::from_value(value)?), Ok(Err(e)) => Err(ClientError::Rpc(e)), @@ -467,17 +623,40 @@ impl Client { /// Send a JSON-RPC notification (fire-and-forget). pub async fn notify(&self, method: &str, params: P) -> Result<(), ClientError> { + if self.shared.outbound.is_closed() { + return Err(ClientError::Shutdown); + } let params_val = serde_json::to_value(¶ms)?; let params_any = if params_val.is_null() { None } else { Some(ahp_types::common::AnyValue::from(params_val)) }; - let msg = JsonRpcMessage::Notification(JsonRpcNotification { + let notification = JsonRpcNotification { jsonrpc: JsonRpcVersion::V2, method: method.into(), params: params_any, - }); + }; + if !ahp_types::notifications::counts_against_window(method) + && self.shared.windowing.has_windowed() + && crate::windowing::utf16_bytes(&serde_json::to_string(¬ification)?) + > crate::windowing::MAX_CONTROL_BYTES + { + return Err(ClientError::Transport(crate::TransportError::Protocol( + "control message exceeds size limit".into(), + ))); + } + if self.is_window_notification(¬ification) { + return self + .shared + .windowing + .enqueue(notification, None) + .await + .map_err(|e| { + ClientError::Transport(crate::TransportError::Protocol(e.to_string())) + }); + } + let msg = JsonRpcMessage::Notification(notification); self.shared .outbound .send(Outbound::Message(msg)) @@ -614,12 +793,68 @@ impl Client { /// this URI. In-flight [`SessionSubscription`] handles will see the /// channel close. pub async fn unsubscribe(&self, uri: String) -> Result<(), ClientError> { - { - let mut subs = self.shared.subscriptions.lock().await; - subs.remove(&uri); + self.unsubscribe_impl(&uri, None).await + } + + pub(crate) async fn unsubscribe_if_current( + &self, + uri: &str, + expected: &Arc, + ) -> Result<(), ClientError> { + self.unsubscribe_impl(uri, Some(expected)).await + } + + async fn unsubscribe_impl( + &self, + uri: &str, + expected: Option<&Arc>, + ) -> Result<(), ClientError> { + let mut client = self.clone(); + loop { + if let Some(current) = expected.and_then(|channel| { + channel + .binding + .lock() + .unwrap_or_else(|error| error.into_inner()) + .as_ref() + .and_then(ClientBinding::upgrade) + }) { + client = current; + } + let sender = client.shared.outbound.clone(); + let permit = sender.reserve().await.ok(); + if let Some(current) = expected.and_then(|channel| { + channel + .binding + .lock() + .unwrap_or_else(|error| error.into_inner()) + .as_ref() + .and_then(ClientBinding::upgrade) + }) { + if !Arc::ptr_eq(¤t.shared, &client.shared) { + drop(permit); + client = current; + continue; + } + } + let message = JsonRpcMessage::Notification(JsonRpcNotification { + jsonrpc: JsonRpcVersion::V2, + method: "unsubscribe".into(), + params: Some(serde_json::to_value(UnsubscribeParams { + channel: uri.into(), + })?), + }); + let removed = client.shared.windowing.remove_and(uri, expected, || { + if let Some(permit) = permit { + permit.send(Outbound::Message(message)); + } + }); + if removed { + client.shared.subscriptions.lock().await.remove(uri); + } + break; } - self.notify("unsubscribe", UnsubscribeParams { channel: uri }) - .await + Ok(()) } /// Subscribe to a top-level stream of every inbound event from this @@ -674,6 +909,215 @@ impl Client { Ok(DispatchHandle { client_seq }) } + /// Subscribe with consumer-driven window credit. A second call updates the same queue. + pub async fn subscribe_windowed( + &self, + uri: impl Into, + options: crate::WindowOptions, + ) -> Result<(SubscribeResult, crate::WindowedSubscription), crate::WindowError> { + if self.shared.outbound.is_closed() { + return Err(crate::WindowError::new("connection closed")); + } + let uri = uri.into(); + let offer = self.shared.windowing.prepare(&uri, options.clone())?; + let channel = self + .shared + .windowing + .subscription(&uri) + .ok_or_else(|| crate::WindowError::new("subscription cancelled"))?; + channel.begin_acquisition(); + let _acquiring = AcquiringSubscription { + channel: channel.clone(), + client: self.clone(), + }; + if uri.starts_with("ahp-tcp:") && offer.is_none() { + self.shared.windowing.cleanup_unaccepted(&channel); + return Err(crate::WindowError::new("TCP requires raw-text transport")); + } + let result = self + .request_inner( + "subscribe", + SubscribeParams { + channel: uri.clone(), + meta: None, + delivery: None, + view: None, + flow_control: offer.map(|receive| { + ahp_types::commands::SubscriptionFlowControlOptions { receive } + }), + }, + Some((options, channel.clone())), + ) + .await; + let result = match result { + Ok(result) => result, + Err(error) => { + self.shared.windowing.cleanup_unaccepted(&channel); + return Err(error.into()); + } + }; + if !self + .shared + .windowing + .subscription(&uri) + .is_some_and(|current| Arc::ptr_eq(¤t, &channel)) + { + return Err(crate::WindowError::new("subscription cancelled")); + } + *channel.binding.lock().unwrap_or_else(|e| e.into_inner()) = Some(self.binding()); + Ok(( + result, + crate::WindowedSubscription::new(channel, self.clone()), + )) + } + + fn binding(&self) -> ClientBinding { + ClientBinding { + shared: Arc::downgrade(&self.shared), + reader: Arc::downgrade(&self._reader), + } + } + + /// Obtain the owning queue after an advanced raw subscribe request. + /// + /// Handles for the same URI share one queue, rather than duplicating credit. + pub fn windowed_subscription( + &self, + uri: &str, + ) -> Result { + if self.shared.outbound.is_closed() { + return Err(crate::WindowError::new("connection closed")); + } + let channel = self + .shared + .windowing + .subscription(uri) + .ok_or_else(|| crate::WindowError::new("unknown window subscription"))?; + *channel.binding.lock().unwrap_or_else(|e| e.into_inner()) = Some(self.binding()); + Ok(crate::WindowedSubscription::new(channel, self.clone())) + } + + pub(crate) fn managed_channels(&self) -> Vec> { + self.shared.windowing.managed_channels() + } + + pub(crate) fn restore_channels( + &self, + channels: &[Arc], + ) -> Result<(), crate::WindowError> { + for channel in channels { + if channel.dropped() || channel.unowned_drop() { + continue; + } + self.shared.windowing.adopt(channel.clone())?; + *channel.binding.lock().unwrap_or_else(|e| e.into_inner()) = Some(self.binding()); + } + Ok(()) + } + + pub(crate) fn fail_subscription( + &self, + channel: &Arc, + error: crate::WindowError, + ) { + self.shared.windowing.fail_subscription(channel, error); + } + + pub(crate) fn is_connected(&self) -> bool { + !self.shared.outbound.is_closed() + } + + pub(crate) fn resume_windows( + &self, + ) -> Option { + self.shared.windowing.resume_windows() + } + + /// Observe lossy connection flow-control telemetry without releasing credit. + pub fn flow_control_events(&self) -> broadcast::Receiver { + self.shared.windowing.events.subscribe() + } + + /// Reconnect with explicit generated window resume parameters. + pub async fn reconnect_with_windows( + &self, + params: ReconnectParams, + ) -> Result { + Ok(self.request("reconnect", params).await?) + } + + fn is_window_notification(&self, n: &JsonRpcNotification) -> bool { + ahp_types::notifications::counts_against_window(&n.method) + && n.params + .as_ref() + .and_then(|p| p["channel"].as_str()) + .is_some_and(|uri| self.shared.windowing.subscription(uri).is_some()) + } + + pub(crate) async fn window_send( + &self, + method: &str, + params: Value, + ) -> Result<(), crate::WindowError> { + if self.shared.outbound.is_closed() { + return Err(crate::WindowError::new("connection closed")); + } + let notification = JsonRpcNotification { + jsonrpc: JsonRpcVersion::V2, + method: method.into(), + params: Some(params), + }; + if !self.is_window_notification(¬ification) { + if crate::windowing::utf16_bytes( + &serde_json::to_string(¬ification) + .map_err(|e| crate::WindowError::new(e.to_string()))?, + ) > crate::windowing::MAX_CONTROL_BYTES + { + return Err(crate::WindowError::new( + "control message exceeds size limit", + )); + } + let (tx, rx) = oneshot::channel(); + self.shared + .outbound + .send(Outbound::AwaitedMessage( + JsonRpcMessage::Notification(notification), + tx, + )) + .await + .map_err(|_| crate::WindowError::new("connection closed"))?; + return rx + .await + .map_err(|_| crate::WindowError::new("connection closed"))?; + } + let (tx, rx) = oneshot::channel(); + self.shared + .windowing + .enqueue(notification, Some(tx)) + .await?; + rx.await + .map_err(|_| crate::WindowError::new("connection closed"))? + } + + pub(crate) async fn window_dispatch( + &self, + channel: String, + action: StateAction, + ) -> Result { + let client_seq = self.shared.next_client_seq.fetch_add(1, Ordering::Relaxed) as i64; + self.window_send( + "dispatchAction", + serde_json::to_value(DispatchActionParams { + channel, + client_seq, + action, + }) + .map_err(|e| crate::WindowError::new(e.to_string()))?, + ) + .await?; + Ok(DispatchHandle { client_seq }) + } + // ─── Server-initiated requests (inbound) ───────────────────────────── /// Install a generic handler for inbound server-initiated requests. @@ -844,11 +1288,32 @@ async fn drive_transport( shared: Arc, mut outbound: mpsc::Receiver, ) { + let mut tick = tokio::time::interval(Duration::from_millis(5)); + let mut failure = None; + let mut send_controls = ControlRate::new(); + let mut receive_controls = ControlRate::new(); loop { tokio::select! { outbound_msg = outbound.recv() => { match outbound_msg { - Some(Outbound::Message(msg)) => { + Some(outbound @ (Outbound::Message(_) | Outbound::Request(..))) => { + let (msg, tracked) = match outbound { + Outbound::Message(message) => (message, false), + Outbound::Request(message, tracked) => (message, tracked), + _ => unreachable!(), + }; + if let JsonRpcMessage::Request(request) = &msg { + let pending = shared.pending.lock().await; + if pending.get(&request.id).map_or(true, oneshot::Sender::is_closed) { + shared.windowing.failed_request(request.id); + continue; + } + if tracked && !shared.windowing.request_sent(request.id) { + continue; + } + } + if matches!(&msg, JsonRpcMessage::Notification(n) if !ahp_types::notifications::counts_against_window(&n.method)) + && shared.windowing.has_windowed() { send_controls.wait().await; } if let Ok(wire) = TransportMessage::encode(&msg) { if let Err(err) = transport.send(wire).await { tracing::warn!(?err, "transport send failed"); @@ -860,14 +1325,77 @@ async fn drive_transport( let _ = transport.close().await; break; } + Some(Outbound::AwaitedMessage(msg, completion)) => { + if completion.is_closed() { continue; } + send_controls.wait().await; + let result = match TransportMessage::encode(&msg) { + Ok(wire) => transport.send(wire).await.map_err(|e| crate::WindowError::new(e.to_string())), + Err(error) => Err(crate::WindowError::new(error.to_string())), + }; + let _ = completion.send(result.clone()); + if let Err(error) = result { + failure = Some(error); let _ = transport.close().await; break; + } + } } } inbound = transport.recv() => { match inbound { Ok(Some(wire)) => { - match wire.into_parsed() { - Ok(msg) => dispatch_inbound(&shared, msg).await, - Err(err) => tracing::warn!(?err, "malformed frame"), + let (size, parsed) = match wire { + TransportMessage::Text(text) => (Some(crate::windowing::utf16_bytes(&text)), + serde_json::from_str(&text).map_err(|error| crate::TransportError::Protocol(error.to_string()))), + TransportMessage::Binary(bytes) => match std::str::from_utf8(&bytes) { + Ok(text) => (Some(crate::windowing::utf16_bytes(text)), + serde_json::from_str(text).map_err(|error| crate::TransportError::Protocol(error.to_string()))), + Err(error) => (None, Err(crate::TransportError::Protocol(error.to_string()))), + }, + TransportMessage::Parsed(message) => (None, Ok(message)), + }; + match parsed { + Ok(msg) => { + let result = match msg { + JsonRpcMessage::Notification(notification) => { + if !ahp_types::notifications::counts_against_window(¬ification.method) && shared.windowing.has_windowed() { + receive_controls.wait().await; + } + shared.windowing.inbound(notification, size).map(|n| n.map(JsonRpcMessage::Notification)) + }, + JsonRpcMessage::SuccessResponse(response) => { + match shared.windowing.response(response.id, &response.result) { + Ok(cleanup) => { + for uri in cleanup { + let text = serde_json::json!({"jsonrpc":"2.0","method":"unsubscribe","params":{"channel":uri}}).to_string(); + if let Err(error) = transport.send(TransportMessage::Text(text)).await { + failure = Some(crate::WindowError::new(error.to_string())); + break; + } + } + Ok(Some(JsonRpcMessage::SuccessResponse(response))) + } + Err(error) => Err(error), + } + }, + JsonRpcMessage::ErrorResponse(response) => { + shared.windowing.failed_request(response.id); + Ok(Some(JsonRpcMessage::ErrorResponse(response))) + } + other => Ok(Some(other)), + }; + match result { + Ok(Some(msg)) => dispatch_inbound(&shared, msg).await, + Ok(None) => {}, + Err(error) => { failure = Some(error); let _ = transport.close().await; break; } + } + } + Err(err) => { + if shared.windowing.has_windowed() { + failure = Some(crate::WindowError::new(err.to_string())); + let _ = transport.close().await; + break; + } + tracing::warn!(?err, "malformed frame"); + } } } Ok(None) => break, @@ -877,10 +1405,38 @@ async fn drive_transport( } } } + _ = shared.windowing.wake.notified() => { + match shared.windowing.next_send() { + Ok(Some(send)) => { + if send.control { send_controls.wait().await; } + if !send.is_live() { + send.complete(Err(crate::WindowError::new("subscription closed"))); + continue; + } + let result = tokio::time::timeout(Duration::from_secs(10), transport.send(TransportMessage::Text(send.text.clone()))).await; + match result { + Ok(Ok(())) => send.complete(Ok(())), + other => { + let error = crate::WindowError::new(format!("windowed transport send failed: {other:?}")); + send.complete(Err(error.clone())); + failure = Some(error); let _ = transport.close().await; break; + } + } + shared.windowing.wake.notify_one(); + } + Ok(None) => {}, + Err(error) => { failure = Some(error); let _ = transport.close().await; break; } + } + } + _ = tick.tick(), if shared.windowing.has_subscriptions() => { shared.windowing.wake.notify_one(); } } } - // Teardown: close everything so waiters see Shutdown. + teardown(&shared, failure).await; +} + +async fn teardown(shared: &Shared, failure: Option) { + shared.windowing.close(failure); let mut pending = shared.pending.lock().await; for (_, tx) in pending.drain() { let _ = tx.send(Err(JsonRpcError { @@ -960,47 +1516,64 @@ async fn handle_server_request(shared: &Shared, req: JsonRpcRequest) { } async fn handle_notification(shared: &Shared, n: JsonRpcNotification) { - let params_val: Value = n.params.unwrap_or(Value::Null); + let tracked = n + .params + .as_ref() + .and_then(|params| params["channel"].as_str()) + .is_some_and(|uri| shared.windowing.subscription(uri).is_some()); + if !tracked + && !matches!( + n.method.as_str(), + "action" + | "root/sessionAdded" + | "root/sessionRemoved" + | "root/sessionSummaryChanged" + | "auth/required" + ) + { + return; + } + if let Ok(event) = decode_notification(n.clone()) { + if let Some(channel) = n.params.as_ref().and_then(|p| p["channel"].as_str()) { + fan_out(shared, &channel.to_string(), event).await; + } + } +} +pub(crate) fn decode_notification(n: JsonRpcNotification) -> Result { + let value = n.params.clone().unwrap_or(Value::Null); + let decode = |error: serde_json::Error| error.to_string(); match n.method.as_str() { - "action" => { - if let Ok(envelope) = serde_json::from_value::(params_val) { - let channel = envelope.channel.clone(); - fan_out(shared, &channel, SubscriptionEvent::Action(envelope)).await; - } - } - "root/sessionAdded" => { - if let Ok(params) = serde_json::from_value::(params_val) { - let channel = params.channel.clone(); - fan_out(shared, &channel, SubscriptionEvent::SessionAdded(params)).await; - } - } - "root/sessionRemoved" => { - if let Ok(params) = serde_json::from_value::(params_val) { - let channel = params.channel.clone(); - fan_out(shared, &channel, SubscriptionEvent::SessionRemoved(params)).await; - } - } - "root/sessionSummaryChanged" => { - if let Ok(params) = serde_json::from_value::(params_val) { - let channel = params.channel.clone(); - fan_out( - shared, - &channel, - SubscriptionEvent::SessionSummaryChanged(params), - ) - .await; - } - } - "auth/required" => { - if let Ok(params) = serde_json::from_value::(params_val) { - let channel = params.channel.clone(); - fan_out(shared, &channel, SubscriptionEvent::AuthRequired(params)).await; - } - } - other => { - tracing::debug!(method = %other, "unhandled notification"); + "channel/snapshot" => { + serde_json::from_value::(value) + .map(|p| SubscriptionEvent::Snapshot(p.snapshot)) + .map_err(decode) } + "tcp/data" => serde_json::from_value(value) + .map(SubscriptionEvent::TcpData) + .map_err(decode), + "tcp/eof" => serde_json::from_value(value) + .map(SubscriptionEvent::TcpEof) + .map_err(decode), + "tcp/reset" => serde_json::from_value(value) + .map(SubscriptionEvent::TcpReset) + .map_err(decode), + "action" => serde_json::from_value(value) + .map(SubscriptionEvent::Action) + .map_err(decode), + "root/sessionAdded" => serde_json::from_value(value) + .map(SubscriptionEvent::SessionAdded) + .map_err(decode), + "root/sessionRemoved" => serde_json::from_value(value) + .map(SubscriptionEvent::SessionRemoved) + .map_err(decode), + "root/sessionSummaryChanged" => serde_json::from_value(value) + .map(SubscriptionEvent::SessionSummaryChanged) + .map_err(decode), + "auth/required" => serde_json::from_value(value) + .map(SubscriptionEvent::AuthRequired) + .map_err(decode), + _ => Ok(SubscriptionEvent::Other(n)), } } diff --git a/clients/rust/crates/ahp/src/hosts/factory.rs b/clients/rust/crates/ahp/src/hosts/factory.rs index 06adc5505..02f10d1f3 100644 --- a/clients/rust/crates/ahp/src/hosts/factory.rs +++ b/clients/rust/crates/ahp/src/hosts/factory.rs @@ -34,6 +34,10 @@ use super::types::HostId; /// use_factory(open_ws); /// ``` pub trait HostTransportFactory: Send + Sync + 'static { + /// Connection windowing defaults. Does not enable any subscription. + fn windowing_config(&self) -> crate::WindowingConfig { + crate::WindowingConfig::default() + } /// Open a fresh transport for `host_id`. /// /// Errors are surfaced as the host's `last_error` and trigger the @@ -57,3 +61,19 @@ where Box::pin(self(host_id)) } } + +pub(super) struct ConfiguredFactory { + pub(super) inner: std::sync::Arc, + pub(super) windowing: crate::WindowingConfig, +} +impl HostTransportFactory for ConfiguredFactory { + fn windowing_config(&self) -> crate::WindowingConfig { + self.windowing.clone() + } + fn open_transport( + &self, + host_id: HostId, + ) -> Pin> + Send + '_>> { + self.inner.open_transport(host_id) + } +} diff --git a/clients/rust/crates/ahp/src/hosts/mod.rs b/clients/rust/crates/ahp/src/hosts/mod.rs index 31fe0a8cf..4711d01d4 100644 --- a/clients/rust/crates/ahp/src/hosts/mod.rs +++ b/clients/rust/crates/ahp/src/hosts/mod.rs @@ -102,6 +102,7 @@ pub use factory::HostTransportFactory; pub use multi::MultiHostClient; pub use policy::{Backoff, ReconnectPolicy}; pub use types::{ - HostClientHandle, HostConfig, HostError, HostEvent, HostEventStream, HostHandle, HostId, - HostState, HostSubscriptionEvent, HostSubscriptionStream, HostedAgent, HostedSessionSummary, + HostClientHandle, HostConfig, HostError, HostEvent, HostEventStream, HostFlowControlEvent, + HostHandle, HostId, HostState, HostSubscriptionEvent, HostSubscriptionStream, HostedAgent, + HostedSessionSummary, }; diff --git a/clients/rust/crates/ahp/src/hosts/multi.rs b/clients/rust/crates/ahp/src/hosts/multi.rs index a35af0fc4..0b41362e2 100644 --- a/clients/rust/crates/ahp/src/hosts/multi.rs +++ b/clients/rust/crates/ahp/src/hosts/multi.rs @@ -51,6 +51,7 @@ struct MultiInner { pending_host_ids: StdMutex>, fan_out: broadcast::Sender, host_events: broadcast::Sender, + flow_events: broadcast::Sender, client_id_store: Arc, } @@ -72,12 +73,14 @@ impl MultiHostClient { pub fn with_client_id_store(client_id_store: Arc) -> Self { let (fan_out, _) = broadcast::channel(DEFAULT_EVENT_BUFFER); let (host_events, _) = broadcast::channel(DEFAULT_EVENT_BUFFER); + let (flow_events, _) = broadcast::channel(DEFAULT_EVENT_BUFFER); Self { inner: Arc::new(MultiInner { hosts: RwLock::new(HashMap::new()), pending_host_ids: StdMutex::new(HashSet::new()), fan_out, host_events, + flow_events, client_id_store, }), } @@ -174,6 +177,7 @@ impl MultiHostClient { resolved, self.inner.fan_out.clone(), self.inner.host_events.clone(), + self.inner.flow_events.clone(), ); let shared = tx.shared.clone(); hosts.insert(id.clone(), tx); @@ -423,6 +427,42 @@ impl MultiHostClient { entry.subscribe(uri).await } + /// Convenience: unsubscribe from `uri` on `host_id`. + pub async fn subscribe_windowed( + &self, + host_id: &HostId, + uri: impl Into, + options: crate::WindowOptions, + ) -> Result< + ( + ahp_types::commands::SubscribeResult, + crate::WindowedSubscription, + ), + HostError, + > { + let entry = self.host_entry(host_id).await?; + let (tx, rx) = tokio::sync::oneshot::channel(); + entry + .cmd_tx + .send(super::runtime::HostCommand::SubscribeWindowed { + uri: uri.into(), + options, + reply: tx, + }) + .await + .map_err(|_| HostError::HostShutDown(host_id.clone()))?; + let (result, mut subscription) = rx + .await + .map_err(|_| HostError::HostShutDown(host_id.clone()))??; + subscription.host_id = Some(host_id.clone()); + Ok((result, subscription)) + } + + /// Observe lossy flow telemetry from every host; never releases credit. + pub fn flow_control_events(&self) -> broadcast::Receiver { + self.inner.flow_events.subscribe() + } + /// Convenience: unsubscribe from `uri` on `host_id`. pub async fn unsubscribe(&self, host_id: &HostId, uri: String) -> Result<(), HostError> { let entry = self.host_entry(host_id).await?; diff --git a/clients/rust/crates/ahp/src/hosts/runtime.rs b/clients/rust/crates/ahp/src/hosts/runtime.rs index 93347a344..f30a83037 100644 --- a/clients/rust/crates/ahp/src/hosts/runtime.rs +++ b/clients/rust/crates/ahp/src/hosts/runtime.rs @@ -27,6 +27,11 @@ use super::types::{ /// Commands the runtime accepts from the [`MultiHostClient`]. pub(super) enum HostCommand { + SubscribeWindowed { + uri: String, + options: crate::WindowOptions, + reply: oneshot::Sender>, + }, Reconnect { reply: oneshot::Sender<()>, }, @@ -82,6 +87,7 @@ pub(super) fn spawn( resolved_client_id: String, fan_out: broadcast::Sender, host_events: broadcast::Sender, + flow_events: broadcast::Sender, ) -> HostHandleTx { let initial = HostInternal { id: config.id.clone(), @@ -102,6 +108,7 @@ pub(super) fn spawn( meta: None, }, subscriptions: config.initial_subscriptions.clone(), + owning_channels: std::collections::HashMap::new(), completion_trigger_characters: vec![], session_summaries: BTreeMap::new(), generation: 0, @@ -112,6 +119,9 @@ pub(super) fn spawn( let (cmd_tx, cmd_rx) = mpsc::channel(32); let runtime = HostRuntime { + managed: vec![], + flow_events, + flow_task: None, client_id: resolved_client_id, config, cmd_rx, @@ -130,7 +140,27 @@ pub(super) fn spawn( } } +fn prune_owned( + state: &mut HostInternal, + uri: &str, + owner: &std::sync::Weak, +) { + if state + .owning_channels + .get(uri) + .is_some_and(|current| current.ptr_eq(owner)) + { + state.owning_channels.remove(uri); + state + .subscriptions + .retain(|subscription| subscription != uri); + } +} + struct HostRuntime { + managed: Vec>, + flow_events: broadcast::Sender, + flow_task: Option>, config: HostConfig, client_id: String, cmd_rx: mpsc::Receiver, @@ -194,7 +224,8 @@ impl HostRuntime { attempt = 0; } let outcome = self.run_connection(events).await; - self.tear_down_client().await; + self.tear_down_client(matches!(outcome, InnerOutcome::Shutdown)) + .await; match outcome { InnerOutcome::Shutdown => break, InnerOutcome::ManualReconnect => { @@ -225,6 +256,9 @@ impl HostRuntime { // Decide whether to retry. if self.config.reconnect_policy.attempts_exhausted(attempt) { + for channel in &self.managed { + channel.finish(Some(crate::WindowError::new("host reconnection failed"))); + } let last_error = self.shared.lock().await.last_error.clone(); let error = last_error.unwrap_or_else(|| Arc::new(ClientError::Shutdown)); self.set_state( @@ -251,16 +285,44 @@ impl HostRuntime { break; } } + self.tear_down_client(true).await; } async fn connect_once(&mut self) -> Result { + if let Some(task) = self.flow_task.take() { + task.abort(); + } let transport = self .config .transport_factory .open_transport(self.config.id.clone()) .await?; - let client = Client::connect(transport, self.config.client_config.clone()).await?; + let client = Client::builder() + .config(self.config.client_config.clone()) + .windowing(self.config.transport_factory.windowing_config()) + .connect(transport) + .await?; + let mut flow = client.flow_control_events(); + let flow_events = self.flow_events.clone(); + let host_id = self.config.id.clone(); + self.flow_task = Some(tokio::spawn(async move { + loop { + match flow.recv().await { + Ok(event) => { + let _ = flow_events.send(super::HostFlowControlEvent { + host_id: host_id.clone(), + event, + }); + } + Err(broadcast::error::RecvError::Lagged(_)) => continue, + Err(broadcast::error::RecvError::Closed) => break, + } + } + })); + client + .restore_channels(&self.managed) + .map_err(|e| ClientError::Transport(crate::TransportError::Protocol(e.to_string())))?; // Attach the events receiver BEFORE the initialize/reconnect // handshake so any notifications the server pushes between the @@ -270,7 +332,23 @@ impl HostRuntime { // Decide between initialize and reconnect based on prior state. let (subscriptions, server_seq_after, init_result, reconnect_result) = { - let snapshot = self.shared.lock().await; + let mut snapshot = self.shared.lock().await; + let dropped: Vec<_> = snapshot + .owning_channels + .iter() + .filter(|(_, owner)| { + owner + .upgrade() + .map_or(true, |channel| channel.dropped() || channel.unowned_drop()) + }) + .map(|(uri, _)| uri.clone()) + .collect(); + for uri in dropped { + snapshot.owning_channels.remove(&uri); + snapshot + .subscriptions + .retain(|subscription| subscription != &uri); + } let can_reconnect = snapshot.server_seq > 0 && !snapshot.subscriptions.is_empty(); let subscriptions = snapshot.subscriptions.clone(); let server_seq = snapshot.server_seq; @@ -278,7 +356,21 @@ impl HostRuntime { if can_reconnect { match client - .reconnect(self.client_id.clone(), server_seq, subscriptions.clone()) + .request( + "reconnect", + ahp_types::commands::ReconnectParams { + channel: ROOT_RESOURCE_URI.into(), + meta: None, + client_id: self.client_id.clone(), + last_seen_server_seq: server_seq, + subscriptions: subscriptions + .iter() + .filter(|uri| !uri.starts_with("ahp-tcp:")) + .cloned() + .collect(), + windows: client.resume_windows(), + }, + ) .await { Ok(result) => (subscriptions, server_seq, None, Some(result)), @@ -310,6 +402,26 @@ impl HostRuntime { } }; + if init_result.is_some() { + for channel in &self.managed { + if channel.dropped() { + continue; + } + if let Err(error) = client + .subscribe_windowed(channel.uri.clone(), channel.options()) + .await + { + if !client.is_connected() { + return Err(ClientError::Transport(crate::TransportError::Protocol( + error.to_string(), + ))); + } + client.fail_subscription(channel, error); + let _ = client.unsubscribe(channel.uri.clone()).await; + } + } + } + // Refresh session summaries from `listSessions` — cheap on first // connect, kept in sync by notifications afterward. Failures are // non-fatal: we just leave the cache as-is and log. @@ -336,6 +448,7 @@ impl HostRuntime { state.server_seq = server_seq_after; } if let Some(init) = init_result.as_ref() { + state.server_seq = init.server_seq; if let Some(snapshot) = init .snapshots .iter() @@ -405,6 +518,29 @@ impl HostRuntime { result: ReconnectResult, prior_subscriptions: &[String], ) { + let kept: Vec = match &result { + ReconnectResult::Replay(result) => result.windows.as_ref(), + ReconnectResult::Snapshot(result) => result.windows.as_ref(), + } + .map(|windows| { + windows + .items + .iter() + .map(|item| item.channel.clone()) + .collect() + }) + .unwrap_or_default(); + let windowed: Vec = self + .managed + .iter() + .filter(|channel| channel.windowed()) + .map(|channel| channel.uri.clone()) + .collect(); + self.shared + .lock() + .await + .subscriptions + .retain(|uri| !windowed.contains(uri) || kept.contains(uri)); match result { ReconnectResult::Replay(replay) => { for envelope in replay.actions { @@ -425,7 +561,7 @@ impl HostRuntime { } ReconnectResult::Snapshot(snap) => { let mut state = self.shared.lock().await; - let mut surviving: Vec = Vec::with_capacity(snap.snapshots.len()); + let mut surviving = kept; for snapshot in snap.snapshots { if snapshot.from_seq > state.server_seq { state.server_seq = snapshot.from_seq; @@ -440,9 +576,11 @@ impl HostRuntime { // Drop subscriptions the server didn't return a snapshot // for — they're effectively `missing` even though the // snapshot arm doesn't carry an explicit list. - state - .subscriptions - .retain(|u| surviving.contains(u) || !prior_subscriptions.contains(u)); + state.subscriptions.retain(|u| { + surviving.contains(u) + || !prior_subscriptions.contains(u) + || u.starts_with("ahp-otlp:") + }); } } } @@ -464,15 +602,42 @@ impl HostRuntime { return InnerOutcome::ManualReconnect; } Some(HostCommand::Subscribe { uri, reply }) => { - let result = self.handle_subscribe(uri).await; + let shutdown = self.shutdown_signal.clone(); + let result = tokio::select! { + result = self.handle_subscribe(uri) => result, + _ = shutdown.notified() => return InnerOutcome::Shutdown, + }; + let _ = reply.send(result); + } + Some(HostCommand::SubscribeWindowed { uri, options, reply }) => { + let client = self.shared.lock().await.current_client.clone(); + let shutdown = self.shutdown_signal.clone(); + let result = match client { + Some(client) => tokio::select! { + result = client.subscribe_windowed(uri.clone(), options) => result.map_err(|error| HostError::Client(ClientError::Transport(crate::TransportError::Protocol(error.to_string())))), + _ = shutdown.notified() => return InnerOutcome::Shutdown, + }, + None => Err(HostError::HostShutDown(self.config.id.clone())), + }; + if let Ok((_, subscription)) = &result { + self.bind_managed(&subscription.channel).await; + } let _ = reply.send(result); } Some(HostCommand::Unsubscribe { uri, reply }) => { - let result = self.handle_unsubscribe(uri).await; + let shutdown = self.shutdown_signal.clone(); + let result = tokio::select! { + result = self.handle_unsubscribe(uri) => result, + _ = shutdown.notified() => return InnerOutcome::Shutdown, + }; let _ = reply.send(result); } Some(HostCommand::Dispatch { channel, action, reply }) => { - let result = self.handle_dispatch(channel, *action).await; + let shutdown = self.shutdown_signal.clone(); + let result = tokio::select! { + result = self.handle_dispatch(channel, *action) => result, + _ = shutdown.notified() => return InnerOutcome::Shutdown, + }; let _ = reply.send(result); } }, @@ -486,6 +651,9 @@ impl HostRuntime { _ = self.shutdown_signal.notified() => return false, cmd = self.cmd_rx.recv() => match cmd { None => return false, + Some(HostCommand::SubscribeWindowed { reply, .. }) => { + let _ = reply.send(Err(HostError::HostShutDown(self.config.id.clone()))); + } Some(HostCommand::Reconnect { reply }) => { let _ = reply.send(()); return true; @@ -529,6 +697,9 @@ impl HostRuntime { _ = &mut sleep => return true, cmd = self.cmd_rx.recv() => match cmd { None => return false, + Some(HostCommand::SubscribeWindowed { reply, .. }) => { + let _ = reply.send(Err(HostError::HostShutDown(self.config.id.clone()))); + } Some(HostCommand::Reconnect { reply }) => { let _ = reply.send(()); return true; @@ -583,6 +754,14 @@ impl HostRuntime { SubscriptionEvent::AuthRequired(_) => { // No cache update; consumers observe via the event stream. } + SubscriptionEvent::Snapshot(snapshot) => { + let mut state = self.shared.lock().await; + state.server_seq = state.server_seq.max(snapshot.from_seq); + if let SnapshotState::Root(root) = &snapshot.state { + state.root_state = root.as_ref().clone(); + } + } + _ => {} } let host_event = HostSubscriptionEvent { @@ -646,6 +825,37 @@ impl HostRuntime { Ok(result) } + async fn bind_managed(&self, channel: &Arc) { + let uri = channel.uri.clone(); + let owner = Arc::downgrade(channel); + { + let mut state = self.shared.lock().await; + if !state.subscriptions.contains(&uri) { + state.subscriptions.push(uri.clone()); + } + state.owning_channels.insert(uri.clone(), owner.clone()); + } + let shared = Arc::downgrade(&self.shared); + let hook = Arc::new(move || { + if let Some(shared) = shared.upgrade() { + if let Ok(mut state) = shared.try_lock() { + prune_owned(&mut state, &uri, &owner); + } else if let Ok(runtime) = tokio::runtime::Handle::try_current() { + let uri = uri.clone(); + let owner = owner.clone(); + let shared = shared.clone(); + runtime.spawn(async move { + let mut state = shared.lock().await; + prune_owned(&mut state, &uri, &owner); + }); + } + } + }); + let mut state = channel.lock(); + state.managed = !channel.uri.starts_with("ahp-tcp:"); + state.drop_hook = Some(hook); + } + async fn handle_unsubscribe(&self, uri: String) -> Result<(), HostError> { let client_opt = { self.shared.lock().await.current_client.clone() }; if let Some(client) = client_opt { @@ -656,6 +866,7 @@ impl HostRuntime { } let mut state = self.shared.lock().await; state.subscriptions.retain(|u| u != &uri); + state.owning_channels.remove(&uri); Ok(()) } @@ -677,14 +888,38 @@ impl HostRuntime { .map_err(HostError::Client) } - async fn tear_down_client(&self) { + async fn tear_down_client(&mut self, permanent: bool) { + if let Some(task) = self.flow_task.take() { + task.abort(); + } let prev = { let mut state = self.shared.lock().await; state.current_client.take() }; if let Some(client) = prev { + self.managed = client.managed_channels(); + if permanent { + for channel in &self.managed { + channel.finish(None); + } + } + self.shared + .lock() + .await + .subscriptions + .retain(|uri| !uri.starts_with("ahp-tcp:")); + self.shared + .lock() + .await + .owning_channels + .retain(|uri, _| !uri.starts_with("ahp-tcp:")); client.shutdown().await; } + if permanent { + for channel in &self.managed { + channel.finish(None); + } + } } async fn set_state(&self, state: HostState, last_error: Option>) { diff --git a/clients/rust/crates/ahp/src/hosts/types.rs b/clients/rust/crates/ahp/src/hosts/types.rs index ec47aaa12..fabab2e85 100644 --- a/clients/rust/crates/ahp/src/hosts/types.rs +++ b/clients/rust/crates/ahp/src/hosts/types.rs @@ -138,6 +138,14 @@ pub struct HostConfig { } impl HostConfig { + /// Set per-host window defaults without extending the public-field struct. + pub fn with_windowing(mut self, windowing: crate::WindowingConfig) -> Self { + self.transport_factory = Arc::new(super::factory::ConfiguredFactory { + inner: self.transport_factory, + windowing, + }); + self + } /// Build a [`HostConfig`] with sensible defaults. /// /// Leaves `client_id` as `None` so the multi-host client resolves @@ -404,6 +412,24 @@ pub struct HostSubscriptionEvent { pub event: SubscriptionEvent, } +/// Observe-only flow telemetry tagged with its host. +#[derive(Clone, Debug)] +#[non_exhaustive] +pub struct HostFlowControlEvent { + /// Originating host. + pub host_id: HostId, + /// Connection event. + pub event: crate::FlowControlEvent, +} +impl Default for HostFlowControlEvent { + fn default() -> Self { + Self { + host_id: HostId::new(""), + event: crate::FlowControlEvent::default(), + } + } +} + /// Stream of [`HostSubscriptionEvent`]s. /// /// Returned by [`super::MultiHostClient::events`]. Each call returns a @@ -533,6 +559,11 @@ impl HostShared { pub(super) async fn lock(&self) -> tokio::sync::MutexGuard<'_, HostInternal> { self.inner.lock().await } + pub(super) fn try_lock( + &self, + ) -> Result, tokio::sync::TryLockError> { + self.inner.try_lock() + } } /// Mutable per-host state. Updated by the runtime task; read on the @@ -550,6 +581,8 @@ pub(super) struct HostInternal { pub(super) automations: Option, pub(super) root_state: RootState, pub(super) subscriptions: Vec, + pub(super) owning_channels: + std::collections::HashMap>, pub(super) completion_trigger_characters: Vec, pub(super) session_summaries: std::collections::BTreeMap, pub(super) generation: u64, diff --git a/clients/rust/crates/ahp/src/lib.rs b/clients/rust/crates/ahp/src/lib.rs index 013229405..b253710c6 100644 --- a/clients/rust/crates/ahp/src/lib.rs +++ b/clients/rust/crates/ahp/src/lib.rs @@ -151,6 +151,12 @@ pub mod hosts; pub mod multi_host_state_mirror; pub mod reducers; pub mod transport; +pub mod windowing; +pub use windowing::{ + ChannelSender, ClientBuilder, Coalescer, CoalescingContext, Delivery, Direction, + FlowControlEvent, FlowControlStatus, WindowError, WindowOptions, WindowTransportCapabilities, + WindowedSubscription, WindowingConfig, +}; pub use ahp_types; diff --git a/clients/rust/crates/ahp/src/multi_host_state_mirror.rs b/clients/rust/crates/ahp/src/multi_host_state_mirror.rs index bcb99f5e5..8c2073743 100644 --- a/clients/rust/crates/ahp/src/multi_host_state_mirror.rs +++ b/clients/rust/crates/ahp/src/multi_host_state_mirror.rs @@ -171,11 +171,24 @@ impl MultiHostStateMirror { /// notifications and auth challenges) are ignored — they don't move any of /// the reducer-tracked state shapes. pub fn apply_event(&mut self, event: &HostSubscriptionEvent) { + if let SubscriptionEvent::Snapshot(snapshot) = &event.event { + self.apply_snapshot(&event.host_id, snapshot); + } + if let SubscriptionEvent::Action(envelope) = &event.event { self.apply_envelope(&event.host_id, envelope); } } + /// Apply a reliable owning delivery. Its caller retains release/Drop control. + pub fn apply_delivery(&mut self, host: &HostId, delivery: &crate::Delivery) { + match delivery.event() { + SubscriptionEvent::Snapshot(snapshot) => self.apply_snapshot(host, snapshot), + SubscriptionEvent::Action(envelope) => self.apply_envelope(host, envelope), + _ => {} + } + } + /// Apply a single action envelope, scoped to `host`. Routing uses /// `envelope.channel`: [`ROOT_RESOURCE_URI`] is the root channel, /// every other channel is identified by the URI the server diff --git a/clients/rust/crates/ahp/src/transport.rs b/clients/rust/crates/ahp/src/transport.rs index 3bf3a552e..1bb367469 100644 --- a/clients/rust/crates/ahp/src/transport.rs +++ b/clients/rust/crates/ahp/src/transport.rs @@ -112,6 +112,10 @@ impl TransportMessage { /// client sends indefinitely until the underlying connection closes, /// and `recv` signals closure by returning `None`. pub trait Transport: Send + 'static { + /// Exact-text windowing support; defaults to ordinary delivery. + fn window_capabilities(&self) -> crate::WindowTransportCapabilities { + crate::WindowTransportCapabilities::default() + } /// Send a single message. /// /// Errors returned here are typically fatal for the transport @@ -154,6 +158,10 @@ pub trait Transport: Send + 'static { /// allocation cost (typically: registries that hold one transport per /// host). pub trait DynTransport: Send + 'static { + /// Object-safe windowing hints. + fn window_capabilities(&self) -> crate::WindowTransportCapabilities { + crate::WindowTransportCapabilities::default() + } /// Object-safe analogue of [`Transport::send`]. fn send<'a>( &'a mut self, @@ -172,6 +180,9 @@ pub trait DynTransport: Send + 'static { } impl DynTransport for T { + fn window_capabilities(&self) -> crate::WindowTransportCapabilities { + Transport::window_capabilities(self) + } fn send<'a>( &'a mut self, msg: TransportMessage, @@ -247,6 +258,9 @@ impl std::fmt::Debug for BoxedTransport { } impl Transport for BoxedTransport { + fn window_capabilities(&self) -> crate::WindowTransportCapabilities { + self.inner.window_capabilities() + } fn send( &mut self, msg: TransportMessage, @@ -264,3 +278,37 @@ impl Transport for BoxedTransport { self.inner.close() } } + +/// Assert exact-text support for a custom transport. +/// +/// Incoming notifications must be unmodified Text/UTF-8 Binary and outgoing +/// Text must be sent unchanged. Parsed windowed data is a protocol error. +pub struct RawTextTransport { + inner: T, + maximum_frame_bytes: Option, +} +impl RawTextTransport { + /// Wrap a transport with an optional safe UTF-16 receive limit. + pub fn new(inner: T, maximum_frame_bytes: Option) -> Self { + Self { + inner, + maximum_frame_bytes, + } + } +} +impl Transport for RawTextTransport { + fn window_capabilities(&self) -> crate::WindowTransportCapabilities { + let mut capabilities = crate::WindowTransportCapabilities::default().raw_text(true); + capabilities.maximum_frame_bytes = self.maximum_frame_bytes; + capabilities + } + async fn send(&mut self, message: TransportMessage) -> Result<(), TransportError> { + self.inner.send(message).await + } + async fn recv(&mut self) -> Result, TransportError> { + self.inner.recv().await + } + async fn close(&mut self) -> Result<(), TransportError> { + self.inner.close().await + } +} diff --git a/clients/rust/crates/ahp/src/windowing.rs b/clients/rust/crates/ahp/src/windowing.rs new file mode 100644 index 000000000..b25a8a44f --- /dev/null +++ b/clients/rust/crates/ahp/src/windowing.rs @@ -0,0 +1,2200 @@ +//! Opt-in, consumer-driven subscription flow control. +use std::collections::{HashMap, VecDeque}; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; + +use ahp_types::commands::{ReconnectResult, SubscribeResult}; +use ahp_types::messages::JsonRpcNotification; +use ahp_types::notifications::counts_against_window; +use ahp_types::state::{ChannelFlowControl, ChannelReceiveLimits}; +use serde::Serialize; +use serde_json::{json, Value}; +use tokio::sync::{broadcast, oneshot, watch, Notify}; + +use crate::{Client, ClientConfig, ClientError, DispatchHandle, SubscriptionEvent, Transport}; + +const SAFE_INTEGER: u64 = (1 << 53) - 1; +pub(crate) const MAX_CONTROL_BYTES: u64 = 64 * 1024; +#[cfg(test)] +#[path = "windowing_tests.rs"] +mod tests; + +/// Measure complete, unmodified JSON text in protocol UTF-16 bytes. +pub fn utf16_bytes(text: &str) -> u64 { + text.encode_utf16().count() as u64 * 2 +} + +/// Additive transport hints. Parsed-only transports retain ordinary delivery. +#[derive(Clone, Debug, Default)] +#[non_exhaustive] +pub struct WindowTransportCapabilities { + /// Whether incoming text and outgoing text preserve the actual wire representation. + pub raw_text: bool, + /// Safe receive limit in protocol UTF-16 bytes. + pub maximum_frame_bytes: Option, +} +impl WindowTransportCapabilities { + /// Assert exact-text support. + pub fn raw_text(mut self, value: bool) -> Self { + self.raw_text = value; + self + } + /// Set the safe protocol frame limit. + pub fn maximum_frame_bytes(mut self, value: u64) -> Self { + self.maximum_frame_bytes = Some(value); + self + } +} + +/// Coalescing direction. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +#[non_exhaustive] +pub enum Direction { + /// Incoming consumer queue. + #[default] + Receive, + /// Outgoing unsent queue. + Send, +} +/// Context passed to custom coalescers. +#[derive(Clone, Debug, Default)] +#[non_exhaustive] +pub struct CoalescingContext { + /// Subscription URI. + pub channel: String, + /// Queue direction. + pub direction: Direction, +} +/// A connection-wide custom notification coalescer. +pub trait Coalescer: Send + Sync { + /// Return a merged notification, or decline. Original inputs remain unchanged. + fn merge( + &self, + earlier: &JsonRpcNotification, + later: &JsonRpcNotification, + context: &CoalescingContext, + ) -> Option; +} +impl Coalescer for F +where + F: Fn( + &JsonRpcNotification, + &JsonRpcNotification, + &CoalescingContext, + ) -> Option + + Send + + Sync, +{ + fn merge( + &self, + earlier: &JsonRpcNotification, + later: &JsonRpcNotification, + context: &CoalescingContext, + ) -> Option { + self(earlier, later, context) + } +} +/// Connection defaults. Merely configuring these does not enable windowing. +#[derive(Clone)] +#[non_exhaustive] +pub struct WindowingConfig { + pub(crate) window_bytes: u64, + pub(crate) maximum_frame_bytes: Option, + pub(crate) max_buffered_send_bytes: u64, + pub(crate) coalescers: Vec>, +} +impl Default for WindowingConfig { + fn default() -> Self { + Self { + window_bytes: 262_144, + maximum_frame_bytes: None, + max_buffered_send_bytes: 1_048_576, + coalescers: vec![], + } + } +} +impl WindowingConfig { + /// Default receive window. + pub fn window_bytes(mut self, value: u64) -> Self { + self.window_bytes = value; + self + } + /// Override transport-derived maximum frame size. + pub fn maximum_frame_bytes(mut self, value: u64) -> Self { + self.maximum_frame_bytes = Some(value); + self + } + /// Bound queued, unsent data per subscription. + pub fn max_buffered_send_bytes(mut self, value: u64) -> Self { + self.max_buffered_send_bytes = value; + self + } + /// Register a custom coalescer before built-ins. + pub fn coalescer(mut self, value: impl Coalescer + 'static) -> Self { + self.coalescers.push(Arc::new(value)); + self + } + /// Replace the custom coalescer list. + pub fn coalescers(mut self, values: Vec>) -> Self { + self.coalescers = values; + self + } +} +/// Overrides for one subscription. Holding deliveries is manual receive control. +#[derive(Clone, Debug)] +#[non_exhaustive] +pub struct WindowOptions { + pub(crate) window_bytes: Option, + pub(crate) maximum_frame_bytes: Option, + pub(crate) coalescing: bool, +} +impl Default for WindowOptions { + fn default() -> Self { + Self { + window_bytes: None, + maximum_frame_bytes: None, + coalescing: true, + } + } +} +impl WindowOptions { + /// Override the connection's receive window. + pub fn window_bytes(mut self, value: u64) -> Self { + self.window_bytes = Some(value); + self + } + /// Override the connection's frame limit. + pub fn maximum_frame_bytes(mut self, value: u64) -> Self { + self.maximum_frame_bytes = Some(value); + self + } + /// Enable or disable local queued-item coalescing. + pub fn coalescing(mut self, value: bool) -> Self { + self.coalescing = value; + self + } +} +/// Inspectable error for the additive windowing API. +#[derive(Clone, Debug, Default, thiserror::Error)] +#[error("{message}")] +#[non_exhaustive] +pub struct WindowError { + /// Human-readable reason. + pub message: String, +} +impl WindowError { + /// Construct an error. + pub fn message(mut self, message: impl Into) -> Self { + self.message = message.into(); + self + } + pub(crate) fn new(message: impl Into) -> Self { + Self::default().message(message) + } +} +impl From for WindowError { + fn from(value: ClientError) -> Self { + Self::new(value.to_string()) + } +} +/// Current subscription flow-control status. Amounts are UTF-16 bytes. +#[derive(Clone, Debug, PartialEq, Serialize)] +#[serde(rename_all = "camelCase")] +#[non_exhaustive] +pub struct FlowControlStatus { + /// `ordinary` or `windowed`. + pub mode: String, + /// `transport` or `declined`, when an opt-in falls back. + #[serde(skip_serializing_if = "Option::is_none")] + pub fallback_reason: Option, + /// Accepted limits. + #[serde(skip_serializing_if = "Option::is_none")] + pub limits: Option, + /// Sent data awaiting host credit. + pub outstanding_send_bytes: u64, + /// Queued, unsent data. + pub buffered_send_bytes: u64, + /// Received data awaiting consumer release. + pub buffered_receive_bytes: u64, + /// Whether window and queue currently permit progress. + pub writable: bool, +} +impl Default for FlowControlStatus { + fn default() -> Self { + Self { + mode: "ordinary".into(), + fallback_reason: None, + limits: None, + outstanding_send_bytes: 0, + buffered_send_bytes: 0, + buffered_receive_bytes: 0, + writable: true, + } + } +} +impl FlowControlStatus { + /// Set the delivery mode. + pub fn mode(mut self, value: impl Into) -> Self { + self.mode = value.into(); + self + } + /// Set a fallback reason. + pub fn fallback_reason(mut self, value: Option) -> Self { + self.fallback_reason = value; + self + } + /// Set accepted limits. + pub fn limits(mut self, value: Option) -> Self { + self.limits = value; + self + } + /// Set outstanding sent bytes. + pub fn outstanding_send_bytes(mut self, value: u64) -> Self { + self.outstanding_send_bytes = value; + self + } + /// Set queued sent bytes. + pub fn buffered_send_bytes(mut self, value: u64) -> Self { + self.buffered_send_bytes = value; + self + } + /// Set unreleased received bytes. + pub fn buffered_receive_bytes(mut self, value: u64) -> Self { + self.buffered_receive_bytes = value; + self + } + /// Set writability. + pub fn writable(mut self, value: bool) -> Self { + self.writable = value; + self + } +} +/// Observe-only connection flow-control event. +#[derive(Clone, Debug, Default, Serialize)] +#[non_exhaustive] +pub struct FlowControlEvent { + /// Subscription URI. + pub channel: String, + /// negotiated, declined, blocked, unblocked, creditSent, creditReceived, + /// fragmentedMessageReceived, fragmentedMessageSent, coalesced, resumed, dropped. + pub kind: String, + /// Status at the event. + pub status: FlowControlStatus, +} +impl FlowControlEvent { + /// Set the subscription URI. + pub fn channel(mut self, value: impl Into) -> Self { + self.channel = value.into(); + self + } + /// Set the event kind. + pub fn kind(mut self, value: impl Into) -> Self { + self.kind = value.into(); + self + } + /// Set the associated status. + pub fn status(mut self, value: FlowControlStatus) -> Self { + self.status = value; + self + } +} +/// Non-breaking client configuration builder. +#[derive(Default)] +#[non_exhaustive] +pub struct ClientBuilder { + config: ClientConfig, + windowing: WindowingConfig, +} +impl ClientBuilder { + /// Configure the existing request timeout. + pub fn default_request_timeout(mut self, value: Option) -> Self { + self.config.default_request_timeout = value; + self + } + /// Configure existing ordinary broadcast buffering. + pub fn subscription_buffer(mut self, value: usize) -> Self { + self.config.subscription_buffer = value; + self + } + /// Replace existing client options. + pub fn config(mut self, value: ClientConfig) -> Self { + self.config = value; + self + } + /// Set opt-in window defaults and coalescers. + pub fn windowing(mut self, value: WindowingConfig) -> Self { + self.windowing = value; + self + } + /// Connect through a pluggable transport. + pub async fn connect(self, transport: T) -> Result { + Client::connect_configured(transport, self.config, self.windowing).await + } +} +/// A notification and its exact accounting receipt. +pub struct Delivery { + event: SubscriptionEvent, + notification: JsonRpcNotification, + bytes: u64, + release_bytes: u64, + epoch: u64, + channel: Arc, +} +impl Delivery { + /// Typed consumer event. + pub fn event(&self) -> &SubscriptionEvent { + &self.event + } + /// Original or locally coalesced notification. + pub fn notification(&self) -> &JsonRpcNotification { + &self.notification + } + /// Sum of original message/fragment wire sizes. + pub fn wire_bytes(&self) -> u64 { + self.bytes + } + /// Release credit; subsequent release or Drop is a no-op. + pub fn release(&mut self) { + let bytes = std::mem::take(&mut self.release_bytes); + if bytes != 0 { + let mut state = self.channel.lock(); + if state.epoch == self.epoch && !state.closed && !state.dropped { + state.released += bytes; + self.channel.publish(&state); + self.channel.wake.notify_one(); + } + } + } +} +impl Drop for Delivery { + fn drop(&mut self) { + self.release(); + } +} +/// Single owning, reliable subscription queue (ordinary fallback remains lossy). +pub struct WindowedSubscription { + pub(crate) channel: Arc, + client: Client, + pub(crate) host_id: Option, +} +impl WindowedSubscription { + pub(crate) fn new(channel: Arc, client: Client) -> Self { + channel.owners.fetch_add(1, Ordering::Relaxed); + channel.drop_requested.store(false, Ordering::Release); + Self { + channel, + client, + host_id: None, + } + } + /// Channel URI. + pub fn uri(&self) -> &str { + &self.channel.uri + } + /// Host of a managed multi-host subscription. + pub fn host_id(&self) -> Option<&crate::hosts::HostId> { + self.host_id.as_ref() + } + /// Take the next delivery. Dropping it releases its credit. + pub async fn next(&mut self) -> Result, WindowError> { + loop { + let notified = self.channel.ready.notified(); + let item = { + let mut state = self.channel.lock(); + if let Some(item) = state.received.pop_front() { + state.initial_delivered = true; + Some(item) + } else if let Some(error) = &state.error { + if state.managed && !state.dropped { + None + } else { + return Err(error.clone()); + } + } else if state.closed { + if state.managed && !state.dropped { + None + } else { + return Ok(None); + } + } else { + None + } + }; + if let Some(item) = item { + let event = crate::client::decode_notification(item.notification.clone()) + .map_err(WindowError::new)?; + return Ok(Some(Delivery { + event, + notification: item.notification, + bytes: item.bytes, + release_bytes: item.release_bytes, + epoch: item.epoch, + channel: self.channel.clone(), + })); + } + notified.await; + } + } + /// A flow-controlled sender for this channel. + pub fn sender(&self) -> ChannelSender { + ChannelSender { + channel: self.channel.clone(), + } + } + /// Current flow-control status. + pub fn status(&self) -> FlowControlStatus { + self.channel.status.borrow().clone() + } + /// A watch stream of status changes. + pub fn status_changes(&self) -> watch::Receiver { + self.channel.status.subscribe() + } + /// Wait until the send window and queue have room. + pub async fn wait_writable(&self) -> Result<(), WindowError> { + self.sender().wait_writable().await + } +} +impl Drop for WindowedSubscription { + fn drop(&mut self) { + if self.channel.owners.fetch_sub(1, Ordering::AcqRel) == 1 { + self.channel.drop_requested.store(true, Ordering::Release); + self.channel.schedule_drop(self.client.clone()); + } + } +} +/// Flow-controlled notification sender. +#[derive(Clone)] +pub struct ChannelSender { + channel: Arc, +} +impl ChannelSender { + fn current_client(&self) -> Result { + self.channel + .binding + .lock() + .unwrap_or_else(|e| e.into_inner()) + .as_ref() + .and_then(crate::client::ClientBinding::upgrade) + .ok_or_else(|| WindowError::new("connection closed")) + } + /// Current status. + pub fn status(&self) -> FlowControlStatus { + self.channel.status.borrow().clone() + } + /// Observe status changes. + pub fn status_changes(&self) -> watch::Receiver { + self.channel.status.subscribe() + } + /// Advisory wait, not a send reservation. + pub async fn wait_writable(&self) -> Result<(), WindowError> { + let mut changes = self.status_changes(); + loop { + { + let state = self.channel.lock(); + if let Some(error) = &state.error { + if !state.managed || state.dropped { + return Err(error.clone()); + } + } + if state.closed && (!state.managed || state.dropped) { + return Err(WindowError::new("subscription closed")); + } + } + if changes.borrow().writable { + return Ok(()); + } + changes + .changed() + .await + .map_err(|_| WindowError::new("subscription closed"))?; + } + } + /// Send and await the actual transport write, not host processing. + pub async fn send(&self, method: &str, params: P) -> Result<(), WindowError> { + if self.channel.lock().closed { + return Err(WindowError::new("subscription connection closed")); + } + let mut params = + serde_json::to_value(params).map_err(|e| WindowError::new(e.to_string()))?; + let object = params + .as_object_mut() + .ok_or_else(|| WindowError::new("notification params must be an object"))?; + if object + .get("channel") + .is_some_and(|value| value != &self.channel.uri) + { + return Err(WindowError::new("notification channel mismatch")); + } + object.insert("channel".into(), self.channel.uri.clone().into()); + self.current_client()?.window_send(method, params).await + } + /// Dispatch an individually sequenced action and await transport write. + pub async fn dispatch( + &self, + action: ahp_types::actions::StateAction, + ) -> Result { + if self.channel.lock().closed { + return Err(WindowError::new("subscription connection closed")); + } + self.current_client()? + .window_dispatch(self.channel.uri.clone(), action) + .await + } +} + +struct Received { + notification: JsonRpcNotification, + bytes: u64, + release_bytes: u64, + epoch: u64, +} +struct Queued { + notification: JsonRpcNotification, + text: String, + offset: usize, + bytes: u64, + parts: Vec, +} +struct QueuedPart { + notification: JsonRpcNotification, + completion: Option>>, +} +pub(crate) struct ChannelState { + options: WindowOptions, + offer: Option, + limits: Option, + fallback: Option, + epoch: u64, + sent: u64, + consumed: u64, + received_bytes: u64, + released: u64, + advertised: u64, + credit_since: Option, + received: VecDeque, + outgoing: VecDeque, + queued_bytes: u64, + fragment: String, + fragment_bytes: u64, + snapshot_expected: bool, + pub(crate) retained_seq: Option, + pub(crate) closed: bool, + error: Option, + dropped: bool, + pending_options: Option, + accepted: bool, + initial_delivered: bool, + pub(crate) managed: bool, + dropped_notified: bool, + pub(crate) drop_hook: Option>, +} +pub(crate) struct Channel { + pub(crate) uri: String, + state: Mutex, + ready: Notify, + wake: Arc, + status: watch::Sender, + events: Mutex>, + max_send: u64, + ordinary_buffer: usize, + pub(crate) binding: Mutex>, + owners: AtomicUsize, + acquisitions: AtomicUsize, + drop_requested: AtomicBool, +} +impl Channel { + pub(crate) fn begin_acquisition(&self) { + self.acquisitions.fetch_add(1, Ordering::AcqRel); + } + pub(crate) fn end_acquisition(self: &Arc, client: Client) { + if self.acquisitions.fetch_sub(1, Ordering::AcqRel) == 1 + && self.drop_requested.load(Ordering::Acquire) + { + self.schedule_drop(client); + } + } + fn schedule_drop(self: &Arc, fallback: Client) { + if let Ok(runtime) = tokio::runtime::Handle::try_current() { + let channel = self.clone(); + runtime.spawn(async move { + let _ = fallback + .unsubscribe_if_current(&channel.uri, &channel) + .await; + }); + } + } + pub(crate) fn lock(&self) -> std::sync::MutexGuard<'_, ChannelState> { + self.state.lock().unwrap_or_else(|e| e.into_inner()) + } + pub(crate) fn options(&self) -> WindowOptions { + self.lock().options.clone() + } + pub(crate) fn windowed(&self) -> bool { + self.lock().limits.is_some() + } + pub(crate) fn dropped(&self) -> bool { + self.lock().dropped + } + pub(crate) fn unowned_drop(&self) -> bool { + self.drop_requested.load(Ordering::Acquire) + && self.owners.load(Ordering::Acquire) == 0 + && self.acquisitions.load(Ordering::Acquire) == 0 + } + pub(crate) fn finish(&self, error: Option) { + let mut state = self.lock(); + state.managed = false; + state.dropped = true; + state.closed = true; + state.error = error; + state.credit_since = None; + self.publish(&state); + self.status.send_modify(|_| {}); + self.ready.notify_waiters(); + if !state.dropped_notified { + state.dropped_notified = true; + self.event("dropped"); + } + let hook = state.drop_hook.take(); + drop(state); + if let Some(hook) = hook { + hook(); + } + } + fn publish(&self, state: &ChannelState) { + let status = FlowControlStatus { + mode: if state.limits.is_some() { + "windowed" + } else { + "ordinary" + } + .into(), + fallback_reason: state.fallback.clone(), + limits: state.limits.clone(), + outstanding_send_bytes: state.sent - state.consumed, + buffered_send_bytes: state.queued_bytes, + buffered_receive_bytes: state.received_bytes - state.released, + writable: state.accepted + && !state.closed + && state + .queued_bytes + .saturating_sub(state.outgoing.front().map_or(0, |item| item.bytes)) + < self.max_send + && state.limits.as_ref().map_or(true, |limits| { + state.sent - state.consumed < limits.host_receive.window_bytes as u64 + }), + }; + let previous = self.status.borrow().clone(); + let was_writable = previous.writable; + if previous != status { + self.status.send_replace(status.clone()); + } + if was_writable != status.writable { + self.event(if status.writable { + "unblocked" + } else { + "blocked" + }); + } + } + fn event(&self, kind: &str) { + let event = FlowControlEvent::default() + .channel(self.uri.clone()) + .kind(kind) + .status(self.status.borrow().clone()); + tracing::debug!(channel = %event.channel, kind, status = ?event.status, "flow control"); + let _ = self + .events + .lock() + .unwrap_or_else(|e| e.into_inner()) + .send(event); + } +} +enum PendingKind { + Subscribe { + uri: String, + channel: std::sync::Weak, + offer: Option, + options: Option, + }, + Reconnect { + windows: Vec, + ordinary: ChannelScope, + }, + Initialize(ChannelScope), +} +struct ChannelScope { + channels: Vec>, + uris: Vec, +} +struct ResumeOffer { + uri: String, + channel: std::sync::Weak, + offer: Option, +} +struct PendingOperation { + kind: PendingKind, + abandoned: bool, + sent: bool, +} +struct EngineState { + channels: HashMap>, + pending: HashMap, + cursor: usize, +} +pub(crate) struct Engine { + state: Mutex, + pub(crate) config: WindowingConfig, + capabilities: WindowTransportCapabilities, + pub(crate) wake: Arc, + pub(crate) events: broadcast::Sender, + ordinary_buffer: usize, +} +pub(crate) struct WireSend { + pub(crate) text: String, + pub(crate) control: bool, + completions: Vec>>, + event: Option<(Arc, &'static str)>, + channel: Arc, + epoch: u64, +} +impl WireSend { + pub(crate) fn is_live(&self) -> bool { + let state = self.channel.lock(); + !state.closed && state.epoch == self.epoch + } + pub(crate) fn complete(self, result: Result<(), WindowError>) { + if result.is_ok() { + if let Some((channel, kind)) = self.event { + channel.event(kind); + } + } + for completion in self.completions { + let _ = completion.send(result.clone()); + } + } +} +impl Engine { + pub(crate) fn new( + config: WindowingConfig, + capabilities: WindowTransportCapabilities, + ordinary_buffer: usize, + ) -> Self { + let (events, _) = broadcast::channel(256); + Self { + state: Mutex::new(EngineState { + channels: HashMap::new(), + pending: HashMap::new(), + cursor: 0, + }), + config, + capabilities, + wake: Arc::new(Notify::new()), + events, + ordinary_buffer, + } + } + fn lock(&self) -> std::sync::MutexGuard<'_, EngineState> { + self.state.lock().unwrap_or_else(|e| e.into_inner()) + } + pub(crate) fn subscription(&self, uri: &str) -> Option> { + self.lock().channels.get(uri).cloned() + } + pub(crate) fn has_subscriptions(&self) -> bool { + !self.lock().channels.is_empty() + } + pub(crate) fn has_windowed(&self) -> bool { + self.lock() + .channels + .values() + .any(|channel| channel.windowed()) + } + fn capture_scope(&self) -> ChannelScope { + let channels = self.managed_channels(); + ChannelScope { + uris: channels.iter().map(|channel| channel.uri.clone()).collect(), + channels: channels.iter().map(Arc::downgrade).collect(), + } + } + pub(crate) fn prepare( + &self, + uri: &str, + options: WindowOptions, + ) -> Result, WindowError> { + if self.lock().pending.len() >= 256 { + return Err(WindowError::new("too many pending window negotiations")); + } + if utf16_bytes(&json!({"jsonrpc":"2.0","method":"channel/credit","params":{"channel":uri,"consumedBytes":SAFE_INTEGER}}).to_string()) > MAX_CONTROL_BYTES { + return Err(WindowError::new("channel URI exceeds control message limit")); + } + let offer = ChannelReceiveLimits { + window_bytes: options.window_bytes.unwrap_or(self.config.window_bytes) as i64, + maximum_frame_bytes: options + .maximum_frame_bytes + .or(self.config.maximum_frame_bytes) + .or(self.capabilities.maximum_frame_bytes) + .unwrap_or(1_048_576) as i64, + }; + validate_limits(&offer)?; + if self.config.max_buffered_send_bytes == 0 { + return Err(WindowError::new("send buffer must be positive")); + } + let mut engine = self.lock(); + let channel = engine + .channels + .entry(uri.into()) + .or_insert_with(|| { + let (status, _) = watch::channel(FlowControlStatus::default()); + Arc::new(Channel { + uri: uri.into(), + state: Mutex::new(ChannelState { + options: options.clone(), + offer: None, + limits: None, + fallback: None, + epoch: 0, + sent: 0, + consumed: 0, + received_bytes: 0, + released: 0, + advertised: 0, + credit_since: None, + received: VecDeque::new(), + outgoing: VecDeque::new(), + queued_bytes: 0, + fragment: String::new(), + fragment_bytes: 0, + snapshot_expected: false, + retained_seq: None, + closed: false, + error: None, + dropped: false, + pending_options: None, + accepted: false, + initial_delivered: false, + managed: false, + dropped_notified: false, + drop_hook: None, + }), + ready: Notify::new(), + wake: self.wake.clone(), + status, + events: Mutex::new(self.events.clone()), + max_send: self.config.max_buffered_send_bytes, + ordinary_buffer: self.ordinary_buffer, + binding: Mutex::new(None), + owners: AtomicUsize::new(0), + acquisitions: AtomicUsize::new(0), + drop_requested: AtomicBool::new(false), + }) + }) + .clone(); + let mut state = channel.lock(); + state.pending_options = Some(options); + if state.limits.is_none() { + state.fallback = (!self.capabilities.raw_text).then(|| "transport".into()); + } + let offer = self.capabilities.raw_text.then_some(offer); + channel.publish(&state); + Ok(offer) + } + pub(crate) fn register_request( + &self, + id: u64, + method: &str, + params: &mut Value, + ) -> Result<(), WindowError> { + if matches!(method, "subscribe" | "reconnect" | "initialize") + && self.lock().pending.len() >= 256 + { + return Err(WindowError::new("too many pending window negotiations")); + } + if method == "subscribe" { + if !self.capabilities.raw_text && params.get("flowControl").is_some() { + let uri = params["channel"] + .as_str() + .ok_or_else(|| WindowError::new("channel missing"))?; + self.prepare(uri, WindowOptions::default())?; + if let Some(object) = params.as_object_mut() { + object.remove("flowControl"); + } + } + let uri = params["channel"] + .as_str() + .ok_or_else(|| WindowError::new("subscribe channel missing"))?; + if uri.starts_with("ahp-tcp:") && params.get("flowControl").is_none() { + return Err(WindowError::new("TCP requires windowing")); + } + let offer: Option = params + .pointer("/flowControl/receive") + .map(|v| { + serde_json::from_value(v.clone()).map_err(|e| WindowError::new(e.to_string())) + }) + .transpose()?; + if let Some(limits) = &offer { + validate_limits(limits)?; + if self.subscription(uri).is_none() { + self.prepare( + uri, + WindowOptions::default() + .window_bytes(limits.window_bytes as u64) + .maximum_frame_bytes(limits.maximum_frame_bytes as u64), + )?; + } + } + if let Some(channel) = self.subscription(uri) { + channel.lock().pending_options = None; + let options = Some(channel.options()); + self.lock().pending.insert( + id, + PendingOperation { + kind: PendingKind::Subscribe { + uri: uri.into(), + channel: Arc::downgrade(&channel), + offer, + options, + }, + abandoned: false, + sent: false, + }, + ); + } + } else if method == "initialize" { + if params["initialSubscriptions"] + .as_array() + .is_some_and(|channels| { + channels + .iter() + .any(|uri| uri.as_str().is_some_and(|uri| uri.starts_with("ahp-tcp:"))) + }) + { + return Err(WindowError::new("TCP cannot be an initial subscription")); + } + if !self.lock().channels.is_empty() { + let channels = self.capture_scope(); + self.lock().pending.insert( + id, + PendingOperation { + kind: PendingKind::Initialize(channels), + abandoned: false, + sent: false, + }, + ); + } + } else if method == "reconnect" { + if params["subscriptions"].as_array().is_some_and(|channels| { + channels + .iter() + .any(|uri| uri.as_str().is_some_and(|uri| uri.starts_with("ahp-tcp:"))) + }) { + return Err(WindowError::new("TCP cannot be resumed")); + } + if !self.capabilities.raw_text { + if let Some(items) = params.pointer("/windows/items").and_then(Value::as_array) { + for item in items { + let uri = item["channel"] + .as_str() + .ok_or_else(|| WindowError::new("resume channel missing"))?; + self.prepare(uri, WindowOptions::default())?; + } + } + if let Some(object) = params.as_object_mut() { + object.remove("windows"); + } + } + if let Some(items) = params.pointer("/windows/items").and_then(Value::as_array) { + let mut channels: Vec = Vec::new(); + for item in items { + let uri = item["channel"] + .as_str() + .ok_or_else(|| WindowError::new("resume channel missing"))?; + if uri.starts_with("ahp-tcp:") || channels.iter().any(|c| c.uri == uri) { + return Err(WindowError::new("invalid resume channel")); + } + if !params["subscriptions"] + .as_array() + .is_some_and(|subscriptions| subscriptions.iter().any(|value| value == uri)) + { + return Err(WindowError::new("resume channel is not subscribed")); + } + let limits: ChannelReceiveLimits = + serde_json::from_value(item["flowControl"]["receive"].clone()) + .map_err(|e| WindowError::new(e.to_string()))?; + let options = self + .subscription(uri) + .map(|channel| channel.lock().options.clone()) + .unwrap_or_default(); + let offer = self.prepare( + uri, + options + .window_bytes(limits.window_bytes as u64) + .maximum_frame_bytes(limits.maximum_frame_bytes as u64), + )?; + if let Some(channel) = self.subscription(uri) { + let mut state = channel.lock(); + state.offer = offer; + if let Some(sequence) = item["lastAppliedServerSeq"].as_i64() { + if sequence < 0 || sequence as u64 > SAFE_INTEGER { + return Err(WindowError::new("invalid retained sequence")); + } + state.retained_seq = + Some(state.retained_seq.unwrap_or(0).max(sequence)); + } + } + let channel = self + .subscription(uri) + .ok_or_else(|| WindowError::new("subscription cancelled"))?; + channels.push(ResumeOffer { + uri: uri.into(), + offer: channel.lock().offer.clone(), + channel: Arc::downgrade(&channel), + }); + } + let ordinary = self.capture_scope(); + self.lock().pending.insert( + id, + PendingOperation { + kind: PendingKind::Reconnect { + windows: channels, + ordinary, + }, + abandoned: false, + sent: false, + }, + ); + } else if self.has_subscriptions() { + let ordinary = self.capture_scope(); + self.lock().pending.insert( + id, + PendingOperation { + kind: PendingKind::Reconnect { + windows: vec![], + ordinary, + }, + abandoned: false, + sent: false, + }, + ); + } + } + Ok(()) + } + pub(crate) fn register_owned( + &self, + id: u64, + channel: &Arc, + options: WindowOptions, + params: &Value, + ) -> Result<(), WindowError> { + let offer = params + .pointer("/flowControl/receive") + .map(|value| { + serde_json::from_value(value.clone()) + .map_err(|error| WindowError::new(error.to_string())) + }) + .transpose()?; + let mut engine = self.lock(); + if engine.pending.len() >= 256 { + return Err(WindowError::new("too many pending window negotiations")); + } + if !engine + .channels + .get(&channel.uri) + .is_some_and(|current| Arc::ptr_eq(current, channel)) + { + return Err(WindowError::new("subscription cancelled")); + } + channel.lock().pending_options = None; + engine.pending.insert( + id, + PendingOperation { + kind: PendingKind::Subscribe { + uri: channel.uri.clone(), + channel: Arc::downgrade(channel), + offer, + options: Some(options), + }, + abandoned: false, + sent: false, + }, + ); + Ok(()) + } + pub(crate) fn response(&self, id: u64, result: &Value) -> Result, WindowError> { + let mut cleanup = vec![]; + let operation = self + .lock() + .pending + .remove(&id) + .map(|operation| operation.kind); + match operation { + Some(PendingKind::Subscribe { + uri, + channel, + offer, + options, + }) => { + let current = self.subscription(&uri); + let Some(channel) = channel.upgrade().filter(|channel| { + current + .as_ref() + .is_some_and(|current| Arc::ptr_eq(current, channel)) + }) else { + return Ok(current.is_none().then_some(uri).into_iter().collect()); + }; + let result: SubscribeResult = serde_json::from_value(result.clone()) + .map_err(|e| WindowError::new(e.to_string()))?; + channel.lock().pending_options = options; + if !self.install(&channel, offer, result.flow_control, true)? { + return Ok(self + .subscription(&uri) + .is_none() + .then_some(uri) + .into_iter() + .collect()); + } + if let Some(snapshot) = result.snapshot { + if channel.lock().limits.is_some() { + return Err(WindowError::new( + "accepted subscription has result snapshot", + )); + } + self.receive_on(&channel, json!({"jsonrpc":"2.0","method":"channel/snapshot","params":{"channel":uri,"snapshot":snapshot}}), 0, 0)?; + } + } + Some(PendingKind::Reconnect { windows, ordinary }) => { + let channels: Vec<_> = windows.iter().map(|window| window.uri.clone()).collect(); + let result: ReconnectResult = serde_json::from_value(result.clone()) + .map_err(|e| WindowError::new(e.to_string()))?; + let resumed = match &result { + ReconnectResult::Replay(result) => result.windows.clone(), + ReconnectResult::Snapshot(result) => result.windows.clone(), + }; + let items = resumed.map(|w| w.items).unwrap_or_default(); + for channel in ordinary + .channels + .iter() + .filter_map(|channel| channel.upgrade()) + .filter(|channel| self.is_current(channel)) + { + if !channels.contains(&channel.uri) { + self.install(&channel, None, None, false)?; + } + } + let mut seen = std::collections::HashSet::new(); + for item in &items { + if !channels.contains(&item.channel) || !seen.insert(item.channel.clone()) { + return Err(WindowError::new("unsolicited or duplicate resumed channel")); + } + } + match &result { + ReconnectResult::Replay(result) => { + for action in &result.actions { + if channels.contains(&action.channel) { + return Err(WindowError::new( + "windowed replay must arrive as data", + )); + } + if ordinary + .channels + .iter() + .filter_map(|channel| channel.upgrade()) + .any(|channel| { + channel.uri == action.channel && self.is_current(&channel) + }) + && !channels.contains(&action.channel) + { + self.receive( + json!({"jsonrpc":"2.0","method":"action","params":action}), + 0, + 0, + )?; + } + } + } + ReconnectResult::Snapshot(result) => { + for snapshot in &result.snapshots { + if channels.contains(&snapshot.resource) { + return Err(WindowError::new( + "windowed snapshot must arrive as data", + )); + } + if ordinary + .channels + .iter() + .filter_map(|channel| channel.upgrade()) + .any(|channel| { + channel.uri == snapshot.resource && self.is_current(&channel) + }) + && !channels.contains(&snapshot.resource) + { + self.receive(json!({"jsonrpc":"2.0","method":"channel/snapshot","params":{"channel":snapshot.resource,"snapshot":snapshot}}), 0, 0)?; + } + } + } + } + for window in windows { + if let Some(channel) = window + .channel + .upgrade() + .filter(|channel| self.is_current(channel)) + { + if let Some(item) = items.iter().find(|item| item.channel == window.uri) { + self.install( + &channel, + window.offer, + Some(item.flow_control.clone()), + item.recovery == ahp_types::state::ChannelRecoveryKind::Snapshot, + )?; + channel.event("resumed"); + } else { + self.fail_subscription( + &channel, + WindowError::new("subscription dropped by reconnect"), + ); + } + } else if self.subscription(&window.uri).is_none() + && items.iter().any(|item| item.channel == window.uri) + { + cleanup.push(window.uri); + } + } + for uri in ordinary.uris { + if self.subscription(&uri).is_none() && !cleanup.contains(&uri) { + cleanup.push(uri); + } + } + } + Some(PendingKind::Initialize(channels)) => { + let result: ahp_types::commands::InitializeResult = + serde_json::from_value(result.clone()) + .map_err(|e| WindowError::new(e.to_string()))?; + for channel in channels + .channels + .iter() + .filter_map(|channel| channel.upgrade()) + .filter(|channel| self.is_current(channel)) + { + channel.lock().retained_seq = None; + self.install(&channel, None, None, false)?; + } + for snapshot in result.snapshots { + if channels + .channels + .iter() + .filter_map(|channel| channel.upgrade()) + .any(|channel| { + channel.uri == snapshot.resource && self.is_current(&channel) + }) + { + self.receive(json!({"jsonrpc":"2.0","method":"channel/snapshot","params":{"channel":snapshot.resource,"snapshot":snapshot}}), 0, 0)?; + } + } + for uri in channels.uris { + if self.subscription(&uri).is_none() { + cleanup.push(uri); + } + } + } + None => {} + } + Ok(cleanup) + } + pub(crate) fn failed_request(&self, id: u64) { + let operation = self.lock().pending.remove(&id); + let channels: Vec<_> = match operation.map(|operation| operation.kind) { + Some(PendingKind::Subscribe { channel, .. }) => channel.upgrade().into_iter().collect(), + Some(PendingKind::Reconnect { ordinary, .. }) + | Some(PendingKind::Initialize(ordinary)) => ordinary + .channels + .iter() + .filter_map(|channel| channel.upgrade()) + .filter(|channel| !channel.lock().managed) + .collect(), + None => vec![], + }; + for channel in channels { + self.cleanup_unaccepted(&channel); + } + } + pub(crate) fn expire_request(&self, id: u64) { + let channels = { + let mut engine = self.lock(); + let Some(operation) = engine.pending.get_mut(&id) else { + return; + }; + operation.abandoned = true; + let channels: Vec<_> = match &operation.kind { + PendingKind::Subscribe { channel, .. } => channel.upgrade().into_iter().collect(), + PendingKind::Reconnect { ordinary, .. } | PendingKind::Initialize(ordinary) => { + ordinary + .channels + .iter() + .filter_map(|channel| channel.upgrade()) + .collect() + } + }; + if !operation.sent { + engine.pending.remove(&id); + } + channels + }; + for channel in channels { + self.cleanup_unaccepted(&channel); + } + } + + pub(crate) fn has_request(&self, id: u64) -> bool { + self.lock().pending.contains_key(&id) + } + + pub(crate) fn request_sent(&self, id: u64) -> bool { + if let Some(operation) = self.lock().pending.get_mut(&id) { + if operation.abandoned { + return false; + } + operation.sent = true; + return true; + } + false + } + fn is_current(&self, channel: &Arc) -> bool { + self.subscription(&channel.uri) + .is_some_and(|current| Arc::ptr_eq(¤t, channel)) + } + + fn has_pending(engine: &EngineState, channel: &Arc) -> bool { + engine.pending.values().any(|operation| !operation.abandoned && matches!(&operation.kind, + PendingKind::Subscribe { channel: pending, .. } if pending.ptr_eq(&Arc::downgrade(channel)))) + } + + pub(crate) fn cleanup_unaccepted(&self, channel: &Arc) { + let mut engine = self.lock(); + if !channel.lock().accepted + && !Self::has_pending(&engine, channel) + && engine + .channels + .get(&channel.uri) + .is_some_and(|current| Arc::ptr_eq(current, channel)) + { + engine.channels.remove(&channel.uri); + channel.lock().pending_options = None; + self.close_channel( + channel, + Some(WindowError::new("subscription negotiation cancelled")), + ); + } + } + fn install( + &self, + channel: &Arc, + offer: Option, + limits: Option, + snapshot: bool, + ) -> Result { + let mut state = channel.lock(); + if state.closed { + return Ok(false); + } + if let Some(limits) = &limits { + validate_limits(&limits.client_receive)?; + validate_limits(&limits.host_receive)?; + let offer = offer + .as_ref() + .ok_or_else(|| WindowError::new("unoffered window acceptance"))?; + if limits.client_receive.window_bytes > offer.window_bytes + || limits.client_receive.maximum_frame_bytes > offer.maximum_frame_bytes + { + return Err(WindowError::new("accepted limits exceed offer")); + } + } else if channel.uri.starts_with("ahp-tcp:") { + return Err(WindowError::new("TCP requires window acceptance")); + } + if let Some(options) = state.pending_options.take() { + state.options = options; + } + if state.limits.is_some() != limits.is_some() { + state.epoch += 1; + state.sent = 0; + state.consumed = 0; + state.received_bytes = 0; + state.released = 0; + state.advertised = 0; + state.credit_since = None; + } + state.offer = offer; + state.limits = limits; + state.accepted = true; + state.snapshot_expected = snapshot && state.limits.is_some(); + if state.limits.is_none() && state.fallback.is_none() { + state.fallback = Some("declined".into()); + } else if state.limits.is_some() { + state.fallback = None; + } + channel.publish(&state); + channel.event(if state.limits.is_some() { + "negotiated" + } else { + "declined" + }); + self.wake.notify_one(); + if channel.unowned_drop() { + let client = channel + .binding + .lock() + .unwrap_or_else(|error| error.into_inner()) + .as_ref() + .and_then(crate::client::ClientBinding::upgrade); + if let Some(client) = client { + channel.schedule_drop(client); + } + } + Ok(true) + } + pub(crate) fn inbound( + &self, + notification: JsonRpcNotification, + bytes: Option, + ) -> Result, WindowError> { + if !counts_against_window(¬ification.method) + && bytes.is_some_and(|bytes| bytes > MAX_CONTROL_BYTES) + { + return Err(WindowError::new("control message exceeds size limit")); + } + let params = notification.params.as_ref().unwrap_or(&Value::Null); + let Some(uri) = params["channel"].as_str() else { + return Ok(Some(notification)); + }; + let Some(channel) = self.subscription(uri) else { + return Ok(Some(notification)); + }; + if channel.lock().closed { + return Ok(None); + } + if !channel.lock().accepted && !channel.windowed() { + if matches!( + notification.method.as_str(), + "channel/frame" | "channel/snapshot" | "channel/credit" + ) { + // Traffic from a retired subscription may still be in flight. + return Ok(None); + } + return Ok(Some(notification)); + } + if notification.method == "channel/credit" { + let credit = params["consumedBytes"] + .as_u64() + .filter(|v| *v <= SAFE_INTEGER) + .ok_or_else(|| WindowError::new("invalid credit"))?; + let mut state = channel.lock(); + if state.closed { + return Ok(None); + } + if state.limits.is_none() { + return Err(WindowError::new("credit on ordinary subscription")); + } + if credit > state.sent { + return Err(WindowError::new("credit exceeds sent bytes")); + } + if credit > state.consumed { + state.consumed = credit; + channel.publish(&state); + channel.event("creditReceived"); + self.wake.notify_one(); + } + return Ok(None); + } + let charged = + channel.lock().limits.is_some() && counts_against_window(¬ification.method); + let bytes = if charged { + bytes.ok_or_else(|| WindowError::new("transport lost raw text"))? + } else { + 0 + }; + if charged { + let mut state = channel.lock(); + if state.closed { + return Ok(None); + } + if !state.accepted { + return Err(WindowError::new("data before negotiation response")); + } + let limits = state + .limits + .as_ref() + .ok_or_else(|| WindowError::new("missing limits"))?; + if bytes > limits.client_receive.maximum_frame_bytes as u64 { + return Err(WindowError::new("data exceeds maximum frame bytes")); + } + state.received_bytes = checked_total(state.received_bytes, bytes)?; + channel.publish(&state); + } + if notification.method == "channel/frame" { + if !charged { + return Err(WindowError::new("frame on ordinary subscription")); + } + let data = params["data"] + .as_str() + .filter(|s| !s.is_empty()) + .ok_or_else(|| WindowError::new("empty fragment"))?; + if params.get("final").is_some_and(|value| !value.is_boolean()) { + return Err(WindowError::new("invalid fragment final flag")); + } + let mut state = channel.lock(); + if state.closed { + return Ok(None); + } + state.fragment.push_str(data); + state.fragment_bytes = checked_total(state.fragment_bytes, bytes)?; + if params["final"].as_bool() != Some(true) { + state.released += bytes; + channel.publish(&state); + self.wake.notify_one(); + return Ok(None); + } + let joined = std::mem::take(&mut state.fragment); + let total = std::mem::take(&mut state.fragment_bytes); + drop(state); + let value: Value = + serde_json::from_str(&joined).map_err(|e| WindowError::new(e.to_string()))?; + if value["params"]["channel"] != uri + || value.get("id").is_some() + || value["method"] + .as_str() + .map_or(true, |m| m == "channel/frame" || !counts_against_window(m)) + { + return Err(WindowError::new("invalid joined fragment notification")); + } + if !self.receive_on(&channel, value.clone(), total, bytes)? { + return Ok(None); + } + channel.event("fragmentedMessageReceived"); + let message: JsonRpcNotification = + serde_json::from_value(value).map_err(|e| WindowError::new(e.to_string()))?; + return Ok(Some(message)); + } + if charged && !channel.lock().fragment.is_empty() { + return Err(WindowError::new("interleaved fragment data")); + } + let value = + serde_json::to_value(¬ification).map_err(|e| WindowError::new(e.to_string()))?; + if !self.receive_on(&channel, value, bytes, bytes)? { + return Ok(None); + } + Ok(Some(notification)) + } + fn receive(&self, value: Value, bytes: u64, release_bytes: u64) -> Result { + let uri = value["params"]["channel"] + .as_str() + .ok_or_else(|| WindowError::new("channel missing"))?; + let Some(channel) = self.subscription(uri) else { + return Ok(false); + }; + self.receive_on(&channel, value, bytes, release_bytes) + } + + fn receive_on( + &self, + channel: &Arc, + value: Value, + bytes: u64, + release_bytes: u64, + ) -> Result { + let uri = channel.uri.as_str(); + let mut state = channel.lock(); + if state.closed { + return Ok(false); + } + let notification: JsonRpcNotification = + serde_json::from_value(value.clone()).map_err(|e| WindowError::new(e.to_string()))?; + crate::client::decode_notification(notification.clone()).map_err(WindowError::new)?; + if state.snapshot_expected && counts_against_window(¬ification.method) { + if notification.method != "channel/snapshot" { + return Err(WindowError::new("first delivery must be snapshot")); + } + state.snapshot_expected = false; + } + if notification.method == "channel/snapshot" + && value["params"]["snapshot"]["resource"] != uri + { + return Err(WindowError::new("snapshot resource mismatch")); + } + let seq = value["params"]["serverSeq"] + .as_i64() + .or_else(|| value["params"]["snapshot"]["fromSeq"].as_i64()); + if let Some(seq) = seq { + state.retained_seq = Some(state.retained_seq.unwrap_or(0).max(seq)); + } + let epoch = state.epoch; + if state.options.coalescing && state.limits.is_some() { + if let Some(tail) = state.received.back_mut() { + if tail.epoch == epoch { + if let Some(merged) = merge( + &self.config.coalescers, + &tail.notification, + ¬ification, + Direction::Receive, + ) { + tail.notification = merged; + tail.bytes += bytes; + tail.release_bytes += release_bytes; + channel.event("coalesced"); + channel.ready.notify_one(); + return Ok(true); + } + } + } + } + let reserved_snapshot = usize::from( + !state.initial_delivered + && state + .received + .front() + .is_some_and(|item| item.notification.method == "channel/snapshot"), + ); + if state.limits.is_none() + && state.received.len().saturating_sub(reserved_snapshot) >= channel.ordinary_buffer + { + state.received.remove(reserved_snapshot); + } + state.received.push_back(Received { + notification, + bytes, + release_bytes, + epoch, + }); + channel.ready.notify_one(); + Ok(true) + } + pub(crate) async fn enqueue( + &self, + notification: JsonRpcNotification, + completion: Option>>, + ) -> Result<(), WindowError> { + let uri = notification + .params + .as_ref() + .and_then(|p| p["channel"].as_str()) + .ok_or_else(|| WindowError::new("channel missing"))?; + let channel = self + .subscription(uri) + .ok_or_else(|| WindowError::new("subscription missing"))?; + let text = + serde_json::to_string(¬ification).map_err(|e| WindowError::new(e.to_string()))?; + let bytes = utf16_bytes(&text); + let mut completion = completion; + let mut changes = channel.status.subscribe(); + loop { + { + let mut state = channel.lock(); + if state.closed { + return Err(state + .error + .clone() + .unwrap_or_else(|| WindowError::new("subscription closed"))); + } + let behind = state + .queued_bytes + .saturating_sub(state.outgoing.front().map_or(0, |item| item.bytes)); + if state.outgoing.is_empty() || behind.saturating_add(bytes) <= channel.max_send { + if state.options.coalescing && state.limits.is_some() { + if let Some(tail) = state.outgoing.back_mut().filter(|tail| { + tail.offset == 0 && tail.bytes.saturating_add(bytes) <= channel.max_send + }) { + if let Some(merged) = merge( + &self.config.coalescers, + &tail.notification, + ¬ification, + Direction::Send, + ) { + let merged_text = serde_json::to_string(&merged) + .map_err(|e| WindowError::new(e.to_string()))?; + tail.notification = merged; + tail.text = merged_text; + tail.bytes += bytes; + tail.parts.push(QueuedPart { + notification: notification.clone(), + completion: completion.take(), + }); + state.queued_bytes += bytes; + channel.publish(&state); + channel.event("coalesced"); + self.wake.notify_one(); + return Ok(()); + } + } + } + state.outgoing.push_back(Queued { + notification: notification.clone(), + text, + offset: 0, + bytes, + parts: vec![QueuedPart { + notification, + completion: completion.take(), + }], + }); + state.queued_bytes += bytes; + channel.publish(&state); + self.wake.notify_one(); + return Ok(()); + } + } + changes + .changed() + .await + .map_err(|_| WindowError::new("subscription closed"))?; + } + } + pub(crate) fn next_send(&self) -> Result, WindowError> { + let mut engine = self.lock(); + let mut channels: Vec<_> = engine.channels.values().cloned().collect(); + channels.sort_by(|a, b| a.uri.cmp(&b.uri)); + for channel in &channels { + let mut state = channel.lock(); + if state.closed || !state.accepted { + continue; + } + if state.limits.is_some() && state.released > state.advertised { + let limits = state + .limits + .as_ref() + .ok_or_else(|| WindowError::new("missing limits"))?; + let threshold = (limits.client_receive.window_bytes as u64 / 4).clamp(1, 16_384); + let stalled = state.received_bytes - state.advertised + >= limits.client_receive.window_bytes as u64; + let since = *state + .credit_since + .get_or_insert_with(tokio::time::Instant::now); + if !stalled + && state.released - state.advertised < threshold + && since.elapsed() < std::time::Duration::from_millis(5) + { + continue; + } + state.advertised = state.released; + state.credit_since = None; + let text = json!({"jsonrpc":"2.0","method":"channel/credit","params":{"channel":channel.uri,"consumedBytes":state.released}}).to_string(); + return Ok(Some(WireSend { + text, + control: true, + completions: vec![], + event: Some((channel.clone(), "creditSent")), + channel: channel.clone(), + epoch: state.epoch, + })); + } + } + for index in 0..channels.len() { + let position = (engine.cursor + index) % channels.len(); + let channel = &channels[position]; + let mut state = channel.lock(); + if state.closed || !state.accepted { + continue; + } + self.remove_cancelled(&mut state)?; + if state.outgoing.is_empty() { + channel.publish(&state); + continue; + } + let outstanding = state.sent - state.consumed; + let limits = state.limits.as_ref().map(|l| l.host_receive.clone()); + if limits + .as_ref() + .is_some_and(|l| outstanding >= l.window_bytes as u64) + { + continue; + } + let mut item = state + .outgoing + .pop_front() + .ok_or_else(|| WindowError::new("queue disappeared"))?; + let (text, final_frame) = if limits + .as_ref() + .is_some_and(|l| utf16_bytes(&item.text) > l.maximum_frame_bytes as u64) + { + let limits = limits + .as_ref() + .ok_or_else(|| WindowError::new("missing send limits"))?; + let (text, end) = fragment( + &channel.uri, + &item.text, + item.offset, + limits.maximum_frame_bytes as u64, + (limits.window_bytes as u64).saturating_sub(outstanding), + )?; + item.offset = end; + (text, end == item.text.len()) + } else { + (item.text.clone(), true) + }; + if limits.is_some() { + state.sent = checked_total(state.sent, utf16_bytes(&text))?; + } + let event = (final_frame && item.offset != 0) + .then(|| (channel.clone(), "fragmentedMessageSent")); + let completions = if final_frame { + state.queued_bytes -= item.bytes; + item.parts + .into_iter() + .filter_map(|part| part.completion) + .collect() + } else { + state.outgoing.push_front(item); + vec![] + }; + channel.publish(&state); + engine.cursor = position + 1; + return Ok(Some(WireSend { + text, + control: false, + completions, + event, + channel: channel.clone(), + epoch: state.epoch, + })); + } + Ok(None) + } + fn close_channel(&self, channel: &Arc, error: Option) { + let mut state = channel.lock(); + state.closed = true; + state.error = error.clone(); + state.credit_since = None; + state.fragment.clear(); + state.fragment_bytes = 0; + for item in state.outgoing.drain(..) { + for completion in item.parts.into_iter().filter_map(|part| part.completion) { + let _ = completion.send(Err(error + .clone() + .unwrap_or_else(|| WindowError::new("connection closed")))); + } + } + state.queued_bytes = 0; + channel.publish(&state); + channel.status.send_modify(|_| {}); + channel.ready.notify_waiters(); + if (!state.managed || state.dropped) && !state.dropped_notified { + state.dropped_notified = true; + channel.event("dropped"); + } + } + pub(crate) fn close(&self, error: Option) { + self.lock().pending.clear(); + let channels: Vec<_> = self.lock().channels.values().cloned().collect(); + for channel in channels { + self.close_channel(&channel, error.clone()); + } + } + pub(crate) fn remove_and( + &self, + uri: &str, + expected: Option<&Arc>, + queued: impl FnOnce(), + ) -> bool { + let mut engine = self.lock(); + if let Some(expected) = expected { + if !engine + .channels + .get(uri) + .is_some_and(|current| Arc::ptr_eq(current, expected)) + || expected.owners.load(Ordering::Acquire) != 0 + || expected.acquisitions.load(Ordering::Acquire) != 0 + || expected.lock().pending_options.is_some() + || Self::has_pending(&engine, expected) + { + return false; + } + } + if let Some(channel) = engine.channels.remove(uri) { + channel.finish(None); + self.close_channel(&channel, None); + let hook = channel.lock().drop_hook.clone(); + if let Some(hook) = hook { + hook(); + } + } + queued(); + true + } + + pub(crate) fn fail_subscription(&self, channel: &Arc, error: WindowError) { + let mut engine = self.lock(); + if engine + .channels + .get(&channel.uri) + .is_some_and(|current| Arc::ptr_eq(current, channel)) + { + engine.channels.remove(&channel.uri); + } + channel.finish(Some(error.clone())); + self.close_channel(channel, Some(error)); + let hook = channel.lock().drop_hook.clone(); + if let Some(hook) = hook { + hook(); + } + } + fn remove_cancelled(&self, state: &mut ChannelState) -> Result<(), WindowError> { + while state.outgoing.front().is_some_and(|item| { + item.offset == 0 + && item.parts.iter().any(|part| { + part.completion + .as_ref() + .is_some_and(oneshot::Sender::is_closed) + }) + }) { + let item = state + .outgoing + .pop_front() + .ok_or_else(|| WindowError::new("queue disappeared"))?; + state.queued_bytes -= item.bytes; + let mut rebuilt: Vec = vec![]; + for part in item.parts.into_iter().filter(|part| { + !part + .completion + .as_ref() + .is_some_and(oneshot::Sender::is_closed) + }) { + let text = serde_json::to_string(&part.notification) + .map_err(|e| WindowError::new(e.to_string()))?; + let bytes = utf16_bytes(&text); + state.queued_bytes += bytes; + if let Some(tail) = rebuilt.last_mut() { + if let Some(merged) = merge( + &self.config.coalescers, + &tail.notification, + &part.notification, + Direction::Send, + ) { + tail.text = serde_json::to_string(&merged) + .map_err(|e| WindowError::new(e.to_string()))?; + tail.notification = merged; + tail.bytes += bytes; + tail.parts.push(part); + continue; + } + } + rebuilt.push(Queued { + notification: part.notification.clone(), + text, + offset: 0, + bytes, + parts: vec![part], + }); + } + for item in rebuilt.into_iter().rev() { + state.outgoing.push_front(item); + } + } + Ok(()) + } + + pub(crate) fn managed_channels(&self) -> Vec> { + self.lock() + .channels + .values() + .filter(|channel| { + !channel.uri.starts_with("ahp-tcp:") + && !channel.lock().dropped + && !channel.unowned_drop() + }) + .cloned() + .collect() + } + + pub(crate) fn adopt(&self, channel: Arc) -> Result<(), WindowError> { + let options = { + let mut state = channel.lock(); + state.epoch += 1; + state.sent = 0; + state.consumed = 0; + state.received_bytes = 0; + state.released = 0; + state.advertised = 0; + state.fragment.clear(); + state.credit_since = None; + state.fragment_bytes = 0; + state.closed = false; + state.error = None; + state.dropped_notified = false; + state.accepted = false; + state.options.clone() + }; + self.lock() + .channels + .insert(channel.uri.clone(), channel.clone()); + *channel.events.lock().unwrap_or_else(|e| e.into_inner()) = self.events.clone(); + let offer = self.prepare(&channel.uri, options)?; + let mut state = channel.lock(); + state.offer = offer; + if state.offer.is_none() { + state.limits = None; + state.snapshot_expected = false; + } + channel.publish(&state); + Ok(()) + } + + pub(crate) fn resume_windows( + &self, + ) -> Option { + use ahp_types::commands::{ + ChannelDeliveryResumeOptions, ChannelDeliveryResumeParams, + SubscriptionFlowControlOptions, + }; + let items: Vec<_> = self + .managed_channels() + .into_iter() + .filter_map(|channel| { + let state = channel.lock(); + state.limits.as_ref()?; + Some(ChannelDeliveryResumeParams { + channel: channel.uri.clone(), + flow_control: SubscriptionFlowControlOptions { + receive: state.offer.clone()?, + }, + last_applied_server_seq: state.retained_seq, + }) + }) + .collect(); + (!items.is_empty()).then_some(ChannelDeliveryResumeOptions { items }) + } +} + +fn validate_limits(limits: &ChannelReceiveLimits) -> Result<(), WindowError> { + if limits.window_bytes < 1 + || limits.maximum_frame_bytes < 1024 + || limits.window_bytes as u64 > SAFE_INTEGER + || limits.maximum_frame_bytes as u64 > SAFE_INTEGER + { + return Err(WindowError::new("invalid window limits")); + } + Ok(()) +} +fn checked_total(value: u64, increment: u64) -> Result { + value + .checked_add(increment) + .filter(|v| *v <= SAFE_INTEGER) + .ok_or_else(|| WindowError::new("window counter overflow")) +} +fn fragment( + uri: &str, + text: &str, + offset: usize, + maximum: u64, + remaining: u64, +) -> Result<(String, usize), WindowError> { + let ends: Vec<_> = text[offset..] + .char_indices() + .map(|(i, c)| offset + i + c.len_utf8()) + .collect(); + let first = *ends + .first() + .ok_or_else(|| WindowError::new("empty outbound fragment"))?; + let encode = |end: usize| { + json!({"jsonrpc":"2.0","method":"channel/frame","params":{"channel":uri,"data":&text[offset..end],"final":end==text.len()}}).to_string() + }; + let minimum = utf16_bytes(&encode(first)); + if minimum > maximum { + return Err(WindowError::new( + "frame envelope cannot carry one character", + )); + } + let target = maximum.min(remaining.max(minimum)); + let mut low = 0; + let mut high = ends.len(); + while low < high { + let middle = (low + high) / 2; + if utf16_bytes(&encode(ends[middle])) <= target { + low = middle + 1; + } else { + high = middle; + } + } + let end = ends[low.saturating_sub(1)]; + Ok((encode(end), end)) +} +fn merge( + rules: &[Arc], + earlier: &JsonRpcNotification, + later: &JsonRpcNotification, + direction: Direction, +) -> Option { + if earlier.method != later.method + || !counts_against_window(&earlier.method) + || matches!( + earlier.method.as_str(), + "channel/snapshot" | "channel/frame" | "tcp/eof" | "dispatchAction" + ) + { + return None; + } + let a = earlier.params.as_ref()?; + let b = later.params.as_ref()?; + if a["channel"] != b["channel"] + || a.get("origin").is_some() + || b.get("origin").is_some() + || a.get("rejectionReason").is_some() + || b.get("rejectionReason").is_some() + { + return None; + } + let context = CoalescingContext { + channel: a["channel"].as_str()?.into(), + direction, + }; + for rule in rules { + if let Some(result) = rule.merge(earlier, later, &context) { + if result.method == earlier.method + && result + .params + .as_ref() + .is_some_and(|params| params["channel"] == a["channel"]) + && crate::client::decode_notification(result.clone()).is_ok() + { + return Some(result); + } + } + } + let mut merged = later.clone(); + let params = merged.params.as_mut()?; + if direction == Direction::Receive && earlier.method == "action" { + let kind = a["action"]["type"].as_str()?; + if kind != b["action"]["type"].as_str()? { + return None; + } + let field = match kind { + "terminal/data" => "data", + "chat/delta" + if a["action"]["turnId"] == b["action"]["turnId"] + && a["action"]["partId"] == b["action"]["partId"] + && a["action"]["_meta"] == b["action"]["_meta"] => + { + "content" + } + _ => return None, + }; + params["action"][field] = format!( + "{}{}", + a["action"][field].as_str()?, + b["action"][field].as_str()? + ) + .into(); + return Some(merged); + } + if direction == Direction::Send && earlier.method == "tcp/data" { + use base64::Engine as _; + let mut data = base64::engine::general_purpose::STANDARD + .decode(a["data"].as_str()?) + .ok()?; + data.extend( + base64::engine::general_purpose::STANDARD + .decode(b["data"].as_str()?) + .ok()?, + ); + params["data"] = base64::engine::general_purpose::STANDARD + .encode(data) + .into(); + return Some(merged); + } + None +} diff --git a/clients/rust/crates/ahp/src/windowing_tests.rs b/clients/rust/crates/ahp/src/windowing_tests.rs new file mode 100644 index 000000000..a3a95a63d --- /dev/null +++ b/clients/rust/crates/ahp/src/windowing_tests.rs @@ -0,0 +1,434 @@ +#![allow(clippy::unwrap_used, clippy::panic)] +use super::*; + +fn notification(uri: &str, method: &str, params: Value) -> JsonRpcNotification { + let mut params = params; + params["channel"] = uri.into(); + serde_json::from_value(json!({"jsonrpc":"2.0","method":method,"params":params})).unwrap() +} +fn engine(config: WindowingConfig) -> Engine { + Engine::new( + config, + WindowTransportCapabilities::default().raw_text(true), + 256, + ) +} +fn accept(engine: &Engine, uri: &str, window: i64) { + let offer = engine + .prepare(uri, WindowOptions::default()) + .unwrap() + .unwrap(); + let channel = engine.subscription(uri).unwrap(); + engine + .install( + &channel, + Some(offer), + Some(ChannelFlowControl { + client_receive: ChannelReceiveLimits { + window_bytes: 1024, + maximum_frame_bytes: 1024, + }, + host_receive: ChannelReceiveLimits { + window_bytes: window, + maximum_frame_bytes: 1024, + }, + }), + true, + ) + .unwrap(); +} + +#[tokio::test] +async fn fair_fragment_scheduling_unicode_and_cancellation() { + let engine = engine(WindowingConfig::default()); + accept(&engine, "ahp-terminal:/a", 128); + accept(&engine, "ahp-terminal:/b", 128); + let payload = "😀漢字\"\\\n".repeat(500); + let a = notification("ahp-terminal:/a", "custom/data", json!({"text":payload})); + let b = notification("ahp-terminal:/b", "custom/data", json!({"text":"small"})); + let (tx, rx) = oneshot::channel(); + engine.enqueue(a.clone(), Some(tx)).await.unwrap(); + engine.enqueue(b, None).await.unwrap(); + let first = engine.next_send().unwrap().unwrap(); + let value: Value = serde_json::from_str(&first.text).unwrap(); + assert_eq!(value["params"]["channel"], "ahp-terminal:/a"); + assert_eq!(value["method"], "channel/frame"); + let mut joined = value["params"]["data"].as_str().unwrap().to_string(); + first.complete(Ok(())); + let second = engine.next_send().unwrap().unwrap(); + let value: Value = serde_json::from_str(&second.text).unwrap(); + assert_eq!(value["params"]["channel"], "ahp-terminal:/b"); + second.complete(Ok(())); + assert!(engine.next_send().unwrap().is_none()); + let mut rx = rx; + assert!(rx.try_recv().is_err()); + loop { + let sent = engine.subscription("ahp-terminal:/a").unwrap().lock().sent; + engine + .inbound( + notification( + "ahp-terminal:/a", + "channel/credit", + json!({"consumedBytes":sent}), + ), + None, + ) + .unwrap(); + let wire = engine.next_send().unwrap().unwrap(); + assert!(utf16_bytes(&wire.text) <= 1024); + let value: Value = serde_json::from_str(&wire.text).unwrap(); + joined.push_str(value["params"]["data"].as_str().unwrap()); + let final_frame = value["params"]["final"] == true; + wire.complete(Ok(())); + if final_frame { + break; + } + } + assert_eq!( + serde_json::from_str::(&joined).unwrap(), + serde_json::to_value(a).unwrap() + ); + rx.await.unwrap().unwrap(); + let channel = engine.subscription("ahp-terminal:/b").unwrap(); + let sent = channel.lock().sent; + engine + .inbound( + notification( + "ahp-terminal:/b", + "channel/credit", + json!({"consumedBytes":sent}), + ), + None, + ) + .unwrap(); + let (tx, rx) = oneshot::channel(); + engine + .enqueue( + notification("ahp-terminal:/b", "custom/data", json!({"text":"cancel"})), + Some(tx), + ) + .await + .unwrap(); + drop(rx); + assert!(engine.next_send().unwrap().is_none()); + assert_eq!(channel.status.borrow().buffered_send_bytes, 0); +} + +#[test] +fn custom_coalescers_validate_identity_and_preserve_credit() { + let config = WindowingConfig::default().coalescer( + |earlier: &JsonRpcNotification, later: &JsonRpcNotification, _: &CoalescingContext| { + if earlier.method != "custom/amount" { + return None; + } + let mut merged = later.clone(); + merged.params.as_mut()?["amount"] = (earlier.params.as_ref()?["amount"].as_i64()? + + later.params.as_ref()?["amount"].as_i64()?) + .into(); + Some(merged) + }, + ); + let engine = engine(config); + let uri = "ahp-terminal:/custom"; + accept(&engine, uri, 1024); + engine + .inbound( + notification( + uri, + "channel/snapshot", + json!({"snapshot":{"resource":uri,"state":{},"fromSeq":0}}), + ), + Some(300), + ) + .unwrap(); + engine + .inbound( + notification(uri, "custom/amount", json!({"amount":2})), + Some(214), + ) + .unwrap(); + engine + .inbound( + notification(uri, "custom/amount", json!({"amount":3})), + Some(220), + ) + .unwrap(); + let channel = engine.subscription(uri).unwrap(); + let state = channel.lock(); + assert_eq!(state.received.len(), 2); + let merged = state.received.back().unwrap(); + assert_eq!(merged.bytes, 434); + assert_eq!(merged.release_bytes, 434); + assert_eq!(merged.notification.params.as_ref().unwrap()["amount"], 5); + drop(state); + let invalid: Vec> = vec![Arc::new( + |_: &JsonRpcNotification, later: &JsonRpcNotification, _: &CoalescingContext| { + let mut value = later.clone(); + value.params.as_mut().unwrap()["channel"] = "ahp-terminal:/wrong".into(); + Some(value) + }, + )]; + assert!(merge( + &invalid, + ¬ification(uri, "custom/amount", json!({"amount":1})), + ¬ification(uri, "custom/amount", json!({"amount":2})), + Direction::Receive + ) + .is_none()); +} + +#[tokio::test] +async fn queue_bound_and_failed_resubscribe_preserve_options() { + let engine = engine(WindowingConfig::default().max_buffered_send_bytes(512)); + let uri = "ahp-terminal:/bounded"; + accept(&engine, uri, 1); + engine + .enqueue( + notification(uri, "custom/data", json!({"text":"x".repeat(512)})), + None, + ) + .await + .unwrap(); + let channel = engine.subscription(uri).unwrap(); + assert!(channel.status.borrow().buffered_send_bytes > 512); + engine + .prepare(uri, WindowOptions::default().coalescing(false)) + .unwrap(); + engine.register_request(2, "subscribe", &mut json!({"channel":uri,"flowControl":{"receive":{"windowBytes":262144,"maximumFrameBytes":1048576}}})).unwrap(); + engine.failed_request(2); + assert!(channel.lock().options.coalescing); + assert!(channel.lock().limits.is_some()); +} + +#[test] +fn invalid_limits_and_fragment_flags_are_protocol_errors() { + let engine = engine(WindowingConfig::default()); + for options in [ + WindowOptions::default().window_bytes(0), + WindowOptions::default().maximum_frame_bytes(100), + WindowOptions::default().window_bytes(SAFE_INTEGER + 1), + ] { + assert!(engine.prepare("ahp-terminal:/bad", options).is_err()); + } + let uri = "ahp-terminal:/a"; + accept(&engine, uri, 1024); + assert!(engine + .inbound( + notification(uri, "channel/frame", json!({"data":"x","final":0})), + Some(240) + ) + .is_err()); +} + +#[test] +fn transport_fallback_protects_initial_snapshot_in_lossy_queue() { + let engine = Engine::new( + WindowingConfig::default(), + WindowTransportCapabilities::default(), + 1, + ); + let uri = "ahp-terminal:/parsed"; + assert!(engine + .prepare(uri, WindowOptions::default()) + .unwrap() + .is_none()); + let mut request = json!({"channel":uri,"flowControl":{"receive":{"windowBytes":512,"maximumFrameBytes":1024}}}); + engine + .register_request(1, "subscribe", &mut request) + .unwrap(); + assert!(request.get("flowControl").is_none()); + engine + .response( + 1, + &json!({"snapshot":{"resource":uri,"state":{},"fromSeq":1}}), + ) + .unwrap(); + for count in 0..10 { + engine + .inbound( + notification(uri, "custom/data", json!({"count":count})), + None, + ) + .unwrap(); + } + let channel = engine.subscription(uri).unwrap(); + let state = channel.lock(); + assert_eq!(state.received.len(), 2); + assert_eq!(state.received[0].notification.method, "channel/snapshot"); + assert_eq!( + state.received[1].notification.params.as_ref().unwrap()["count"], + 9 + ); + assert_eq!( + channel.status.borrow().fallback_reason.as_deref(), + Some("transport") + ); +} + +#[tokio::test] +async fn cancelling_one_coalesced_send_preserves_other_sends() { + let engine = engine(WindowingConfig::default()); + let uri = "ahp-tcp:/a"; + accept(&engine, uri, 1); + let (keep_tx, keep_rx) = oneshot::channel(); + let (cancel_tx, cancel_rx) = oneshot::channel(); + engine + .enqueue( + notification(uri, "tcp/data", json!({"data":"aGVs"})), + Some(keep_tx), + ) + .await + .unwrap(); + engine + .enqueue( + notification(uri, "tcp/data", json!({"data":"bG8="})), + Some(cancel_tx), + ) + .await + .unwrap(); + assert_eq!(engine.subscription(uri).unwrap().lock().outgoing.len(), 1); + drop(cancel_rx); + let wire = engine.next_send().unwrap().unwrap(); + let value: Value = serde_json::from_str(&wire.text).unwrap(); + assert_eq!(value["params"]["data"], "aGVs"); + wire.complete(Ok(())); + keep_rx.await.unwrap().unwrap(); + assert!(engine.next_send().unwrap().is_none()); +} + +#[test] +fn upgrading_ordinary_subscription_starts_new_owner_at_snapshot() { + let engine = engine(WindowingConfig::default()); + let uri = "ahp-terminal:/upgrade"; + let offer = engine + .prepare(uri, WindowOptions::default()) + .unwrap() + .unwrap(); + let ordinary_action = notification( + uri, + "action", + json!({"serverSeq":1,"action":{"type":"terminal/data","data":"old"}}), + ); + assert!(engine + .inbound(ordinary_action, Some(280)) + .unwrap() + .is_some()); + let channel = engine.subscription(uri).unwrap(); + assert!(channel.lock().received.is_empty()); + engine + .install( + &channel, + Some(offer), + Some(ChannelFlowControl { + client_receive: ChannelReceiveLimits { + window_bytes: 1024, + maximum_frame_bytes: 1024, + }, + host_receive: ChannelReceiveLimits { + window_bytes: 1024, + maximum_frame_bytes: 1024, + }, + }), + true, + ) + .unwrap(); + engine + .inbound( + notification( + uri, + "channel/snapshot", + json!({"snapshot":{"resource":uri,"state":{},"fromSeq":2}}), + ), + Some(300), + ) + .unwrap(); + assert_eq!( + channel.lock().received.front().unwrap().notification.method, + "channel/snapshot" + ); +} + +#[test] +fn cancelled_negotiation_bookkeeping_is_bounded_and_pruned() { + let engine = engine(WindowingConfig::default()); + for id in 0..256 { + let uri = format!("ahp-terminal:/cancel-{id}"); + let offer = engine + .prepare(&uri, WindowOptions::default()) + .unwrap() + .unwrap(); + engine + .register_request( + id, + "subscribe", + &mut json!({"channel":uri,"flowControl":{"receive":offer}}), + ) + .unwrap(); + engine.request_sent(id); + engine.expire_request(id); + assert!(!engine.has_subscriptions()); + } + assert_eq!(engine.lock().pending.len(), 256); + assert!(engine + .prepare("ahp-terminal:/bounded", WindowOptions::default()) + .is_err()); + for id in 0..256 { + engine.failed_request(id); + } + assert!(engine.lock().pending.is_empty()); + for _ in 0..100 { + accept(&engine, "ahp-terminal:/same", 1024); + engine.remove_and("ahp-terminal:/same", None, || {}); + } + assert!(!engine.has_subscriptions()); + assert!(engine.lock().pending.is_empty()); +} + +#[tokio::test(start_paused = true)] +async fn unsubscribe_clears_credit_timers_and_old_receipts_cannot_credit_replacements() { + let engine = engine(WindowingConfig::default()); + let uri = "ahp-terminal:/pending"; + accept(&engine, uri, 1024); + let channel = engine.subscription(uri).unwrap(); + let snapshot = notification( + uri, + "channel/snapshot", + json!({"snapshot":{"resource":uri,"state":{},"fromSeq":0}}), + ); + engine.inbound(snapshot, Some(300)).unwrap(); + fn take(channel: &Arc) -> Delivery { + let item = channel.lock().received.pop_front().unwrap(); + Delivery { + event: crate::client::decode_notification(item.notification.clone()).unwrap(), + notification: item.notification, + bytes: item.bytes, + release_bytes: item.release_bytes, + epoch: item.epoch, + channel: channel.clone(), + } + } + take(&channel).release(); + engine.next_send().unwrap().unwrap().complete(Ok(())); + for _ in 0..2 { + let note = notification(uri, "n", json!({})); + let bytes = utf16_bytes(&serde_json::to_string(¬e).unwrap()); + assert!(bytes < 256); + engine.inbound(note, Some(bytes)).unwrap(); + } + take(&channel).release(); + let mut old = take(&channel); + assert!(engine.next_send().unwrap().is_none()); + assert!(channel.lock().credit_since.is_some()); + engine.remove_and(uri, None, || {}); + assert!(channel.lock().credit_since.is_none()); + let released = channel.lock().released; + old.release(); + assert_eq!(channel.lock().released, released); + tokio::time::advance(std::time::Duration::from_millis(10)).await; + assert!(engine.next_send().unwrap().is_none()); + accept(&engine, uri, 1024); + let current = engine.subscription(uri).unwrap(); + assert!(!Arc::ptr_eq(¤t, &channel)); + assert_eq!(current.lock().received_bytes, 0); + assert_eq!(current.lock().released, 0); +} diff --git a/clients/rust/crates/ahp/tests/windowing_corpus.rs b/clients/rust/crates/ahp/tests/windowing_corpus.rs new file mode 100644 index 000000000..92fdf4a61 --- /dev/null +++ b/clients/rust/crates/ahp/tests/windowing_corpus.rs @@ -0,0 +1,707 @@ +#![allow(clippy::unwrap_used, clippy::panic)] + +use std::collections::{HashMap, VecDeque}; +use std::path::PathBuf; +use std::sync::Arc; +use std::time::Duration; + +use ahp::transport::TransportMessage; +use ahp::{ + Client, Delivery, Transport, TransportError, WindowOptions, WindowedSubscription, + WindowingConfig, +}; +use serde_json::{json, Value}; +use tokio::sync::{mpsc, Mutex, Notify}; + +struct MemoryTransport { + inbound: mpsc::UnboundedReceiver, + outbound: mpsc::UnboundedSender, + closed: Arc, +} +impl Transport for MemoryTransport { + fn window_capabilities(&self) -> ahp::WindowTransportCapabilities { + ahp::WindowTransportCapabilities::default().raw_text(true) + } + async fn send(&mut self, message: TransportMessage) -> Result<(), TransportError> { + let TransportMessage::Text(text) = message else { + panic!("expected exact text") + }; + self.outbound.send(text).map_err(|_| TransportError::Closed) + } + async fn recv(&mut self) -> Result, TransportError> { + Ok(self.inbound.recv().await) + } + async fn close(&mut self) -> Result<(), TransportError> { + self.closed.notify_one(); + Ok(()) + } +} +fn matches(actual: &Value, expected: &Value) -> bool { + if let Some(object) = expected.as_object() { + object.iter().all(|(key, expected)| { + actual + .get(key) + .is_some_and(|value| matches(value, expected)) + }) + } else { + actual == expected + } +} +fn corpus() -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../../../../types/test-cases/windowed-delivery") +} +#[tokio::test] +async fn all_shared_sizes() { + let data: Value = + serde_json::from_str(&std::fs::read_to_string(corpus().join("sizes.json")).unwrap()) + .unwrap(); + for case in data["cases"].as_array().unwrap() { + assert_eq!( + ahp::windowing::utf16_bytes(case["text"].as_str().unwrap()), + case["utf16Bytes"].as_u64().unwrap(), + "{}", + case["name"] + ); + } +} +#[tokio::test] +async fn all_shared_scenarios() { + let mut files: Vec<_> = std::fs::read_dir(corpus().join("scenarios")) + .unwrap() + .map(|entry| entry.unwrap().path()) + .collect(); + files.sort(); + assert_eq!(files.len(), 24); + for file in files { + let scenario: Value = + serde_json::from_str(&std::fs::read_to_string(&file).unwrap()).unwrap(); + println!("scenario {}", file.display()); + tokio::time::timeout(Duration::from_secs(15), run_scenario(scenario)) + .await + .unwrap_or_else(|_| panic!("scenario timed out: {}", file.display())); + } +} +async fn run_scenario(scenario: Value) { + let (host_tx, inbound) = mpsc::unbounded_channel(); + let (outbound, mut host_rx) = mpsc::unbounded_channel(); + let closed = Arc::new(Notify::new()); + let client = Client::builder() + .windowing( + WindowingConfig::default() + .window_bytes( + scenario["windowing"]["windowBytes"] + .as_u64() + .unwrap_or(262_144), + ) + .max_buffered_send_bytes( + scenario["windowing"]["maxBufferedSendBytes"] + .as_u64() + .unwrap_or(1_048_576), + ) + .maximum_frame_bytes( + scenario["windowing"]["maximumFrameBytes"] + .as_u64() + .unwrap_or(16_384), + ), + ) + .connect(MemoryTransport { + inbound, + outbound, + closed: closed.clone(), + }) + .await + .unwrap(); + let subscriptions: Arc>> = + Arc::new(Mutex::new(HashMap::new())); + let mut tasks = vec![]; + let mut deliveries: HashMap>> = HashMap::new(); + let mut manual = HashMap::new(); + let mut requests = VecDeque::new(); + let mut buffered = VecDeque::::new(); + let mut sent_totals = HashMap::::new(); + let mut frame_limits = HashMap::::new(); + for step in scenario["steps"].as_array().unwrap() { + let uri = step["channel"].as_str().unwrap_or("").to_string(); + match step["op"].as_str().unwrap() { + "subscribe" => { + manual.insert(uri.clone(), step["receive"] == "manual"); + let client = client.clone(); + let subscriptions = subscriptions.clone(); + let window = step.get("window").cloned(); + let coalescing = step["coalescing"].as_bool().unwrap_or(true); + tasks.push(tokio::spawn(async move { + if let Some(window) = window { + let mut options = WindowOptions::default().coalescing(coalescing); + if let Some(value) = window["windowBytes"].as_u64() { + options = options.window_bytes(value); + } + if let Some(value) = window["maximumFrameBytes"].as_u64() { + options = options.maximum_frame_bytes(value); + } + if let Ok((_, subscription)) = + client.subscribe_windowed(uri.clone(), options).await + { + subscriptions.lock().await.insert(uri, subscription); + } + } else { + let _ = client.subscribe(uri).await; + } + })); + } + "respond" => { + let position = requests + .iter() + .position(|request: &Value| request["method"] == step["to"]) + .unwrap(); + let request = requests.remove(position).unwrap(); + if let Some(limit) = step + .pointer("/result/flowControl/hostReceive/maximumFrameBytes") + .and_then(Value::as_u64) + { + frame_limits + .insert(request["params"]["channel"].as_str().unwrap().into(), limit); + } + host_tx + .send(TransportMessage::Text( + json!({"jsonrpc":"2.0","id":request["id"],"result":step["result"]}) + .to_string(), + )) + .unwrap(); + } + "hostSends" => { + host_tx + .send(TransportMessage::Text( + step["text"].as_str().unwrap().into(), + )) + .unwrap(); + } + "expectSent" => { + let deadline = tokio::time::Instant::now() + Duration::from_secs(1); + let mut joined = String::new(); + loop { + let value = if let Some(value) = buffered.pop_front() { + value + } else { + let text = tokio::time::timeout_at(deadline, host_rx.recv()) + .await + .unwrap_or_else(|_| panic!("missing sent message {step}")) + .unwrap(); + record(&text, &mut requests, &mut sent_totals) + }; + if value["method"] == "channel/frame" && step["reassemble"] == true { + let channel = value["params"]["channel"].as_str().unwrap(); + assert!( + ahp::windowing::utf16_bytes(&value.to_string()) + <= frame_limits[channel] + ); + joined.push_str(value["params"]["data"].as_str().unwrap()); + if value["params"]["final"] != true { + continue; + } + let joined: Value = serde_json::from_str(&joined).unwrap(); + assert!(matches(&joined, &step["match"]), "{joined} != {step}"); + break; + } + if matches(&value, &step["match"]) { + if let Some(absent) = step["absent"].as_array() { + for path in absent { + let pointer = + format!("/{}", path.as_str().unwrap().replace('.', "/")); + assert!(value.pointer(&pointer).is_none()); + } + } + break; + } + assert_eq!( + value["method"], "channel/credit", + "unexpected sent {value}; expected {step}" + ); + } + } + "expectNotSent" => { + assert!(!buffered.iter().any(|value| matches(value, &step["match"]))); + let deadline = tokio::time::Instant::now() + + Duration::from_millis(step["withinMs"].as_u64().unwrap()); + while let Ok(Some(text)) = tokio::time::timeout_at(deadline, host_rx.recv()).await { + let value = record(&text, &mut requests, &mut sent_totals); + assert!(!matches(&value, &step["match"]), "unexpected sent {value}"); + buffered.push_back(value); + } + } + "next" => { + if !manual[&uri] { + if let Some(previous) = deliveries.get_mut(&uri).and_then(|d| d.last_mut()) { + previous.take(); + } + } + let delivery = tokio::time::timeout(Duration::from_secs(1), async { + loop { + let mut subscriptions = subscriptions.lock().await; + if let Some(subscription) = subscriptions.get_mut(&uri) { + break subscription.next().await.unwrap().unwrap(); + } + drop(subscriptions); + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + let value = serde_json::to_value(delivery.notification()).unwrap(); + assert!( + matches(&value, &step["expect"]), + "delivery {value} != {step}" + ); + if let Some(bytes) = step["expectBytes"].as_u64() { + assert_eq!(delivery.wire_bytes(), bytes); + } + deliveries.entry(uri).or_default().push(Some(delivery)); + } + "unsubscribe" => { + client.unsubscribe(uri.clone()).await.unwrap(); + subscriptions.lock().await.remove(&uri); + sent_totals.remove(&uri); + } + "advance" => { + deliveries.get_mut(&uri).unwrap().last_mut().unwrap().take(); + } + "release" => { + deliveries.get_mut(&uri).unwrap()[step["index"].as_u64().unwrap() as usize] + .as_mut() + .unwrap() + .release(); + } + "send" => { + let sender = tokio::time::timeout(Duration::from_secs(1), async { + loop { + if let Some(subscription) = subscriptions.lock().await.get(&uri) { + break subscription.sender(); + } + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + let method = step["method"].as_str().unwrap().to_string(); + let params = step["params"].clone(); + tasks.push(tokio::spawn(async move { + sender.send(&method, params).await.unwrap(); + })); + tokio::task::yield_now().await; + } + "hostCreditsAll" => { + while let Ok(text) = host_rx.try_recv() { + buffered.push_back(record(&text, &mut requests, &mut sent_totals)); + } + host_tx.send(TransportMessage::Text(json!({"jsonrpc":"2.0","method":"channel/credit","params":{"channel":uri,"consumedBytes":sent_totals.get(&uri).copied().unwrap_or(0)}}).to_string())).unwrap(); + } + "expectStatus" => { + tokio::time::timeout(Duration::from_secs(1), async { + loop { + if let Some(subscription) = subscriptions.lock().await.get(&uri) { + let status = serde_json::to_value(subscription.status()).unwrap(); + if matches(&status, &step["match"]) { + break; + } + } + tokio::task::yield_now().await; + } + }) + .await + .unwrap_or_else(|_| panic!("status mismatch {step}")); + } + "expectClosed" => { + tokio::time::timeout(Duration::from_secs(1), closed.notified()) + .await + .unwrap(); + } + other => panic!("unsupported op {other}"), + } + } + client.shutdown().await; + for task in tasks { + task.abort(); + } +} +fn record(text: &str, requests: &mut VecDeque, totals: &mut HashMap) -> Value { + let value: Value = serde_json::from_str(text).unwrap(); + if value.get("id").is_some() && value.get("method").is_some() { + requests.push_back(value.clone()); + } else if let Some(method) = value["method"].as_str() { + if ahp_types::notifications::counts_against_window(method) { + if let Some(channel) = value["params"]["channel"].as_str() { + *totals.entry(channel.into()).or_default() += ahp::windowing::utf16_bytes(text); + } + } + } + value +} + +#[tokio::test] +async fn managed_recovery_retains_queue_and_rebinds_sender() { + for recovery in [ + "snapshot", + "replay", + "live", + "dropped", + "consumer-drop", + "resubscribe-failed", + "resubscribe-hung", + "subscribe-hung", + ] { + managed_recovery(recovery).await; + } +} + +async fn managed_recovery(recovery: &'static str) { + use ahp::hosts::{HostConfig, HostId, MultiHostClient, ReconnectPolicy}; + use ahp::transport::BoxedTransport; + use std::sync::atomic::{AtomicUsize, Ordering}; + let attempts = Arc::new(AtomicUsize::new(0)); + let loss = Arc::new(Notify::new()); + let reconnect_gate = Arc::new(Notify::new()); + let records = Arc::new(Mutex::new(Vec::<(usize, Value)>::new())); + let recovery_text = recovery; + let factory = { + let attempts = attempts.clone(); + let loss = loss.clone(); + let reconnect_gate = reconnect_gate.clone(); + let records = records.clone(); + move |_| { + let attempt = attempts.fetch_add(1, Ordering::Relaxed); + let loss = loss.clone(); + let reconnect_gate = reconnect_gate.clone(); + let records = records.clone(); + async move { + if attempt > 0 { + reconnect_gate.notified().await; + } + let (in_tx, inbound) = mpsc::unbounded_channel(); + let (outbound, mut out_rx) = mpsc::unbounded_channel::(); + tokio::spawn(async move { + let limits = json!({"clientReceive":{"windowBytes":1024,"maximumFrameBytes":1024},"hostReceive":{"windowBytes":128,"maximumFrameBytes":1024}}); + loop { + let text = tokio::select! { + _ = loss.notified(), if attempt == 0 => break, + message = out_rx.recv() => match message { Some(text) => text, None => break }, + }; + let request: Value = serde_json::from_str(&text).unwrap(); + records.lock().await.push((attempt, request.clone())); + if request.get("id").is_none() { + continue; + } + let resubscribing = recovery_text.starts_with("resubscribe-"); + if (request["method"] == "reconnect" && resubscribing) + || (request["method"] == "subscribe" + && attempt > 0 + && recovery_text == "resubscribe-failed") + { + in_tx.send(TransportMessage::Text(json!({"jsonrpc":"2.0","id":request["id"],"error":{"code":-32602,"message":"refused"}}).to_string())).unwrap(); + continue; + } + if request["method"] == "subscribe" + && (recovery_text == "subscribe-hung" + || (attempt > 0 && recovery_text == "resubscribe-hung")) + { + continue; + } + let root = + json!({"resource":"ahp-root://","state":{"agents":[]},"fromSeq":1}); + let uri = request["params"]["channel"].as_str().unwrap(); + let result = match request["method"].as_str().unwrap() { + "initialize" => { + json!({"protocolVersion":"0.9.0","serverSeq":1,"snapshots":[root]}) + } + "listSessions" => json!({"items":[]}), + "subscribe" => json!({"flowControl":limits}), + "reconnect" => { + if request["params"].get("windows").is_none() { + in_tx.send(TransportMessage::Text(json!({"jsonrpc":"2.0","id":request["id"],"result":{"type":"snapshot","snapshots":[root]}}).to_string())).unwrap(); + continue; + } + let windows = if recovery_text == "dropped" { + json!([]) + } else { + json!([{"channel":"ahp-terminal:/retained","flowControl":limits,"recovery":recovery_text}]) + }; + json!({"type":"snapshot","snapshots":[root],"windows":{"items":windows}}) + } + _ => Value::Null, + }; + in_tx + .send(TransportMessage::Text( + json!({"jsonrpc":"2.0","id":request["id"],"result":result}) + .to_string(), + )) + .unwrap(); + if request["method"] == "subscribe" { + let snapshot = json!({"jsonrpc":"2.0","method":"channel/snapshot","params":{"channel":uri,"snapshot":{"resource":uri,"state":{},"fromSeq":1}}}); + in_tx + .send(TransportMessage::Text(snapshot.to_string())) + .unwrap(); + if !uri.starts_with("ahp-tcp:") { + in_tx.send(TransportMessage::Text(json!({"jsonrpc":"2.0","method":"action","params":{"channel":uri,"serverSeq":2,"action":{"type":"terminal/data","data":"retained"}}}).to_string())).unwrap(); + } + } else if request["method"] == "reconnect" + && matches!(recovery_text, "snapshot" | "replay" | "live") + { + let value = if recovery_text == "snapshot" { + json!({"jsonrpc":"2.0","method":"channel/snapshot","params":{"channel":"ahp-terminal:/retained","snapshot":{"resource":"ahp-terminal:/retained","state":{},"fromSeq":3}}}) + } else { + json!({"jsonrpc":"2.0","method":"action","params":{"channel":"ahp-terminal:/retained","serverSeq":3,"action":{"type":"terminal/data","data":"new"}}}) + }; + in_tx + .send(TransportMessage::Text(value.to_string())) + .unwrap(); + } + } + }); + Ok(BoxedTransport::new(MemoryTransport { + inbound, + outbound, + closed: Arc::new(Notify::new()), + })) + } + } + }; + let host_id = HostId::new("window-host"); + let config = HostConfig::new(host_id.clone(), "window host", factory) + .with_client_config(ahp::ClientConfig { + default_request_timeout: None, + ..Default::default() + }) + .with_reconnect_policy(ReconnectPolicy::immediate_forever()) + .with_windowing( + WindowingConfig::default() + .window_bytes(1024) + .maximum_frame_bytes(1024), + ); + let (multi, _) = MultiHostClient::single(config).await.unwrap(); + tokio::time::timeout(Duration::from_secs(1), async { + while multi.host(&host_id).await.unwrap().generation == 0 { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + let mut telemetry = multi.flow_control_events(); + if recovery == "subscribe-hung" { + let child = multi.clone(); + let child_id = host_id.clone(); + let subscribe = tokio::spawn(async move { + child + .subscribe_windowed(&child_id, "ahp-terminal:/hung", WindowOptions::default()) + .await + }); + tokio::time::timeout(Duration::from_secs(1), async { + while !records + .lock() + .await + .iter() + .any(|(_, request)| request["method"] == "subscribe") + { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + tokio::time::timeout(Duration::from_secs(1), multi.remove_host(&host_id)) + .await + .unwrap() + .unwrap(); + assert!(subscribe.await.unwrap().is_err()); + return; + } + let (_, mut subscription) = multi + .subscribe_windowed(&host_id, "ahp-terminal:/retained", WindowOptions::default()) + .await + .unwrap(); + assert_eq!(subscription.host_id(), Some(&host_id)); + let old_snapshot = subscription.next().await.unwrap().unwrap(); + let old_action = subscription.next().await.unwrap().unwrap(); + let (_, mut tcp) = multi + .subscribe_windowed(&host_id, "ahp-tcp:/ephemeral", WindowOptions::default()) + .await + .unwrap(); + let old_tcp = tcp.next().await.unwrap().unwrap(); + if recovery == "consumer-drop" { + drop(subscription); + tokio::time::timeout(Duration::from_secs(1), async { + while multi + .host(&host_id) + .await + .unwrap() + .subscriptions + .contains(&"ahp-terminal:/retained".into()) + { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + loss.notify_one(); + tokio::time::timeout(Duration::from_secs(1), async { + while attempts.load(Ordering::Relaxed) < 2 { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + reconnect_gate.notify_one(); + tokio::time::timeout(Duration::from_secs(1), async { + while multi.host(&host_id).await.unwrap().generation < 2 { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + let records = records.lock().await; + let request = &records + .iter() + .find(|(attempt, request)| *attempt == 1 && request["method"] == "reconnect") + .unwrap() + .1; + assert!(request["params"].get("windows").is_none()); + assert!(!request["params"]["subscriptions"] + .as_array() + .unwrap() + .iter() + .any(|uri| uri == "ahp-terminal:/retained")); + drop(records); + drop(old_snapshot); + drop(old_action); + drop(old_tcp); + multi.remove_host(&host_id).await.unwrap(); + return; + } + loss.notify_one(); + tokio::time::timeout(Duration::from_secs(1), async { + while attempts.load(Ordering::Relaxed) < 2 { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + assert!( + tokio::time::timeout(Duration::from_millis(10), subscription.next()) + .await + .is_err() + ); + assert!( + tokio::time::timeout(Duration::from_millis(10), subscription.wait_writable()) + .await + .is_err() + ); + reconnect_gate.notify_one(); + if recovery == "resubscribe-hung" { + tokio::time::timeout(Duration::from_secs(1), async { + while !records + .lock() + .await + .iter() + .any(|(attempt, request)| *attempt == 1 && request["method"] == "subscribe") + { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + tokio::time::timeout(Duration::from_secs(1), multi.remove_host(&host_id)) + .await + .unwrap() + .unwrap(); + return; + } + tokio::time::timeout(Duration::from_secs(1), async { + while multi.host(&host_id).await.unwrap().generation < 2 { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + drop(old_snapshot); + drop(old_action); + drop(old_tcp); + assert!(tcp.next().await.unwrap().is_none()); + let recorded = records.lock().await.clone(); + let reconnect = recorded + .iter() + .find(|(attempt, value)| *attempt == 1 && value["method"] == "reconnect") + .unwrap() + .1 + .clone(); + assert_eq!( + reconnect["params"]["windows"]["items"][0]["lastAppliedServerSeq"], + 2 + ); + assert!(!reconnect["params"]["subscriptions"] + .as_array() + .unwrap() + .iter() + .any(|uri| uri.as_str().unwrap().starts_with("ahp-tcp:"))); + if recovery == "dropped" || recovery == "resubscribe-failed" { + assert!(subscription.next().await.is_err()); + assert!(!multi + .host(&host_id) + .await + .unwrap() + .subscriptions + .contains(&"ahp-terminal:/retained".into())); + if recovery == "resubscribe-failed" { + assert_eq!(attempts.load(Ordering::Relaxed), 2); + assert!(multi.host(&host_id).await.unwrap().state.is_connected()); + multi + .client(&host_id) + .await + .unwrap() + .raw_client() + .ping() + .await + .unwrap(); + } + } else { + let delivery = tokio::time::timeout(Duration::from_secs(1), subscription.next()) + .await + .unwrap() + .unwrap() + .unwrap(); + assert_eq!( + subscription.status().buffered_receive_bytes, + delivery.wire_bytes() + ); + let new_bytes = delivery.wire_bytes(); + drop(delivery); + subscription + .sender() + .send( + "custom/write", + json!({"channel":"ahp-terminal:/retained","value":1}), + ) + .await + .unwrap(); + tokio::time::timeout(Duration::from_secs(1), async { + loop { + if records.lock().await.iter().any(|(attempt, value)| { + *attempt == 1 + && value["method"] == "channel/credit" + && value["params"]["consumedBytes"] == new_bytes + }) { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + assert!(records + .lock() + .await + .iter() + .any(|(attempt, value)| *attempt == 1 && value["method"] == "custom/write")); + } + let event = tokio::time::timeout(Duration::from_secs(1), telemetry.recv()) + .await + .unwrap() + .unwrap(); + assert_eq!(event.host_id, host_id); + multi.remove_host(&host_id).await.unwrap(); +} diff --git a/clients/rust/crates/ahp/tests/windowing_lifecycle.rs b/clients/rust/crates/ahp/tests/windowing_lifecycle.rs new file mode 100644 index 000000000..a463972da --- /dev/null +++ b/clients/rust/crates/ahp/tests/windowing_lifecycle.rs @@ -0,0 +1,369 @@ +#![allow(clippy::unwrap_used, clippy::panic)] +use ahp::{ + Client, Transport, TransportError, TransportMessage, WindowOptions, WindowedSubscription, +}; +use serde_json::{json, Value}; +use std::sync::{ + atomic::{AtomicBool, Ordering}, + Arc, +}; +use std::time::Duration; +use tokio::sync::{mpsc, Notify}; + +struct Memory { + inbound: mpsc::UnboundedReceiver, + outbound: mpsc::UnboundedSender, +} +impl Transport for Memory { + fn window_capabilities(&self) -> ahp::WindowTransportCapabilities { + ahp::WindowTransportCapabilities::default().raw_text(true) + } + async fn send(&mut self, message: TransportMessage) -> Result<(), TransportError> { + let TransportMessage::Text(text) = message else { + panic!("expected text") + }; + self.outbound.send(text).map_err(|_| TransportError::Closed) + } + async fn recv(&mut self) -> Result, TransportError> { + Ok(self.inbound.recv().await) + } +} +struct Script { + tx: mpsc::UnboundedSender, + rx: mpsc::UnboundedReceiver, +} +impl Script { + async fn connect(timeout: Option) -> (Client, Self) { + Self::configured(timeout, Default::default()).await + } + async fn configured(timeout: Option, config: ahp::WindowingConfig) -> (Client, Self) { + let (tx, inbound) = mpsc::unbounded_channel(); + let (outbound, rx) = mpsc::unbounded_channel(); + let client = Client::builder() + .default_request_timeout(timeout) + .windowing(config) + .connect(Memory { inbound, outbound }) + .await + .unwrap(); + (client, Self { tx, rx }) + } + async fn next(&mut self) -> Value { + serde_json::from_str( + &tokio::time::timeout(Duration::from_secs(1), self.rx.recv()) + .await + .unwrap() + .unwrap(), + ) + .unwrap() + } + fn send(&self, value: Value) { + self.tx + .send(TransportMessage::Text(value.to_string())) + .unwrap(); + } + fn accept(&self, request: &Value) { + self.send( + json!({"jsonrpc":"2.0","id":request["id"],"result":{"flowControl":{ + "clientReceive":{"windowBytes":1024,"maximumFrameBytes":1024}, + "hostReceive":{"windowBytes":1024,"maximumFrameBytes":1024} + }}}), + ); + } + async fn subscribe(&mut self, client: &Client, uri: &str) -> WindowedSubscription { + let client = client.clone(); + let uri = uri.to_string(); + let pending = tokio::spawn(async move { + client + .subscribe_windowed(uri, WindowOptions::default()) + .await + }); + let request = self.next().await; + assert_eq!(request["method"], "subscribe"); + self.accept(&request); + pending.await.unwrap().unwrap().1 + } + async fn ping(&mut self, client: &Client) { + let client = client.clone(); + let pending = tokio::spawn(async move { client.ping().await }); + let request = self.next().await; + assert_eq!(request["method"], "ping"); + self.send(json!({"jsonrpc":"2.0","id":request["id"],"result":null})); + pending.await.unwrap().unwrap(); + } +} + +#[tokio::test] +async fn failed_cancelled_and_timed_out_subscribes_do_not_leave_dead_queues() { + for failure in ["rpc", "cancel", "timeout"] { + let (client, mut host) = Script::connect(Some(Duration::from_millis(30))).await; + let child = client.clone(); + let pending = tokio::spawn(async move { + child + .subscribe_windowed("ahp-terminal:/failed", WindowOptions::default()) + .await + }); + let request = host.next().await; + match failure { + "rpc" => { + host.send(json!({"jsonrpc":"2.0","id":request["id"],"error":{"code":-32602,"message":"refused"}})); + assert!(pending.await.unwrap().is_err()); + } + "cancel" => { + pending.abort(); + assert!(pending.await.is_err()); + } + _ => { + assert!(pending.await.unwrap().is_err()); + } + } + client + .notify( + "custom/after", + json!({"channel":"ahp-terminal:/failed","value":1}), + ) + .await + .unwrap(); + assert_eq!(host.next().await["method"], "custom/after"); + if failure != "rpc" { + host.accept(&request); + assert_eq!(host.next().await["method"], "unsubscribe"); + } + host.ping(&client).await; + client.shutdown().await; + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn unsubscribe_racing_subscribe_and_drop_followed_by_subscribe_are_nonfatal() { + let (client, mut host) = Script::connect(None).await; + let child = client.clone(); + let pending = tokio::spawn(async move { + child + .subscribe_windowed("ahp-terminal:/race", WindowOptions::default()) + .await + }); + let request = host.next().await; + client + .unsubscribe("ahp-terminal:/race".into()) + .await + .unwrap(); + assert_eq!(host.next().await["method"], "unsubscribe"); + host.accept(&request); + assert_eq!(host.next().await["method"], "unsubscribe"); + assert!(pending.await.unwrap().is_err()); + let mut sub = host.subscribe(&client, "ahp-terminal:/race").await; + for _ in 0..30 { + drop(sub); + let child = client.clone(); + let pending = tokio::spawn(async move { + child + .subscribe_windowed("ahp-terminal:/race", WindowOptions::default()) + .await + }); + let mut request = host.next().await; + while request["method"] == "unsubscribe" { + request = host.next().await; + } + assert_eq!(request["method"], "subscribe"); + host.accept(&request); + sub = pending.await.unwrap().unwrap().1; + } + host.ping(&client).await; + client.shutdown().await; +} + +#[tokio::test] +async fn late_resubscribe_and_reconnect_results_install_without_closing() { + let (client, mut host) = Script::connect(Some(Duration::from_millis(30))).await; + let sub = host.subscribe(&client, "ahp-terminal:/late").await; + let child = client.clone(); + let pending = tokio::spawn(async move { + child + .subscribe_windowed( + "ahp-terminal:/late", + WindowOptions::default().window_bytes(2048), + ) + .await + }); + let request = host.next().await; + assert!(pending.await.unwrap().is_err()); + host.accept(&request); + host.ping(&client).await; + assert_eq!(sub.status().mode, "windowed"); + let child = client.clone(); + let reconnect = tokio::spawn(async move { + child.reconnect_with_windows(serde_json::from_value(json!({ + "channel":"ahp-root://","clientId":"late","lastSeenServerSeq":0, + "subscriptions":["ahp-terminal:/late"], + "windows":{"items":[{"channel":"ahp-terminal:/late","flowControl":{"receive":{"windowBytes":2048,"maximumFrameBytes":1024}},"lastAppliedServerSeq":0}]} + })).unwrap()).await + }); + let request = host.next().await; + assert!(reconnect.await.unwrap().is_err()); + host.send(json!({"jsonrpc":"2.0","id":request["id"],"result":{"type":"replay","actions":[],"missing":[],"windows":{"items":[{ + "channel":"ahp-terminal:/late","flowControl":{ + "clientReceive":{"windowBytes":1024,"maximumFrameBytes":1024}, + "hostReceive":{"windowBytes":1024,"maximumFrameBytes":1024} + },"recovery":"replay"}]}}})); + host.ping(&client).await; + client.shutdown().await; +} + +#[tokio::test] +async fn late_initialize_installs_without_closing() { + let (client, mut host) = Script::connect(Some(Duration::from_millis(30))).await; + let _sub = host.subscribe(&client, "ahp-terminal:/late").await; + let child = client.clone(); + let initialize = tokio::spawn(async move { + child + .initialize( + "late".into(), + vec!["1.0.0".into()], + vec!["ahp-terminal:/late".into()], + ) + .await + }); + let request = host.next().await; + assert!(initialize.await.unwrap().is_err()); + host.send(json!({"jsonrpc":"2.0","id":request["id"],"result":{ + "protocolVersion":"1.0.0","serverSeq":1, + "snapshots":[{"resource":"ahp-terminal:/late","state":{},"fromSeq":1}] + }})); + host.ping(&client).await; + client.shutdown().await; +} + +#[tokio::test] +async fn control_traffic_bypasses_a_full_one_byte_send_budget() { + let (client, mut host) = Script::configured( + None, + ahp::WindowingConfig::default().max_buffered_send_bytes(1), + ) + .await; + let child = client.clone(); + let subscribe = tokio::spawn(async move { + child + .subscribe_windowed("ahp-tcp:/control", WindowOptions::default()) + .await + }); + let request = host.next().await; + host.send( + json!({"jsonrpc":"2.0","id":request["id"],"result":{"flowControl":{ + "clientReceive":{"windowBytes":1024,"maximumFrameBytes":1024}, + "hostReceive":{"windowBytes":1,"maximumFrameBytes":1024} + }}}), + ); + let mut sub = subscribe.await.unwrap().unwrap().1; + host.send(json!({"jsonrpc":"2.0","method":"channel/snapshot","params":{"channel":sub.uri(),"snapshot":{"resource":sub.uri(),"state":{},"fromSeq":0}}})); + sub.next().await.unwrap().unwrap().release(); + let credit = host.next().await; + assert_eq!(credit["method"], "channel/credit"); + assert!(credit["params"]["consumedBytes"].as_u64().unwrap() > 1); + let sender = sub.sender(); + let first_sender = sender.clone(); + let first = + tokio::spawn(async move { first_sender.send("tcp/data", json!({"data":"eA=="})).await }); + assert_eq!(host.next().await["method"], "tcp/data"); + first.await.unwrap().unwrap(); + assert!(!sender.status().writable); + let data_sender = sender.clone(); + let blocked = + tokio::spawn(async move { data_sender.send("tcp/data", json!({"data":"eA=="})).await }); + tokio::task::yield_now().await; + let control = tokio::spawn(async move { sender.send("tcp/reset", json!({})).await }); + assert_eq!(host.next().await["method"], "tcp/reset"); + control.await.unwrap().unwrap(); + client.unsubscribe("ahp-tcp:/control".into()).await.unwrap(); + assert_eq!(host.next().await["method"], "unsubscribe"); + assert!(blocked.await.unwrap().is_err()); + host.ping(&client).await; + client.shutdown().await; +} + +#[tokio::test] +async fn ordinary_messages_are_drained_before_shutdown() { + let (client, mut host) = Script::connect(None).await; + let mut events = client.events(); + let mut subscription = client.attach_subscription("ahp-root://").await; + for index in 0..100 { + client + .notify( + "ordinary/data", + json!({"channel":"ahp-root://","index":index}), + ) + .await + .unwrap(); + } + client + .dispatch( + "ahp-root://".into(), + ahp::ahp_types::actions::StateAction::RootActiveSessionsChanged( + ahp::ahp_types::actions::RootActiveSessionsChangedAction { active_sessions: 1 }, + ), + ) + .await + .unwrap(); + client.unsubscribe("ahp-root://".into()).await.unwrap(); + client.shutdown().await; + for index in 0..100 { + assert_eq!(host.next().await["params"]["index"], index); + } + assert_eq!(host.next().await["method"], "dispatchAction"); + assert_eq!(host.next().await["method"], "unsubscribe"); + assert!(events.recv().await.is_none()); + assert!(subscription.recv().await.is_none()); +} + +struct Stalled { + entered: Arc, + dropped: Arc, +} +impl Drop for Stalled { + fn drop(&mut self) { + self.dropped.store(true, Ordering::Release); + } +} +impl Transport for Stalled { + async fn send(&mut self, _: TransportMessage) -> Result<(), TransportError> { + self.entered.notify_one(); + std::future::pending().await + } + async fn recv(&mut self) -> Result, TransportError> { + std::future::pending().await + } +} +#[tokio::test(start_paused = true)] +async fn forced_shutdown_with_full_outbound_queue_still_tears_down_readers() { + let entered = Arc::new(Notify::new()); + let dropped = Arc::new(AtomicBool::new(false)); + let client = Client::connect( + Stalled { + entered: entered.clone(), + dropped: dropped.clone(), + }, + Default::default(), + ) + .await + .unwrap(); + let mut events = client.events(); + let mut subscription = client.attach_subscription("ahp-root://").await; + client + .notify("ordinary/data", json!({"channel":"ahp-root://"})) + .await + .unwrap(); + entered.notified().await; + for _ in 0..64 { + client + .notify("ordinary/data", json!({"channel":"ahp-root://"})) + .await + .unwrap(); + } + let child = client.clone(); + let shutdown = tokio::spawn(async move { child.shutdown().await }); + tokio::task::yield_now().await; + tokio::time::advance(Duration::from_secs(11)).await; + shutdown.await.unwrap(); + assert!(dropped.load(Ordering::Acquire)); + assert!(events.recv().await.is_none()); + assert!(subscription.recv().await.is_none()); +} diff --git a/clients/swift/AgentHostProtocol/README.md b/clients/swift/AgentHostProtocol/README.md index bab640157..0eb03f095 100644 --- a/clients/swift/AgentHostProtocol/README.md +++ b/clients/swift/AgentHostProtocol/README.md @@ -78,6 +78,161 @@ for snapshot in initialized.snapshots { `AHPStateMirror` is a convenience for simple consumers. Larger apps can keep their own state store and route snapshots/actions through the generated reducers directly. +## Windowed delivery and backpressure + +Windowing is off by default. Connection configuration supplies defaults but +does not enable it. Existing `subscribe(uri, delivery:view:)` calls are unchanged. +On a connected client, opt in with a managed subscription: + +```swift +let subscription = try await client.subscribe( + "ahp-terminal:/t1", + options: .windowed(windowBytes: 262_144) +) +let result = await subscription.result +for try await delivery in subscription.events { + await mirror.apply(delivery: delivery) + // Auto mode releases this delivery when the loop asks for the next one. +} +``` + +The first delivery is the snapshot: `channel/snapshot` after acceptance, or +the ordinary result's snapshot when the host declines (if that snapshot exists). +There is no protocol-version gate. A parsed-only custom transport falls back +to ordinary delivery with status `fallbackReason == "transport"`. A host +declining the offer reports `"declined"`. TCP requires acceptance and never +falls back. Ordinary receive queues retain their existing unbounded behavior. + +For manual credit, use `.windowed(receive: .manual)` and `deliveries`: + +```swift +let subscription = try await client.subscribe( + "ahp-terminal:/t1", options: .windowed(receive: .manual) +) +for try await delivery in subscription.deliveries { + await mirror.apply(delivery: delivery) + await delivery.release() // idempotent; call after processing, not on receipt +} +``` + +For terminal or TCP byte sinks, release after bytes have been written or accepted +by a bounded application buffer, not merely handed to another unbounded queue. +Non-final protocol fragments are credited when stored so that messages larger +than the window can complete; the final fragment is held until consumer release. +There is no additional logical-message or reassembly cap. Stop with +`try await subscription.close()`; breaking an auto loop alone does not release +its last item. Re-subscribe or `subscription.update(options:)` changes the same +handle. Windowed-to-windowed counters continue; changing mode resets counters. + +Send methods await transport completion, not server processing: + +```swift +try await subscription.waitUntilWritable() // advisory, not a reservation +try await subscription.send( + method: "dispatchAction", + params: DispatchActionParams( + channel: subscription.channel, clientSeq: 1, + action: .terminalInput(.init(type: .terminalInput, data: "hello\n")) + ) +) +let status = await subscription.status() +print(status.outstandingSendBytes, status.bufferedSendBytes) +``` + +The per-subscription queue defaults to 1 MiB for data waiting behind another +message; exceeding that queued allowance throws `AHPWindowingError.backpressure`. +A single message is accepted when the queue is empty even if it is larger than +the allowance: fragmentation and credit pace its transmission. Controls, +scheduler frames, and RPCs are never rejected by this user send cap. +Cancellation before emission removes the +send. Once fragmentation starts, the message is completed. Sends/dispatches +through the original client API also respect accepted windows. Requests and +generated controls (`channel/credit`, `unsubscribe`, `tcp/reset`) bypass full +data windows. Observe `subscription.statusChanges` and +`await client.flowControlEvents` without granting credit. + +Configure defaults through `AHPClientConfig(windowing:)` or +`HostConfig.withWindowing(_:)`. All protocol sizes count the **complete raw JSON +text** as `text.utf16.count * 2`, not its UTF-8 transport payload: + +```swift +let configuration = AHPClientConfig(windowing: AHPWindowingConfig( + windowBytes: 262_144, maximumFrameBytes: 16_384, + maxBufferedSendBytes: 1_048_576 +)) +``` + +When no frame override is supplied, URLSession derives the default from the +actual task's `maximumMessageSize` as `2 * (limit / 3)`: **699,050** UTF-16 +bytes for the usual 1-MiB payload limit. The native NWConnection transport has +no message cap (its 65,536-byte read chunk is not one), so it uses the +**1,048,576**-byte fallback. Custom transports retain source compatibility: +adopt `AHPWindowingTransport` and declare raw-text capabilities to opt in; +parsed frames need an original JSON text envelope, never reconstructed text. + +Coalescing is enabled by default only on accepted windowed subscriptions. +Adjacent queued terminal output and same-part/same-metadata chat deltas merge +locally; queued outbound TCP data merges decoded bytes. No latency preference +is added. Origins, clientSeq sends, snapshots, EOF/reset and epochs are barriers. +Disable with `.windowed(coalescing: false)`. Custom coalescers are registered +per connection and run before the built-ins: + +```swift +let terminalCoalescer = AHPCoalescer(method: "action") { earlier, later, context in + guard context.direction == "inbound" else { return nil } + var first = try earlier.decode(ActionEnvelope.self) + let second = try later.decode(ActionEnvelope.self) + guard case .terminalData(let a) = first.action, + case .terminalData(let b) = second.action else { return nil } + first.action = .terminalData(.init(type: .terminalData, data: a.data + b.data)) + first.serverSeq = second.serverSeq + return try AHPNotificationMessage(method: "action", params: first) +} +let configuration = AHPClientConfig( + windowing: .init(coalescers: [terminalCoalescer]) +) +``` + +A custom merge must preserve channel-required events and state. The library +checks channel/method identity and control exemptions; received credit still +uses the sum of original wire sizes, not the merged serialization. + +Multi-host consumers use the same shape: +`multi.subscribe(host: hostId, uri: uri, options: .windowed())`. Its owner queue +survives reconnect, with per-channel retained sequence checkpoints and fresh +counters. Old deliveries remain readable but never credit the new connection. +Snapshot/replay/live recovery is per channel; dropped handles throw +`subscriptionDropped`. TCP streams close on connection loss and are never +resumed. `multi.flowControlEvents()` tags advisory flow events with a host ID. +Cancellation and request timeouts end only the caller's wait; a late negotiation +response is still processed in order without disconnecting the host. Unaccepted +subscriptions created by an abandoned call are removed, and late acceptance is +unsubscribed. Closing a handle cancels its credit timer and removes its replay +registration. A failed re-subscribe during reconnect closes only that handle, +not the other subscriptions or the host connection. +State mirrors can consume `apply(delivery:)` / `apply(host:delivery:)`; these +helpers never release credit implicitly. The sample app is not changed. + +### Examples + +Two runnable examples use an in-memory demo host that follows the window rules, +so they need no server. They live in a separate local package, +[`clients/swift/Examples`](../Examples), so the published package is unchanged: + +- [`WindowingBasic`](../Examples/Sources/WindowingBasic/main.swift) opts a + terminal subscription into windowed delivery and shows the host pausing when + the window is full and resuming as the client returns credit. +- [`WindowingBackpressure`](../Examples/Sources/WindowingBackpressure/main.swift) + adds a slow renderer, a custom coalescer that keeps only the latest terminal + resize, and a paste that waits for `waitUntilWritable()` instead of queueing. + +From the repository root: + +```sh +swift run --package-path clients/swift/Examples WindowingBasic +swift run --package-path clients/swift/Examples WindowingBackpressure +``` + ## Multi-Host Client Use `MultiHostClient` when one app talks to more than one AHP host, or when you want the same supervisor model for a single host. It owns per-host transport creation, reconnect backoff, stable `clientId` lookup, event fan-in, session-summary caches, generation-checked client handles, and deterministic aggregated views. diff --git a/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/AHPClient.swift b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/AHPClient.swift index 7d84dba5b..67fed54cf 100644 --- a/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/AHPClient.swift +++ b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/AHPClient.swift @@ -72,7 +72,29 @@ public actor AHPClient { /// Outbound channel carries raw JSON bytes. Going via JSONSerialization /// preserves the NSNumber Bool/Int distinction that AnyCodable would /// otherwise erase on Apple platforms. - private var outboundContinuation: AsyncStream.Continuation? + private struct OutboundWrite: Sendable { + let data: Data + let completion: UUID? + let beforeWrite: AHPBeforeWrite? + init(data: Data, completion: UUID?, beforeWrite: AHPBeforeWrite? = nil) { + self.data = data + self.completion = completion + self.beforeWrite = beforeWrite + } + } + private var outboundContinuation: AsyncStream.Continuation? + private var windowingEngine: AHPWindowingEngine? + private var requestContexts: [Int: (String, Data)] = [:] + private var requestOptions: [Int: AHPSubscriptionOptions] = [:] + private var flowNegotiations: Set = [] + private struct SubscriptionRequest { + let owner: AHPSubscription + let created: Bool + } + private var requestSubscriptions: [Int: SubscriptionRequest] = [:] + private var abandonedRequests: Set = [] + private var windowWrites: [UUID: CheckedContinuation] = [:] + private var callerWrites: Set = [] private var writerTask: Task? private var keepAliveTask: Task? @@ -119,15 +141,25 @@ public actor AHPClient { // from inside the actor and the writer Task drains and serializes them // to `transport.send`. This avoids two `await transport.send` calls // racing under actor reentrancy. - var cont: AsyncStream.Continuation! - let outbound = AsyncStream(bufferingPolicy: .unbounded) { c in + var cont: AsyncStream.Continuation! + let outbound = AsyncStream(bufferingPolicy: .unbounded) { c in cont = c } self.outboundContinuation = cont let transport = self.transport self.writerTask = Task { [weak self] in - for await data in outbound { - guard let text = String(data: data, encoding: .utf8) else { + for await packet in outbound { + if let before = packet.beforeWrite { + do { try await before() } + catch { + if let id = packet.completion { await self?.finishWindowWrite(id, error: error) } + continue + } + } + guard let text = String(data: packet.data, encoding: .utf8) else { + if let id = packet.completion { + await self?.finishWindowWrite(id, error: TransportError.protocol("invalid UTF-8")) + } await self?.handleTransportFailure( TransportError.protocol("outbound bytes are not valid UTF-8") ) @@ -135,7 +167,10 @@ public actor AHPClient { } do { try await transport.send(.text(text)) + if let id = packet.completion { await self?.finishWindowWrite(id, error: nil) } } catch { + if let id = packet.completion { await self?.finishWindowWrite(id, error: error) } + if case AHPWindowingError.backpressure = error { continue } await self?.handleTransportFailure(error) return } @@ -158,13 +193,9 @@ public actor AHPClient { public func shutdown() async { if didShutdown { return } didShutdown = true - - // Stop the writer (no more sends). Order matters: finishing the - // outbound stream lets the writer loop exit naturally; cancelling the - // receive task aborts the inbound loop. outboundContinuation?.finish() outboundContinuation = nil - writerTask?.cancel() + let writer = writerTask writerTask = nil keepAliveTask?.cancel() keepAliveTask = nil @@ -172,9 +203,19 @@ public actor AHPClient { receiveTask?.cancel() receiveTask = nil + await windowingEngine?.disconnect(AHPClientError.shutdown) + failAllPending(with: .shutdown) + // Preserve already-enqueued ordinary notifications on graceful close. + // Closing a stalled transport also bounds teardown of a pending write. + let transport = self.transport + let watchdog = Task { + try? await Task.sleep(for: .seconds(1)) + if !Task.isCancelled { try? await transport.close() } + } + await writer?.value + watchdog.cancel() try? await transport.close() - failAllPending(with: .shutdown) finishAllSubscriptions() finishAllEventListeners() // Final state transition (also fans out). @@ -260,6 +301,9 @@ public actor AHPClient { protocolVersions: [String], initialSubscriptions: [String] = [] ) async throws -> InitializeResult { + guard !initialSubscriptions.contains(where: { $0.hasPrefix("ahp-tcp:") }) else { + throw TransportError.protocol("TCP cannot use initialSubscriptions") + } let params = InitializeParams( channel: RootResourceURI, protocolVersions: protocolVersions, @@ -305,6 +349,197 @@ public actor AHPClient { return result } + @discardableResult + public func reconnect(clientId: String, lastSeenServerSeq: Int, subscriptions: [String], + windows: ChannelDeliveryResumeOptions) async throws -> ReconnectResult { + let result: ReconnectResult = try await request(method: "reconnect", params: ReconnectParams( + channel: RootResourceURI, clientId: clientId, lastSeenServerSeq: lastSeenServerSeq, + subscriptions: subscriptions, windows: windows)) + switch result { + case .replay(let value): self.lastSeenServerSeq = max(self.lastSeenServerSeq, value.actions.map(\.serverSeq).max() ?? 0) + case .snapshot(let value): self.lastSeenServerSeq = max(self.lastSeenServerSeq, value.snapshots.map(\.fromSeq).max() ?? 0) + } + return result + } + + /// Observe flow control without owning or releasing any data. + public var flowControlEvents: AsyncStream { + get async { await flowEngine().events } + } + + /// Obtain the owner installed by a typed or raw flow-control handshake. + public func subscription(for channel: String) async -> AHPSubscription? { + await windowingEngine?.subscription(channel) + } + + internal func dropManagedSubscription(_ channel: String, error: Error) async { + await windowingEngine?.drop(channel, error: error) + } + + private func flowEngine() -> AHPWindowingEngine { + if let windowingEngine { return windowingEngine } + let engine = AHPWindowingEngine(config: config.windowing, write: { [weak self] data, before, limited in + guard let self else { throw AHPClientError.shutdown } + try await self.writeWindowData(data, beforeWrite: before, callerInitiated: limited) + }, fail: { [weak self] error in + await self?.handleTransportFailure(error) + }, dispatch: { [weak self] action, channel, seq in + guard let self else { throw AHPClientError.shutdown } + return try await self.dispatchAndWait(action, channel: channel, clientSeq: seq) + }, update: { [weak self] channel, options in + guard let self else { throw AHPClientError.shutdown } + return try await self.subscribe(channel, options: options) + }) + windowingEngine = engine + return engine + } + + private func writeWindowData(_ data: Data, beforeWrite: AHPBeforeWrite? = nil, + callerInitiated: Bool = false) async throws { + guard !didShutdown, let continuation = outboundContinuation else { throw AHPClientError.shutdown } + guard !callerInitiated || callerWrites.count < 256 else { throw AHPWindowingError.backpressure } + try await withCheckedThrowingContinuation { completion in + let id = UUID() + windowWrites[id] = completion + if callerInitiated { callerWrites.insert(id) } + continuation.yield(.init(data: data, completion: id, beforeWrite: beforeWrite)) + } + } + + private func finishWindowWrite(_ id: UUID, error: Error?) { + callerWrites.remove(id) + guard let reply = windowWrites.removeValue(forKey: id) else { return } + if let error { reply.resume(throwing: error) } else { reply.resume() } + } + + private func dispatchAndWait(_ action: StateAction, channel: String, clientSeq: Int?) + async throws -> DispatchHandle { + let seq = clientSeq ?? nextClientSeq + if seq >= nextClientSeq { nextClientSeq = seq + 1 } + try await send(method: "dispatchAction", + params: DispatchActionParams(channel: channel, clientSeq: seq, action: action)) + return DispatchHandle(clientSeq: seq) + } + + /// Send a notification and await its actual transport write, including ordinary fallback. + public func send(method: String, params: P) async throws { + try await flowEngine().send(AHPNotificationMessage(method: method, params: params)) + } + + private func transportWindowCapabilities() async -> AHPTransportCapabilities { + guard let transport = transport as? AHPWindowingTransport else { return .init() } + return await transport.windowingCapabilities + } + + private func prepareRequestContext(method: String, params: Data) + async throws -> (Data, AHPSubscriptionOptions?) { + if method == "subscribe", let request = try? decoder.decode(SubscribeParams.self, from: params), + request.channel.hasPrefix("ahp-tcp:"), request.flowControl == nil { + throw TransportError.protocol("TCP requires windowed subscription options") + } + if method == "reconnect", let request = try? decoder.decode(ReconnectParams.self, from: params), + request.subscriptions.contains(where: { $0.hasPrefix("ahp-tcp:") }) { + throw TransportError.protocol("TCP cannot reconnect") + } + if method == "subscribe", let request = try? decoder.decode(SubscribeParams.self, from: params), + windowingEngine != nil || request.flowControl != nil { + let engine = flowEngine() + let capabilities = await transportWindowCapabilities() + if let offer = request.flowControl?.receive, !capabilities.preservesOriginalJSONText { + let options = AHPSubscriptionOptions.windowed( + windowBytes: offer.windowBytes, maximumFrameBytes: offer.maximumFrameBytes, + delivery: request.delivery, view: request.view) + _ = try await engine.prepare(channel: request.channel, options: options, capabilities: capabilities) + var object = try JSONSerialization.jsonObject(with: params) as! [String: Any] + object.removeValue(forKey: "flowControl") + let adjusted = try JSONSerialization.data(withJSONObject: object) + try await engine.prepareGeneric(adjusted, capabilities: capabilities) + return (adjusted, options) + } + do { + try await engine.prepareGeneric(params, capabilities: capabilities) + if request.flowControl == nil { try await engine.pauseForOrdinarySubscribe(request.channel) } + } catch { + await engine.abortSubscribe(request.channel, error: error) + throw error + } + } else if method == "reconnect", + windowingEngine != nil || (try? decoder.decode(ReconnectParams.self, from: params).windows) != nil { + try await flowEngine().prepareReconnect(params, capabilities: transportWindowCapabilities()) + } + return (params, nil) + } + + private func flowNegotiation(method: String, params: Data) async -> Bool { + if method == "subscribe", let request = try? decoder.decode(SubscribeParams.self, from: params) { + if request.flowControl != nil { return true } + return await windowingEngine?.isWindowed(request.channel) ?? false + } + if method == "reconnect" { + return (try? decoder.decode(ReconnectParams.self, from: params).windows) != nil + } + return false + } + + private func checkRequestAdmission(method: String, params: Data) async throws { + do { + try Task.checkCancellation() + guard !didShutdown, outboundContinuation != nil else { throw AHPClientError.shutdown } + } catch { + if method == "subscribe", let request = try? decoder.decode(SubscribeParams.self, from: params) { + await windowingEngine?.abortSubscribe(request.channel) + } + throw error + } + } + + public func subscribe(_ uri: String, options: AHPSubscriptionOptions) async throws -> AHPSubscription { + try await subscribeManaged(uri, options: options, observer: nil, onPrepared: nil) + } + + internal func subscribeManaged(_ uri: String, options: AHPSubscriptionOptions, + observer: (@Sendable (AHPDeliveryPayload) async -> Void)?, + onPrepared: (@Sendable (AHPSubscription) async throws -> Void)?) + async throws -> AHPSubscription { + guard !didShutdown, outboundContinuation != nil else { throw AHPClientError.shutdown } + let (subscription, offer, created) = try await flowEngine().prepare( + channel: uri, options: options, capabilities: transportWindowCapabilities()) + do { + if let observer { + await subscription.markManaged() + await subscription.setObserver(observer) + } + try await onPrepared?(subscription) + let _: SubscribeResult = try await requestWithOptions(method: "subscribe", params: SubscribeParams( + channel: uri, delivery: options.delivery, view: options.view, + flowControl: offer.map { SubscriptionFlowControlOptions(receive: $0) }), + subscriptionOptions: options, owner: subscription, created: created) + return subscription + } catch { + await windowingEngine?.abortPrepared(subscription, error: error) + throw error + } + } + + internal func restoreManagedSubscriptions(_ subscriptions: [AHPSubscription]) + async throws -> ChannelDeliveryResumeOptions { + var items: [ChannelDeliveryResumeParams] = [] + for subscription in subscriptions { + let pendingOffer = await subscription.pendingWindowOffer + let wasWindowed = await subscription.status().mode == "windowed" || pendingOffer + let options = await subscription.options + let (_, offer, _) = try await flowEngine().prepare( + channel: subscription.channel, options: options, + capabilities: transportWindowCapabilities(), retained: subscription) + if let offer, wasWindowed { + items.append(.init(channel: subscription.channel, + flowControl: .init(receive: offer), + lastAppliedServerSeq: await subscription.lastAppliedServerSeq)) + } + } + return .init(items: items) + } + /// Protocol-level liveness `ping`. /// /// Verifies the connection is still alive and keeps it from being closed @@ -397,11 +632,17 @@ public actor AHPClient { public func unsubscribe(_ uri: String) async throws { let listeners = perUriListeners.removeValue(forKey: uri) ?? [] for l in listeners { l.continuation.finish() } - try await notify(method: "unsubscribe", params: UnsubscribeParams(channel: uri)) + if let engine = windowingEngine, await engine.subscription(uri) != nil { + try await engine.unsubscribe(uri) + } else { + try await notify(method: "unsubscribe", params: UnsubscribeParams(channel: uri)) + } } /// Fire a write-ahead `dispatchAction` notification on `channel`. Returns - /// a handle carrying the assigned `clientSeq`. + /// a handle carrying the assigned `clientSeq`. On an accepted windowed + /// channel this waits for transport completion; ordinary delivery retains + /// its fire-and-forget behavior. @discardableResult public func dispatch(_ action: StateAction, channel: String) async throws -> DispatchHandle { let seq = nextClientSeq @@ -579,6 +820,13 @@ public actor AHPClient { public func request( method: String, params: P + ) async throws -> R { + try await requestWithOptions(method: method, params: params, subscriptionOptions: nil) + } + + private func requestWithOptions( + method: String, params: P, subscriptionOptions: AHPSubscriptionOptions?, + owner: AHPSubscription? = nil, created: Bool? = nil ) async throws -> R { if didShutdown { throw AHPClientError.shutdown } guard let cont = outboundContinuation else { @@ -588,6 +836,15 @@ public actor AHPClient { // Avoids minting a request id and pushing wire bytes for a // request whose result will be thrown away immediately. try Task.checkCancellation() + let encodedParams = try encoder.encode(params) + let channel = method == "subscribe" ? (try? decoder.decode(SubscribeParams.self, from: encodedParams).channel) : nil + var binding = owner + if binding == nil, let channel { binding = await windowingEngine?.subscription(channel) } + let createdHere = created ?? (binding == nil) + let (contextParams, fallbackOptions) = try await prepareRequestContext(method: method, params: encodedParams) + if binding == nil, let channel { binding = await windowingEngine?.subscription(channel) } + let isFlowNegotiation = await flowNegotiation(method: method, params: contextParams) + try await checkRequestAdmission(method: method, params: contextParams) // Order matters: register the pending entry BEFORE pushing the // outbound message, so an immediate response (which a fast in-memory @@ -600,10 +857,14 @@ public actor AHPClient { try await withCheckedThrowingContinuation { (continuation: CheckedContinuation) in let entry = PendingEntry(continuation: continuation) pending[id] = entry + requestContexts[id] = (method, contextParams) + requestOptions[id] = subscriptionOptions ?? fallbackOptions + if let binding { requestSubscriptions[id] = .init(owner: binding, created: createdHere) } + if isFlowNegotiation { flowNegotiations.insert(id) } let wireData: Data do { - wireData = try buildRequestWire(id: id, method: method, params: params) + wireData = try buildRawRequestWire(id: id, method: method, paramsData: contextParams) } catch { pending.removeValue(forKey: id) continuation.resume(throwing: AHPClientError.decoding( @@ -611,7 +872,7 @@ public actor AHPClient { )) return } - cont.yield(wireData) + cont.yield(.init(data: wireData, completion: nil)) // Schedule a timeout. The timeout task races against the response; // whichever resolves first wins. The task is tracked on the entry @@ -665,6 +926,14 @@ public actor AHPClient { throw AHPClientError.shutdown } try Task.checkCancellation() + let channel = method == "subscribe" ? (try? decoder.decode(SubscribeParams.self, from: paramsData).channel) : nil + var binding: AHPSubscription? + if let channel { binding = await windowingEngine?.subscription(channel) } + let createdHere = binding == nil + let (paramsData, fallbackOptions) = try await prepareRequestContext(method: method, params: paramsData) + if binding == nil, let channel { binding = await windowingEngine?.subscription(channel) } + let isFlowNegotiation = await flowNegotiation(method: method, params: paramsData) + try await checkRequestAdmission(method: method, params: paramsData) let id = nextRequestId nextRequestId += 1 @@ -673,6 +942,10 @@ public actor AHPClient { try await withCheckedThrowingContinuation { (continuation: CheckedContinuation) in let entry = PendingEntry(continuation: continuation) pending[id] = entry + requestContexts[id] = (method, paramsData) + requestOptions[id] = fallbackOptions + if let binding { requestSubscriptions[id] = .init(owner: binding, created: createdHere) } + if isFlowNegotiation { flowNegotiations.insert(id) } let wireData: Data do { @@ -684,7 +957,7 @@ public actor AHPClient { )) return } - cont.yield(wireData) + cont.yield(.init(data: wireData, completion: nil)) let timeoutDuration = config.requestTimeout entry.timeoutTask = Task { [weak self] in @@ -706,26 +979,49 @@ public actor AHPClient { /// no-op when the entry has already been removed by the response or /// timeout path — only the first of those three to find the entry /// wins, so the continuation is never double-resumed. - private func cancelPending(id: Int) { - if let entry = pending.removeValue(forKey: id) { - entry.timeoutTask?.cancel() - entry.continuation.resume(throwing: CancellationError()) + private func cancelPending(id: Int) async { + await abandonPending(id: id, error: CancellationError()) + } + + private func abandonPending(id: Int, error: Error) async { + guard let entry = pending.removeValue(forKey: id) else { return } + entry.timeoutTask?.cancel() + if let context = requestContexts[id], ["subscribe", "reconnect"].contains(context.0) { + abandonedRequests.insert(id) + if let subscription = requestSubscriptions[id] { + await windowingEngine?.abandonSubscribe(subscription.owner, created: subscription.created, error: error) + } + } else { + requestContexts.removeValue(forKey: id) + requestOptions.removeValue(forKey: id) + requestSubscriptions.removeValue(forKey: id) + flowNegotiations.remove(id) } + entry.continuation.resume(throwing: error) } - /// Send a JSON-RPC notification (fire-and-forget). + /// Send a JSON-RPC notification. Ordinary delivery is fire-and-forget; + /// accepted windowed data waits for credit and transport completion. public func notify(method: String, params: P) async throws { if didShutdown { throw AHPClientError.shutdown } guard let cont = outboundContinuation else { throw AHPClientError.shutdown } - let wireData: Data + let paramsData: Data do { - wireData = try buildNotificationWire(method: method, params: params) + paramsData = try encoder.encode(params) } catch { throw AHPClientError.decoding("failed to encode params for \(method): \(error)") } - cont.yield(wireData) + if let engine = windowingEngine { + if let message = try? AHPNotificationMessage(method: method, paramsJSON: paramsData), + await engine.isWindowed(message.channel) { + try await engine.send(message) + return + } + } + let wireData = try buildRawNotificationWire(method: method, paramsData: paramsData) + cont.yield(.init(data: wireData, completion: nil)) } // MARK: - Private: wire-building @@ -770,7 +1066,10 @@ public actor AHPClient { private func buildNotificationWire( method: String, params: P ) throws -> Data { - let paramsData = try encoder.encode(params) + try buildRawNotificationWire(method: method, paramsData: encoder.encode(params)) + } + + private func buildRawNotificationWire(method: String, paramsData: Data) throws -> Data { let paramsAny = try JSONSerialization.jsonObject( with: paramsData, options: [.fragmentsAllowed] ) @@ -787,8 +1086,20 @@ public actor AHPClient { private func runReceiveLoop() async { while !Task.isCancelled { let msg: TransportMessage? + let rawText: String? do { - msg = try await transport.recv() + if let measured = transport as? AHPWindowingTransport { + let wire = try await measured.recvWireMessage() + msg = wire?.message + rawText = wire?.originalJSONText + } else { + msg = try await transport.recv() + switch msg { + case .text(let text): rawText = text + case .binary(let data): rawText = String(data: data, encoding: .utf8) + default: rawText = nil + } + } } catch { await handleTransportFailure(error) return @@ -804,7 +1115,7 @@ public actor AHPClient { let data: Data switch msg { case .text(let s): - guard let d = s.data(using: .utf8) else { + guard let d = (rawText ?? s).data(using: .utf8) else { #if DEBUG print("[AHPClient] dropped malformed text frame") #endif @@ -812,7 +1123,7 @@ public actor AHPClient { } data = d case .binary(let d): - data = d + data = rawText.map { Data($0.utf8) } ?? d case .parsed(let parsed): // Slow path: re-serialize via Codable so the rest of the // pipeline only deals with raw bytes. Note that any @@ -820,6 +1131,10 @@ public actor AHPClient { // have been corrupted by an earlier `JSONDecoder` pass — // see `TransportMessage` docs and microsoft/agent-host-protocol#123. // Transports SHOULD prefer `.text`/`.binary` for inbound frames. + if let rawText { + data = Data(rawText.utf8) + break + } guard let d = try? encoder.encode(parsed) else { #if DEBUG print("[AHPClient] dropped unencodable parsed frame") @@ -829,12 +1144,16 @@ public actor AHPClient { data = d } guard let frame = parseRawFrame(from: data) else { + if await windowingEngine?.hasWindows == true { + await handleTransportFailure(TransportError.protocol("malformed JSON-RPC on windowed connection")) + return + } #if DEBUG print("[AHPClient] dropped malformed frame") #endif continue } - await dispatchInbound(frame) + await dispatchInbound(frame, rawText: rawText) } } @@ -900,14 +1219,47 @@ public actor AHPClient { return nil } - private func dispatchInbound(_ frame: RawFrame) async { + private func dispatchInbound(_ frame: RawFrame, rawText: String? = nil) async { switch frame { case .successResponse(let id, let resultData): + flowNegotiations.remove(id) + let abandoned = abandonedRequests.remove(id) != nil + let subscription = requestSubscriptions.removeValue(forKey: id) + if let (method, params) = requestContexts.removeValue(forKey: id), let engine = windowingEngine { + do { + if method == "subscribe" { + let p = try decoder.decode(SubscribeParams.self, from: params) + try await engine.installSubscribe(request: p, options: requestOptions.removeValue(forKey: id), + resultData: resultData, expected: subscription?.owner, + retain: !abandoned) + if abandoned, let subscription { + await engine.cleanAbandoned(subscription.owner, created: subscription.created) + } else if let subscription, await engine.subscription(p.channel) !== subscription.owner { + await engine.cleanAbandoned(subscription.owner, created: true) + } + } else if method == "reconnect" { + try await engine.installReconnect(request: params, resultData: resultData) + } + } catch { + await handleTransportFailure(TransportError.protocol("\(error)")) + return + } + } if let entry = pending.removeValue(forKey: id) { entry.timeoutTask?.cancel() entry.continuation.resume(returning: resultData) } case .errorResponse(let id, let error): + flowNegotiations.remove(id) + abandonedRequests.remove(id) + let subscription = requestSubscriptions.removeValue(forKey: id) + if let context = requestContexts.removeValue(forKey: id), context.0 == "subscribe", + let request = try? decoder.decode(SubscribeParams.self, from: context.1) { + await windowingEngine?.abortSubscribe(request.channel, + error: AHPClientError.rpc(code: error.code, message: error.message, data: error.data), + expected: subscription?.owner) + } + requestOptions.removeValue(forKey: id) if let entry = pending.removeValue(forKey: id) { entry.timeoutTask?.cancel() entry.continuation.resume(throwing: AHPClientError.rpc( @@ -915,6 +1267,26 @@ public actor AHPClient { )) } case .notification(let method, let params): + if let engine = windowingEngine, let params { + do { + if let message = try? AHPNotificationMessage(method: method, paramsJSON: params), + let subscription = await engine.subscription(message.channel) { + if let delivered = try await engine.receive(message, rawText: rawText) { + let body = try JSONSerialization.jsonObject(with: delivered.paramsJSON) as? [String: Any] + let seq = body?["serverSeq"] as? Int ?? + ((body?["snapshot"] as? [String: Any])?["fromSeq"] as? Int) ?? 0 + lastSeenServerSeq = max(lastSeenServerSeq, seq) + if await !subscription.managed { + await handleNotification(method: delivered.method, paramsData: delivered.paramsJSON) + } + } + return + } + } catch { + await handleTransportFailure(TransportError.protocol("\(error)")) + return + } + } await handleNotification(method: method, paramsData: params) case .request(let id, let method, let params): // Server-initiated requests (the symmetrical `resource*` family) @@ -1028,6 +1400,7 @@ public actor AHPClient { } else { clientError = .shutdown } + await windowingEngine?.disconnect(clientError) outboundContinuation?.finish() outboundContinuation = nil @@ -1048,6 +1421,15 @@ public actor AHPClient { } private func failAllPending(with error: AHPClientError) { + requestContexts.removeAll() + requestOptions.removeAll() + requestSubscriptions.removeAll() + abandonedRequests.removeAll() + flowNegotiations.removeAll() + let writes = windowWrites.values + windowWrites.removeAll() + callerWrites.removeAll() + for reply in writes { reply.resume(throwing: error) } let entries = pending pending.removeAll() for (_, entry) in entries { @@ -1056,11 +1438,8 @@ public actor AHPClient { } } - private func timeoutPending(id: Int) { - if let entry = pending.removeValue(forKey: id) { - entry.timeoutTask = nil - entry.continuation.resume(throwing: AHPClientError.requestTimeout) - } + private func timeoutPending(id: Int) async { + await abandonPending(id: id, error: AHPClientError.requestTimeout) } /// Answer an inbound server-initiated request via the installed @@ -1093,7 +1472,7 @@ public actor AHPClient { "result": resultAny, ] guard let data = try? JSONSerialization.data(withJSONObject: dict) else { return } - cont.yield(data) + cont.yield(.init(data: data, completion: nil)) } /// Yield a JSON-RPC error response with the given code/message/data. @@ -1114,7 +1493,7 @@ public actor AHPClient { "error": errorObject, ] guard let wire = try? JSONSerialization.data(withJSONObject: dict) else { return } - cont.yield(wire) + cont.yield(.init(data: wire, completion: nil)) } /// Adapt a typed `ResourceRequestHandlers` into a generic @@ -1189,7 +1568,7 @@ public actor AHPClient { ] as [String: Any], ] guard let data = try? JSONSerialization.data(withJSONObject: dict) else { return } - cont.yield(data) + cont.yield(.init(data: data, completion: nil)) } private func finishAllSubscriptions() { @@ -1261,6 +1640,11 @@ public actor AHPClient { internal func _pendingCount() -> Int { return pending.count } + + internal func _windowWriteCount() -> Int { windowWrites.count } + internal func _callerWriteCount() -> Int { callerWrites.count } + internal func _windowOrderCount() async -> Int { await windowingEngine?.bookkeepingCount() ?? 0 } + internal func _requestContextCount() -> Int { requestContexts.count } } // MARK: - Server-initiated request handling diff --git a/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/AHPClientConfig.swift b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/AHPClientConfig.swift index f20a8cc6a..339cda01a 100644 --- a/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/AHPClientConfig.swift +++ b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/AHPClientConfig.swift @@ -42,6 +42,18 @@ public struct AHPClientConfig: Sendable { /// Optional transport liveness policy. Defaults to `.disabled`. public var keepAlive: AHPKeepAlivePolicy + public var windowing: AHPWindowingConfig = .init() + + public init( + windowing: AHPWindowingConfig, + requestTimeout: Duration = .seconds(30), + subscriptionBufferSize: Int = 256, + keepAlive: AHPKeepAlivePolicy = .disabled + ) { + self.init(requestTimeout: requestTimeout, subscriptionBufferSize: subscriptionBufferSize, + keepAlive: keepAlive) + self.windowing = windowing + } public init( requestTimeout: Duration = .seconds(30), diff --git a/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/AHPClientEvents.swift b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/AHPClientEvents.swift index 338f13b2e..549cc8825 100644 --- a/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/AHPClientEvents.swift +++ b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/AHPClientEvents.swift @@ -45,7 +45,8 @@ public enum ConnectionState: Sendable, Equatable { /// Handle returned from `AHPClient.dispatch`. /// -/// Dispatch is fire-and-forget by design; the handle records the +/// Ordinary dispatch is fire-and-forget; windowed dispatch waits for the +/// transport. The handle records the /// client-assigned `clientSeq` so callers can correlate optimistic local /// updates with server echoes (`ActionEnvelope.origin.clientSeq`). public struct DispatchHandle: Sendable, Equatable { diff --git a/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/AHPStateMirror.swift b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/AHPStateMirror.swift index 04f6c00d3..a6a6bc22e 100644 --- a/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/AHPStateMirror.swift +++ b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/AHPStateMirror.swift @@ -25,6 +25,14 @@ public actor AHPStateMirror { public init() {} + public func apply(delivery: AHPDelivery) { + switch delivery.payload { + case .snapshot(let snapshot): applySnapshot(snapshot) + case .event(.action(let envelope)): apply(envelope) + default: break + } + } + /// Apply a single action envelope, routing by `envelope.channel`. /// /// - Root state lives at `RootResourceURI`; all other channels are diff --git a/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/HostClientHandle.swift b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/HostClientHandle.swift index d0e599fec..592c0f4ed 100644 --- a/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/HostClientHandle.swift +++ b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/HostClientHandle.swift @@ -122,4 +122,20 @@ public struct HostClientHandle: Sendable { public func rawClient() async -> AHPClient { client } + + public func send(method: String, params: P) async throws { + try await checkAlive() + try await client.send(method: method, params: params) + } + + public func flowControlStatus(channel: String) async throws -> AHPFlowControlStatus? { + try await checkAlive() + return await client.subscription(for: channel)?.status() + } + + public func waitUntilWritable(channel: String) async throws { + try await checkAlive() + guard let subscription = await client.subscription(for: channel) else { throw AHPClientError.shutdown } + try await subscription.waitUntilWritable() + } } diff --git a/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/HostConfig.swift b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/HostConfig.swift index fa6617f74..d72013c63 100644 --- a/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/HostConfig.swift +++ b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/HostConfig.swift @@ -78,6 +78,12 @@ public struct HostConfig: Sendable { return copy } + public func withWindowing(_ config: AHPWindowingConfig) -> Self { + var copy = self + copy.clientConfig.windowing = config + return copy + } + /// Override the reconnect policy. public func withReconnectPolicy(_ policy: ReconnectPolicy) -> Self { var copy = self diff --git a/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/HostRuntime.swift b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/HostRuntime.swift index 882a0b4c5..982fdbaba 100644 --- a/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/HostRuntime.swift +++ b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/HostRuntime.swift @@ -28,6 +28,7 @@ internal final class HostRuntime: Sendable { private let fanOut: @Sendable (HostSubscriptionEvent) async -> Void /// Sink for connection-level events. Awaited for the same reason. private let hostEventSink: @Sendable (HostEvent) async -> Void + private let flowEventSink: @Sendable (AHPHostedFlowControlEvent) async -> Void private let cmdContinuation: AsyncStream.Continuation private let cmdStream: AsyncStream @@ -44,11 +45,13 @@ internal final class HostRuntime: Sendable { config: HostConfig, clientIdStore: ClientIdStore, fanOut: @escaping @Sendable (HostSubscriptionEvent) async -> Void, - hostEventSink: @escaping @Sendable (HostEvent) async -> Void + hostEventSink: @escaping @Sendable (HostEvent) async -> Void, + flowEventSink: @escaping @Sendable (AHPHostedFlowControlEvent) async -> Void = { _ in } ) async { self.config = config self.fanOut = fanOut self.hostEventSink = hostEventSink + self.flowEventSink = flowEventSink let resolved: String if let explicit = config.clientId { @@ -133,6 +136,33 @@ internal final class HostRuntime: Sendable { try await withCheckedThrowingContinuation { (continuation: CheckedContinuation) in cmdContinuation.yield(.subscribe(uri: uri, reply: continuation)) } + + } + + func subscribe(_ uri: String, options: AHPSubscriptionOptions) async throws -> AHPSubscription { + guard let client = await shared.currentClient() else { throw HostError.hostShutDown(config.id) } + let subscription = try await client.subscribeManaged(uri, options: options, observer: { [weak self] payload in + await self?.handleManagedPayload(payload) + }, onPrepared: { [shared, config] subscription in + await subscription.setCloseObserver { [shared, weak subscription] in + guard let subscription else { return } + await shared.update { state in + guard state.managedSubscriptions[uri] === subscription else { return } + state.managedSubscriptions.removeValue(forKey: uri) + state.subscriptions.removeAll { $0 == uri } + } + } + var current = false + await shared.update { state in + current = state.currentClient === client && !state.shuttingDown + if current { + state.managedSubscriptions[uri] = subscription + if !state.subscriptions.contains(uri) { state.subscriptions.append(uri) } + } + } + guard current else { throw HostError.hostShutDown(config.id) } + }) + return subscription } /// Unsubscribe from `uri`. Stops replay of `uri` across reconnects. Safe @@ -165,9 +195,17 @@ internal final class HostRuntime: Sendable { /// Tear down the supervisor: cancel any in-flight connect/sleep, close /// the current `AHPClient` if any, and finish the supervisor task. func shutdown() async { + await shared.update { $0.shuttingDown = true } cmdContinuation.yield(.shutdown) cmdContinuation.finish() + supervisorTask?.cancel() + let state = await shared.internalState + await state.currentClient?.shutdown() + if state.connectingClient !== state.currentClient { await state.connectingClient?.shutdown() } await supervisorTask?.value + let subscriptions = await shared.internalState.managedSubscriptions + for subscription in subscriptions.values { await subscription.finish(AHPClientError.shutdown) } + await shared.update { $0.managedSubscriptions.removeAll() } } // MARK: - Supervisor loop @@ -277,7 +315,12 @@ internal final class HostRuntime: Sendable { /// the next attempt while the previous client's tasks remain alive). private func connectOnce() async throws -> ConnectionStreams { let transport = try await config.transportFactory(config.id) + guard await !shared.internalState.shuttingDown else { + try? await transport.close() + throw HostError.hostShutDown(config.id) + } let client = AHPClient(transport: transport, config: config.clientConfig) + await shared.update { $0.connectingClient = client } // Attach the events and state-change taps BEFORE the // initialize/reconnect handshake so any notifications the server @@ -292,6 +335,7 @@ internal final class HostRuntime: Sendable { // to `.disconnected`. let events = await client.events let stateChanges = await client.stateChanges + let flowEvents = await client.flowControlEvents try await client.connect() @@ -299,7 +343,8 @@ internal final class HostRuntime: Sendable { // propagating — `client.connect()` started writer/receive tasks // that hold the transport. `withClientShutdownOnThrow` enforces it. return try await withClientShutdownOnThrow(client) { - try await self.completeHandshake(client: client, events: events, stateChanges: stateChanges) + try await self.completeHandshake(client: client, events: events, + stateChanges: stateChanges, flowEvents: flowEvents) } } @@ -310,13 +355,20 @@ internal final class HostRuntime: Sendable { private func completeHandshake( client: AHPClient, events: AsyncStream, - stateChanges: AsyncStream + stateChanges: AsyncStream, + flowEvents: AsyncStream ) async throws -> ConnectionStreams { // Decide between initialize and reconnect based on prior state. let priorSnapshot = await shared.internalState let canReconnect = priorSnapshot.serverSeq > 0 && !priorSnapshot.subscriptions.isEmpty - let priorSubscriptions = priorSnapshot.subscriptions + let priorSubscriptions = priorSnapshot.subscriptions.filter { !$0.hasPrefix("ahp-tcp:") } let priorSeq = priorSnapshot.serverSeq + let retained = priorSnapshot.managedSubscriptions.values.filter { !$0.channel.hasPrefix("ahp-tcp:") } + let windows = try await client.restoreManagedSubscriptions(Array(retained)) + await shared.update { state in + state.subscriptions.removeAll { $0.hasPrefix("ahp-tcp:") } + state.managedSubscriptions = state.managedSubscriptions.filter { !$0.key.hasPrefix("ahp-tcp:") } + } var initResult: InitializeResult? = nil var reconnectResult: ReconnectResult? = nil @@ -324,11 +376,14 @@ internal final class HostRuntime: Sendable { if canReconnect { do { - reconnectResult = try await client.reconnect( - clientId: clientId, - lastSeenServerSeq: priorSeq, - subscriptions: priorSubscriptions - ) + if windows.items.isEmpty { + reconnectResult = try await client.reconnect( + clientId: clientId, lastSeenServerSeq: priorSeq, subscriptions: priorSubscriptions) + } else { + reconnectResult = try await client.reconnect( + clientId: clientId, lastSeenServerSeq: priorSeq, + subscriptions: priorSubscriptions, windows: windows) + } } catch let error as AHPClientError { if case .rpc = error { let init1 = try await client.initialize( @@ -366,6 +421,7 @@ internal final class HostRuntime: Sendable { await shared.update { state in state.generation = state.generation &+ 1 state.currentClient = client + state.connectingClient = nil state.lastConnectedAt = Date() state.lastError = nil state.serverSeq = newSeq @@ -394,11 +450,21 @@ internal final class HostRuntime: Sendable { if let reconnectResult { await applyReconnectResult(reconnectResult, priorSubscriptions: priorSubscriptions) await hostEventSink(.reconnectResult(config.id, reconnectResult)) + } else if initResult != nil { + for subscription in retained { + do { + _ = try await client.subscribe(subscription.channel, options: await subscription.options) + } catch { + await subscription.removeManagedRegistration() + await subscription.finish(error) + await client.dropManagedSubscription(subscription.channel, error: error) + } + } } await transition(to: .connected, error: nil) await hostEventSink(.connected(config.id, generation: newGeneration)) - return ConnectionStreams(events: events, stateChanges: stateChanges) + return ConnectionStreams(events: events, stateChanges: stateChanges, flowEvents: flowEvents) } /// Apply the result of a successful `reconnect` call. @@ -414,6 +480,7 @@ internal final class HostRuntime: Sendable { switch result { case .replay(let replay): for envelope in replay.actions { + if await shared.internalState.managedSubscriptions[envelope.channel] != nil { continue } await applyAction(envelope) if config.fanOutReconnectReplayActions { await fanOut(HostSubscriptionEvent( @@ -428,7 +495,9 @@ internal final class HostRuntime: Sendable { state.subscriptions.removeAll { replay.missing.contains($0) } } } + await pruneManaged(replay.windows?.items.map(\.channel) ?? []) case .snapshot(let snapshotResult): + let windows = snapshotResult.windows?.items.map(\.channel) ?? [] await shared.update { state in var surviving: [String] = [] surviving.reserveCapacity(snapshotResult.snapshots.count) @@ -443,9 +512,25 @@ internal final class HostRuntime: Sendable { surviving.append(snapshot.resource) } state.subscriptions.removeAll { uri in - priorSubscriptions.contains(uri) && !surviving.contains(uri) + priorSubscriptions.contains(uri) && !surviving.contains(uri) && !windows.contains(uri) } } + await pruneManaged(windows) + } + } + + private func pruneManaged(_ kept: [String]) async { + let current = await shared.internalState.managedSubscriptions + var dropped: [String] = [] + for (channel, subscription) in current { + if await subscription.status().mode == "windowed", !kept.contains(channel) { + dropped.append(channel) + await subscription.finish(AHPWindowingError.subscriptionDropped(channel)) + } + } + await shared.update { state in + for channel in dropped { state.managedSubscriptions.removeValue(forKey: channel) } + state.subscriptions.removeAll { dropped.contains($0) } } } @@ -483,9 +568,18 @@ internal final class HostRuntime: Sendable { } } + let flowPump = Task { [flowEventSink, config] in + if let events = streams.flowEvents { + for await event in events { + await flowEventSink(.init(hostId: config.id, event: event)) + } + } + } + defer { pumpTask.cancel() dropDetector.cancel() + flowPump.cancel() } while let cmd = await iter.next() { @@ -508,8 +602,18 @@ internal final class HostRuntime: Sendable { let result = await handleUnsubscribe(uri) resumeCommand(reply: reply, with: result) case .dispatch(let action, let channel, let clientSeq, let reply): - let result = await handleDispatch(action, channel: channel, clientSeq: clientSeq) - resumeCommand(reply: reply, with: result) + if let client = await shared.currentClient(), + let subscription = await client.subscription(for: channel), + await subscription.status().mode == "windowed" { + // Do not monopolize the supervisor while a send waits for credit. + Task { [self] in + let result = await handleDispatch(action, channel: channel, clientSeq: clientSeq, client: client) + resumeCommand(reply: reply, with: result) + } + } else { + let result = await handleDispatch(action, channel: channel, clientSeq: clientSeq) + resumeCommand(reply: reply, with: result) + } } } return .shutdown @@ -588,6 +692,8 @@ internal final class HostRuntime: Sendable { // MARK: - Event handling private func handleEvent(_ event: ClientEvent) async { + if let channel = event.resource, + await shared.internalState.managedSubscriptions[channel] != nil { return } // Mutate per-host mirrors before broadcasting so observers reading // the next snapshot see the post-event state. switch event.event { @@ -597,6 +703,7 @@ internal final class HostRuntime: Sendable { await shared.update { state in state.sessionSummaries[n.summary.resource] = n.summary } + case .sessionRemoved(let n): await shared.update { state in state.sessionSummaries.removeValue(forKey: n.session) @@ -619,6 +726,40 @@ internal final class HostRuntime: Sendable { await fanOut(hostEvent) } + private func handleManagedPayload(_ payload: AHPDeliveryPayload) async { + switch payload { + case .snapshot(let snapshot): + await shared.update { state in + state.serverSeq = max(state.serverSeq, snapshot.fromSeq) + if snapshot.resource == RootResourceURI, case .root(let root) = snapshot.state { + state.rootState = root + } + } + case .event(let event): + let channel: String + switch event { + case .action(let envelope): channel = envelope.channel; await applyAction(envelope) + case .sessionAdded(let value): + channel = value.channel + await shared.update { $0.sessionSummaries[value.summary.resource] = value.summary } + case .sessionRemoved(let value): + channel = value.channel + await shared.update { $0.sessionSummaries.removeValue(forKey: value.session) } + case .sessionSummaryChanged(let value): + channel = value.channel + await shared.update { state in + if var summary = state.sessionSummaries[value.session] { + applySummaryChanges(&summary, changes: value.changes) + state.sessionSummaries[value.session] = summary + } + } + case .authRequired(let value): channel = value.channel + } + await fanOut(.init(hostId: config.id, resource: channel, event: event)) + default: break + } + } + private func applyAction(_ envelope: ActionEnvelope) async { await shared.update { state in if envelope.serverSeq > state.serverSeq { @@ -662,12 +803,16 @@ internal final class HostRuntime: Sendable { } } await shared.removeSubscription(uri) + await shared.update { $0.managedSubscriptions.removeValue(forKey: uri) } return .success(()) } - private func handleDispatch(_ action: StateAction, channel: String, clientSeq: Int?) async -> Result { - guard let client = await shared.currentClient() else { - return .failure(.hostShutDown(config.id)) + private func handleDispatch(_ action: StateAction, channel: String, clientSeq: Int?, + client explicitClient: AHPClient? = nil) async -> Result { + let connection: AHPClient? + if let explicitClient { connection = explicitClient } else { connection = await shared.currentClient() } + guard let client = connection else { + return .failure(HostError.hostShutDown(config.id)) } do { let handle: DispatchHandle @@ -678,9 +823,11 @@ internal final class HostRuntime: Sendable { } return .success(handle) } catch let error as AHPClientError { - return .failure(.client(error)) + return .failure(HostError.client(error)) + } catch let error as AHPWindowingError { + return .failure(error) } catch { - return .failure(.client(.transport(.io(String(describing: error))))) + return .failure(HostError.client(.transport(.io(String(describing: error))))) } } @@ -746,12 +893,13 @@ private enum RunOutcome { internal struct ConnectionStreams: @unchecked Sendable { let events: AsyncStream let stateChanges: AsyncStream + var flowEvents: AsyncStream? = nil } /// Resume a `CheckedContinuation` with a `Result`. -private func resumeCommand( +private func resumeCommand( reply: CheckedContinuation, - with result: Result + with result: Result ) { switch result { case .success(let value): reply.resume(returning: value) diff --git a/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/HostShared.swift b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/HostShared.swift index cc5e3676e..95295265c 100644 --- a/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/HostShared.swift +++ b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/HostShared.swift @@ -31,6 +31,9 @@ internal struct HostInternal { /// The currently-installed `AHPClient`, when connected. `nil` between /// connections. var currentClient: AHPClient? + var connectingClient: AHPClient? = nil + var shuttingDown = false + var managedSubscriptions: [String: AHPSubscription] = [:] func snapshot() -> HostHandle { let summaries = sessionSummaries.values diff --git a/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/MultiHostClient.swift b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/MultiHostClient.swift index 74afdb7ea..f9febebd0 100644 --- a/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/MultiHostClient.swift +++ b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Hosts/MultiHostClient.swift @@ -41,6 +41,7 @@ public actor MultiHostClient { private var nextListenerId: UInt64 = 1 private var subscriptionListeners: [UInt64: AsyncStream.Continuation] = [:] private var hostEventListeners: [UInt64: AsyncStream.Continuation] = [:] + private var flowListeners: [UInt64: AsyncStream.Continuation] = [:] /// Per-`(hostId, uri)` listener registry for `events(host:uri:)`. /// Lives on this actor (rather than inside `HostRuntime`) so listeners /// outlive any single `AHPClient` instance and survive reconnects — @@ -111,7 +112,8 @@ public actor MultiHostClient { config: config, clientIdStore: clientIdStore, fanOut: fanOut, - hostEventSink: hostEventSink + hostEventSink: hostEventSink, + flowEventSink: { [weak self] event in await self?.broadcastFlowEvent(event) } ) // We may have been shut down while awaiting the runtime init; bail @@ -257,6 +259,24 @@ public actor MultiHostClient { return try await runtime.subscribe(uri) } + public func subscribe(host: HostId, uri: String, options: AHPSubscriptionOptions) + async throws -> AHPSubscription { + guard let runtime = hosts[host] else { throw HostError.unknownHost(host) } + return try await runtime.subscribe(uri, options: options) + } + + public func flowControlEvents() -> AsyncStream { + let id = bumpListenerId() + return AsyncStream(bufferingPolicy: .bufferingNewest(defaultFanInBuffer)) { continuation in + flowListeners[id] = continuation + continuation.onTermination = { [weak self] _ in Task { await self?.removeFlowListener(id) } } + } + } + private func broadcastFlowEvent(_ event: AHPHostedFlowControlEvent) { + for continuation in flowListeners.values { continuation.yield(event) } + } + private func removeFlowListener(_ id: UInt64) { flowListeners.removeValue(forKey: id) } + /// Unsubscribe from `uri` on `host`. Drops the URI from the replay set. public func unsubscribe(host: HostId, uri: String) async throws { guard let runtime = hosts[host] else { @@ -697,6 +717,8 @@ public actor MultiHostClient { subscriptionListeners.removeAll() for cont in hostEventListeners.values { cont.finish() } hostEventListeners.removeAll() + for continuation in flowListeners.values { continuation.finish() } + flowListeners.removeAll() for bucket in perResourceListeners.values { for listener in bucket.values { listener.continuation.finish() } } diff --git a/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/MultiHostStateMirror.swift b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/MultiHostStateMirror.swift index d6cac1d58..b85c8a226 100644 --- a/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/MultiHostStateMirror.swift +++ b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/MultiHostStateMirror.swift @@ -56,6 +56,14 @@ public actor MultiHostStateMirror { public init() {} + public func apply(host: HostId, delivery: AHPDelivery) { + switch delivery.payload { + case .snapshot(let snapshot): applySnapshot(host: host, snapshot: snapshot) + case .event(.action(let envelope)): apply(host: host, envelope: envelope) + default: break + } + } + /// Convenience: apply a `HostSubscriptionEvent` produced by /// `MultiHostClient.events()`. Action envelopes are routed through /// the reducer; non-action events are dropped (they don't affect diff --git a/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Transport/NWConnectionWebSocketTransport.swift b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Transport/NWConnectionWebSocketTransport.swift index e62531ff6..2e7fc99e5 100644 --- a/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Transport/NWConnectionWebSocketTransport.swift +++ b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Transport/NWConnectionWebSocketTransport.swift @@ -130,6 +130,7 @@ public actor NWConnectionWebSocketTransport: AHPTransport, AHPKeepAliveTransport do { try await connectIfNeeded() while true { + let bufferedBytes = readBuffer.count switch try await nextMessageFromBuffer() { case .message(let message): switch message { @@ -140,6 +141,9 @@ public actor NWConnectionWebSocketTransport: AHPTransport, AHPKeepAliveTransport try await closeAfterPeerClose() return nil case .none: + // Continuations and controls can share a socket read with + // the next frame. Drain them before waiting for more bytes. + if readBuffer.count < bufferedBytes { continue } let chunk = try await receiveRaw() if chunk.isEmpty { markClosed() diff --git a/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Transport/URLSessionWebSocketTransport.swift b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Transport/URLSessionWebSocketTransport.swift index ca921b633..8b530e25f 100644 --- a/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Transport/URLSessionWebSocketTransport.swift +++ b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Transport/URLSessionWebSocketTransport.swift @@ -8,15 +8,23 @@ import Foundation /// Default WebSocket transport built on `URLSessionWebSocketTask`. -public final class URLSessionWebSocketTransport: AHPTransport, @unchecked Sendable { +public final class URLSessionWebSocketTransport: AHPWindowingTransport, @unchecked Sendable { private let url: URL private let headers: [String: String] private let session: URLSession private let lock = NSLock() private var task: URLSessionWebSocketTask? + private var started = false private var closed = false + public var windowingCapabilities: AHPTransportCapabilities { + get async { + let limit = try? ensureTask(start: false).maximumMessageSize + return .init(preservesOriginalJSONText: true, maximumReceiveMessageUTF8Bytes: limit) + } + } + /// Creates a new WebSocket transport. The connection is opened lazily on /// the first call to `send` or `recv`. public init( @@ -90,17 +98,20 @@ public final class URLSessionWebSocketTransport: AHPTransport, @unchecked Sendab return task } - private func ensureTask() throws -> URLSessionWebSocketTask { + private func ensureTask(start: Bool = true) throws -> URLSessionWebSocketTask { lock.lock() defer { lock.unlock() } if closed { throw TransportError.closed } - if let task { return task } + if let task { + if start && !started { task.resume(); started = true } + return task + } var request = URLRequest(url: url) for (k, v) in headers { request.setValue(v, forHTTPHeaderField: k) } let task = session.webSocketTask(with: request) - task.resume() + if start { task.resume(); started = true } self.task = task return task } diff --git a/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Windowing.swift b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Windowing.swift new file mode 100644 index 000000000..ddacf3e8f --- /dev/null +++ b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Windowing.swift @@ -0,0 +1,519 @@ +import Foundation +import AgentHostProtocol + +/// Defaults only: windowing is enabled by a subscription, never by configuration. +public struct AHPWindowingConfig: Sendable { + public var windowBytes: Int + public var maximumFrameBytes: Int? + public var maxBufferedSendBytes: Int + public var coalescers: [AHPCoalescer] + + public init( + windowBytes: Int = 262_144, + maximumFrameBytes: Int? = nil, + maxBufferedSendBytes: Int = 1_048_576, + coalescers: [AHPCoalescer] = [] + ) { + self.windowBytes = windowBytes + self.maximumFrameBytes = maximumFrameBytes + self.maxBufferedSendBytes = maxBufferedSendBytes + self.coalescers = coalescers + } +} + +public enum AHPReceiveMode: Sendable { case auto, manual } + +public struct AHPWindowOptions: Sendable { + public var windowBytes: Int? + public var maximumFrameBytes: Int? + public var receive: AHPReceiveMode + public var coalescing: Bool + + public init(windowBytes: Int? = nil, maximumFrameBytes: Int? = nil, + receive: AHPReceiveMode = .auto, coalescing: Bool = true) { + self.windowBytes = windowBytes + self.maximumFrameBytes = maximumFrameBytes + self.receive = receive + self.coalescing = coalescing + } +} + +public struct AHPSubscriptionOptions: Sendable { + public var window: AHPWindowOptions? + public var delivery: SubscriptionDeliveryOptions? + public var view: SubscribeView? + + public init(window: AHPWindowOptions? = nil, + delivery: SubscriptionDeliveryOptions? = nil, view: SubscribeView? = nil) { + self.window = window + self.delivery = delivery + self.view = view + } + + public static func windowed(windowBytes: Int? = nil, maximumFrameBytes: Int? = nil, + receive: AHPReceiveMode = .auto, coalescing: Bool = true, + delivery: SubscriptionDeliveryOptions? = nil, + view: SubscribeView? = nil) -> Self { + .init(window: .init(windowBytes: windowBytes, maximumFrameBytes: maximumFrameBytes, + receive: receive, coalescing: coalescing), delivery: delivery, view: view) + } +} + +public struct AHPFlowControlStatus: Codable, Sendable { + public var mode: String = "ordinary" + public var fallbackReason: String? + public var limits: ChannelFlowControl? + public var outstandingSendBytes: Int = 0 + public var bufferedSendBytes: Int = 0 + public var bufferedReceiveBytes: Int = 0 + public var writable: Bool = true +} + +public struct AHPFlowControlEvent: Sendable { + public let channel: String + public let kind: String + public let status: AHPFlowControlStatus +} + +public struct AHPHostedFlowControlEvent: Sendable { + public let hostId: HostId + public let event: AHPFlowControlEvent +} + +public enum AHPWindowingError: Error, Sendable { + case backpressure + case subscriptionDropped(String) + case consumerAlreadyWaiting +} + +public enum AHPDeliveryPayload: Sendable { + case event(SubscriptionEvent) + case snapshot(Snapshot) + case tcpData(TcpDataParams) + case tcpEof(TcpEofParams) + case tcpReset(TcpResetParams) + case progress(ProgressParams) + case exportLogs(OtlpExportLogsParams) + case exportTraces(OtlpExportTracesParams) + case exportMetrics(OtlpExportMetricsParams) + case dispatchAction(DispatchActionParams) + case notification(method: String, paramsJSON: Data) +} + +/// A notification, with raw parameters for extensible/custom coalescing. +/// Known messages are also decoded to native models on `AHPDelivery.payload`. +public struct AHPNotificationMessage: Sendable { + public let channel: String + public let method: String + public let paramsJSON: Data + + public init(method: String, params: P) throws { + try self.init(method: method, paramsJSON: JSONEncoder().encode(params)) + } + + public init(method: String, paramsJSON: Data) throws { + guard let params = try JSONSerialization.jsonObject(with: paramsJSON) as? [String: Any], + let channel = params["channel"] as? String else { + throw TransportError.protocol("notification has no channel") + } + self.channel = channel + self.method = method + self.paramsJSON = paramsJSON + } + + public func decode(_ type: P.Type) throws -> P { + try JSONDecoder().decode(type, from: paramsJSON) + } + + internal func wire() throws -> Data { + let params = try JSONSerialization.jsonObject(with: paramsJSON) + return try JSONSerialization.data(withJSONObject: [ + "jsonrpc": "2.0", "method": method, "params": params, + ]) + } + + internal func payload() throws -> AHPDeliveryPayload { + switch method { + case "action": return .event(.action(try decode(ActionEnvelope.self))) + case "channel/snapshot": return .snapshot(try decode(ChannelSnapshotParams.self).snapshot) + case "root/sessionAdded": return .event(.sessionAdded(try decode(SessionAddedParams.self))) + case "root/sessionRemoved": return .event(.sessionRemoved(try decode(SessionRemovedParams.self))) + case "root/sessionSummaryChanged": + return .event(.sessionSummaryChanged(try decode(SessionSummaryChangedParams.self))) + case "auth/required": return .event(.authRequired(try decode(AuthRequiredParams.self))) + case "tcp/data": return .tcpData(try decode(TcpDataParams.self)) + case "tcp/eof": return .tcpEof(try decode(TcpEofParams.self)) + case "tcp/reset": return .tcpReset(try decode(TcpResetParams.self)) + case "root/progress": return .progress(try decode(ProgressParams.self)) + case "otlp/exportLogs": return .exportLogs(try decode(OtlpExportLogsParams.self)) + case "otlp/exportTraces": return .exportTraces(try decode(OtlpExportTracesParams.self)) + case "otlp/exportMetrics": return .exportMetrics(try decode(OtlpExportMetricsParams.self)) + case "dispatchAction": return .dispatchAction(try decode(DispatchActionParams.self)) + default: return .notification(method: method, paramsJSON: paramsJSON) + } + } +} + +public struct AHPCoalescingContext: Sendable { + public let channel: String + public let direction: String +} + +public struct AHPCoalescer: Sendable { + public let method: String + public let combine: @Sendable ( + AHPNotificationMessage, AHPNotificationMessage, AHPCoalescingContext + ) throws -> AHPNotificationMessage? + + public init(method: String, combine: @escaping @Sendable ( + AHPNotificationMessage, AHPNotificationMessage, AHPCoalescingContext + ) throws -> AHPNotificationMessage?) { + self.method = method + self.combine = combine + } +} + +public struct AHPDelivery: Sendable { + public let message: AHPNotificationMessage + public let payload: AHPDeliveryPayload + public let receivedBytes: Int + internal let epoch: UInt64 + internal let releaseBytes: Int + internal let releaseBody: @Sendable () async -> Void + + public func release() async { await releaseBody() } +} + +internal actor AHPReleaseToken { + var released = false + let body: @Sendable () async -> Void + init(_ body: @escaping @Sendable () async -> Void) { self.body = body } + func release() async { + guard !released else { return } + released = true + await body() + } +} + +/// One reliable owner queue. Auto mode releases the previous item on `next()`. +public actor AHPSubscription { + public nonisolated let channel: String + public private(set) var result = SubscribeResult() + public private(set) var options: AHPSubscriptionOptions + public private(set) var lastAppliedServerSeq: Int = 0 + internal var managed = false + internal var pendingWindowOffer = false + private var engine: AHPWindowingEngine? + private var queue: [AHPDelivery] = [] + private var nextWaiter: (UUID, CheckedContinuation)? + private var consumptionOwner: UUID? + private let directOwner = UUID() + private var previous: AHPDelivery? + private var error: Error? + private var finished = false + private var flowStatus = AHPFlowControlStatus() + private var statusListeners: [UUID: AsyncStream.Continuation] = [:] + private var writableWaiters: [UUID: CheckedContinuation] = [:] + private var observer: (@Sendable (AHPDeliveryPayload) async -> Void)? + private var closeObserver: (@Sendable () async -> Void)? + private var heldBytes = 0 + private var statusRevision: UInt64 = 0 + + internal init(channel: String, options: AHPSubscriptionOptions) { + self.channel = channel + self.options = options + } + + public nonisolated var deliveries: AHPDeliverySequence { .init(subscription: self, manual: true) } + public nonisolated var events: AHPDeliverySequence { .init(subscription: self, manual: false) } + public func status() -> AHPFlowControlStatus { flowStatus } + public var statusChanges: AsyncStream { + let id = UUID() + return AsyncStream(bufferingPolicy: .bufferingNewest(1)) { continuation in + statusListeners[id] = continuation + continuation.yield(flowStatus) + continuation.onTermination = { [weak self] _ in + Task { await self?.removeStatusListener(id) } + } + } + } + + public func next() async throws -> AHPDelivery? { + try await nextDelivery(manual: options.window?.receive == .manual, owner: directOwner) + } + + /// Begin waiting for the next delivery without taking it (auto mode). + public func advance() async { + let item = previous + previous = nil + await item?.release() + } + + public func waitUntilWritable() async throws { + try Task.checkCancellation() + if let error { throw error } + if engine == nil { throw AHPClientError.shutdown } + if flowStatus.writable { return } + let id = UUID() + try await withTaskCancellationHandler { + try await withCheckedThrowingContinuation { writableWaiters[id] = $0 } + } onCancel: { + Task { await self.cancelWritable(id) } + } + } + + public func send(method: String, params: P) async throws { + let message = try AHPNotificationMessage(method: method, params: params) + try await send(message) + } + + public func send(_ message: AHPNotificationMessage) async throws { + guard let engine else { throw AHPClientError.shutdown } + guard message.channel == channel else { throw TransportError.protocol("send channel mismatch") } + try await engine.send(message) + } + + @discardableResult + public func dispatch(_ action: StateAction, clientSeq: Int? = nil) async throws -> DispatchHandle { + guard let engine else { throw AHPClientError.shutdown } + return try await engine.dispatch(action, channel: channel, clientSeq: clientSeq) + } + + public func close() async throws { + await closeObserver?() + if let engine { try await engine.unsubscribe(channel) } + else { finish(nil) } + } + + @discardableResult + public func update(options: AHPSubscriptionOptions) async throws -> SubscribeResult { + guard let engine else { throw AHPClientError.shutdown } + let subscription = try await engine.update(channel: channel, options: options) + return await subscription.result + } + + internal func bind(_ engine: AHPWindowingEngine) { + guard !finished else { return } + if self.engine !== engine { statusRevision = 0 } + self.engine = engine + error = nil + } + internal func install(_ result: SubscribeResult, options: AHPSubscriptionOptions) { + self.result = result + self.options = options + pendingWindowOffer = false + } + internal func markManaged() { managed = true } + internal func setPendingWindowOffer(_ value: Bool) { pendingWindowOffer = value } + internal func setObserver(_ observer: @escaping @Sendable (AHPDeliveryPayload) async -> Void) { + self.observer = observer + } + internal func setCloseObserver(_ observer: @escaping @Sendable () async -> Void) { + closeObserver = observer + } + internal func removeManagedRegistration() async { await closeObserver?() } + internal func observe(_ payload: AHPDeliveryPayload) async { await observer?(payload) } + + internal func enqueue(_ source: AHPDelivery, coalescers: [AHPCoalescer], coalescing: Bool) async throws -> Bool { + guard !finished else { return false } + let token = AHPReleaseToken { [weak self] in + await self?.releaseRetained(source.releaseBytes) + await source.release() + } + let delivery = AHPDelivery(message: source.message, payload: source.payload, + receivedBytes: source.receivedBytes, epoch: source.epoch, + releaseBytes: source.releaseBytes, releaseBody: { await token.release() }) + await observer?(delivery.payload) + guard !finished else { return false } + if let waiter = nextWaiter { + retain(delivery) + nextWaiter = nil + waiter.1.resume(returning: delivery) + return false + } + if coalescing, let tail = queue.last, tail.epoch == delivery.epoch, + let merged = try coalesce(tail.message, delivery.message, direction: "inbound", custom: coalescers) { + let payload = try merged.payload() + retain(delivery) + queue.removeLast() + let token = AHPReleaseToken { await tail.release(); await delivery.release() } + queue.append(AHPDelivery(message: merged, payload: payload, + receivedBytes: tail.receivedBytes + delivery.receivedBytes, + epoch: delivery.epoch, releaseBytes: tail.releaseBytes + delivery.releaseBytes, + releaseBody: { await token.release() })) + return true + } else { + retain(delivery) + queue.append(delivery) + return false + } + } + + private func retain(_ delivery: AHPDelivery) { + heldBytes += delivery.releaseBytes + switch delivery.payload { + case .snapshot(let snapshot): lastAppliedServerSeq = snapshot.fromSeq + case .event(.action(let action)): lastAppliedServerSeq = max(lastAppliedServerSeq, action.serverSeq) + default: break + } + } + + internal func discardQueuedData() async { + let items = queue.filter { + if case .snapshot = $0.payload { return false } + return true + } + queue.removeAll { + if case .snapshot = $0.payload { return false } + return true + } + for item in items { await item.release() } + } + + private func releaseRetained(_ bytes: Int) { + heldBytes -= bytes + flowStatus.bufferedReceiveBytes = heldBytes + for listener in statusListeners.values { listener.yield(flowStatus) } + } + + internal func updateStatus(_ status: AHPFlowControlStatus, engine: AHPWindowingEngine, + revision: UInt64) -> AHPFlowControlStatus? { + // Actor hops can complete out of order, including after a reconnect. + guard self.engine === engine, revision >= statusRevision else { return nil } + statusRevision = revision + flowStatus = status + flowStatus.bufferedReceiveBytes = heldBytes + for continuation in statusListeners.values { continuation.yield(flowStatus) } + if status.writable { + let waiters = writableWaiters.values + writableWaiters.removeAll() + for waiter in waiters { waiter.resume() } + } + return flowStatus + } + + internal func detach() { + engine = nil + failWritable(AHPClientError.shutdown) + } + + internal func finish(_ error: Error?) { + engine = nil + self.error = error + finished = true + flowStatus.writable = false + if let waiter = nextWaiter { + nextWaiter = nil + if let error { waiter.1.resume(throwing: error) } else { waiter.1.resume(returning: nil) } + } + failWritable(error ?? AHPClientError.shutdown) + for c in statusListeners.values { c.yield(flowStatus); c.finish() } + statusListeners.removeAll() + } + + private func nextDelivery(manual: Bool, owner: UUID) async throws -> AHPDelivery? { + try Task.checkCancellation() + if let consumptionOwner, consumptionOwner != owner { throw AHPWindowingError.consumerAlreadyWaiting } + consumptionOwner = owner + if !manual { await advance() } + let item: AHPDelivery? + if !queue.isEmpty { item = queue.removeFirst() } + else if let error { throw error } + else if finished { return nil } + else { + guard nextWaiter == nil else { throw AHPWindowingError.consumerAlreadyWaiting } + let ticket = UUID() + item = try await withTaskCancellationHandler { + try await withCheckedThrowingContinuation { nextWaiter = (ticket, $0) } + } onCancel: { + Task { await self.cancelNext(ticket) } + } + } + if !manual { previous = item } + return item + } + + private func cancelNext(_ id: UUID) { + guard nextWaiter?.0 == id else { return } + nextWaiter?.1.resume(throwing: CancellationError()) + nextWaiter = nil + } + private func cancelWritable(_ id: UUID) { + writableWaiters.removeValue(forKey: id)?.resume(throwing: CancellationError()) + } + private func failWritable(_ error: Error) { + let waiters = writableWaiters.values + writableWaiters.removeAll() + for waiter in waiters { waiter.resume(throwing: error) } + } + private func removeStatusListener(_ id: UUID) { statusListeners.removeValue(forKey: id) } +} + +public struct AHPDeliverySequence: AsyncSequence, Sendable { + public typealias Element = AHPDelivery + internal let subscription: AHPSubscription + internal let manual: Bool + public struct AsyncIterator: AsyncIteratorProtocol { + internal let subscription: AHPSubscription + internal let manual: Bool + internal let owner: UUID + public mutating func next() async throws -> AHPDelivery? { + try await subscription.nextFromIterator(manual: manual, owner: owner) + } + } + public func makeAsyncIterator() -> AsyncIterator { .init(subscription: subscription, manual: manual, owner: UUID()) } +} + +extension AHPSubscription { + fileprivate func nextFromIterator(manual: Bool, owner: UUID) async throws -> AHPDelivery? { + try await nextDelivery(manual: manual || options.window?.receive == .manual, owner: owner) + } +} + +internal func coalesce(_ earlier: AHPNotificationMessage, _ later: AHPNotificationMessage, + direction: String, custom: [AHPCoalescer]) throws -> AHPNotificationMessage? { + guard earlier.channel == later.channel, earlier.method == later.method, + countsAgainstWindow(earlier.method), + !["channel/snapshot", "channel/frame", "dispatchAction", "tcp/eof"].contains(earlier.method) + else { return nil } + var a = try JSONSerialization.jsonObject(with: earlier.paramsJSON) as! [String: Any] + let b = try JSONSerialization.jsonObject(with: later.paramsJSON) as! [String: Any] + guard a["origin"] == nil, b["origin"] == nil, + a["rejectionReason"] == nil, b["rejectionReason"] == nil else { return nil } + for c in custom where c.method == earlier.method { + if let result = try c.combine(earlier, later, .init(channel: earlier.channel, direction: direction)) { + guard result.channel == earlier.channel, result.method == earlier.method, + countsAgainstWindow(result.method) else { + throw TransportError.protocol("custom coalescer changed channel or method") + } + return result + } + } + if direction == "outbound", earlier.method == "tcp/data", + let x = a["data"] as? String, let y = b["data"] as? String, + var data = Data(base64Encoded: x), let more = Data(base64Encoded: y) { + data.append(more) + a["data"] = data.base64EncodedString() + } else if direction == "inbound", earlier.method == "action", + var x = a["action"] as? [String: Any], let y = b["action"] as? [String: Any], + let type = x["type"] as? String, type == y["type"] as? String { + if type == "terminal/data", let one = x["data"] as? String, let two = y["data"] as? String, + try sameJSON(x.filter { $0.key != "data" }, y.filter { $0.key != "data" }) { + x["data"] = one + two + } else if type == "chat/delta", let one = x["content"] as? String, + let two = y["content"] as? String, + try sameJSON(x.filter { $0.key != "content" }, y.filter { $0.key != "content" }) { + x["content"] = one + two + } else { return nil } + a["action"] = x + a["serverSeq"] = b["serverSeq"] + } else { return nil } + return try .init(method: earlier.method, paramsJSON: JSONSerialization.data(withJSONObject: a)) +} + +internal func ahpUTF16Bytes(_ text: String) -> Int { text.utf16.count * 2 } + +private func sameJSON(_ a: [String: Any], _ b: [String: Any]) throws -> Bool { + try JSONSerialization.data(withJSONObject: a, options: [.sortedKeys]) == + JSONSerialization.data(withJSONObject: b, options: [.sortedKeys]) +} diff --git a/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/WindowingEngine.swift b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/WindowingEngine.swift new file mode 100644 index 000000000..3a46997c9 --- /dev/null +++ b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/WindowingEngine.swift @@ -0,0 +1,875 @@ +import Foundation +import AgentHostProtocol + +internal actor AHPWindowingEngine { + private final class SendEntry { + let id: UUID + var message: AHPNotificationMessage + var wire: String + var bytes: Int + var bufferedBytes: Int + var headAllowance: Int + var scalars: [Unicode.Scalar]? + var offset = 0 + var started = false + var emitting = false + var emissionID: UUID? + var replies: [CheckedContinuation] + var ids: [UUID] + var messages: [AHPNotificationMessage] + init(id: UUID, message: AHPNotificationMessage, wire: String, + reply: CheckedContinuation) { + self.id = id + self.message = message + self.wire = wire + self.bytes = ahpUTF16Bytes(wire) + self.bufferedBytes = self.bytes + self.headAllowance = self.bytes + self.replies = [reply] + self.ids = [id] + self.messages = [message] + } + } + + private final class Channel { + let subscription: AHPSubscription + var options: AHPSubscriptionOptions + var pendingRequests = 0 + var installed = false + var ready = false + var hasConsumer = false + var fallback: String? + var limits: ChannelFlowControl? + var epoch = UInt64.random(in: 0...UInt64.max) + var sent = 0 + var credited = 0 + var received = 0 + var released = 0 + var creditSent = 0 + var creditTask: Task? + var awaitingSnapshot = false + var assembly: String? + var assemblyBytes = 0 + var sends: [SendEntry] = [] + var buffered = 0 + var sentEOF = false + var receivedEOF = false + var reset = false + var pauseNewData = false + var modeWaiters: [UUID: CheckedContinuation] = [:] + var statusRevision: UInt64 = 0 + var lastWritable = true + init(_ subscription: AHPSubscription, options: AHPSubscriptionOptions) { + self.subscription = subscription + self.options = options + } + } + + private let config: AHPWindowingConfig + private let writeBody: @Sendable (Data, AHPBeforeWrite?, Bool) async throws -> Void + private let fail: @Sendable (Error) async -> Void + private let dispatchBody: @Sendable (StateAction, String, Int?) async throws -> DispatchHandle + private let updateBody: @Sendable (String, AHPSubscriptionOptions) async throws -> AHPSubscription + private var channels: [String: Channel] = [:] + private var order: [String] = [] + private var cursor = 0 + private var scheduler: Task? + private var closed = false + private var telemetry: [UUID: AsyncStream.Continuation] = [:] + + init(config: AHPWindowingConfig, + write: @escaping @Sendable (Data, AHPBeforeWrite?, Bool) async throws -> Void, + fail: @escaping @Sendable (Error) async -> Void, + dispatch: @escaping @Sendable (StateAction, String, Int?) async throws -> DispatchHandle, + update: @escaping @Sendable (String, AHPSubscriptionOptions) async throws -> AHPSubscription) { + self.config = config + self.writeBody = write + self.fail = fail + self.dispatchBody = dispatch + self.updateBody = update + } + + var events: AsyncStream { + let id = UUID() + return AsyncStream(bufferingPolicy: .bufferingNewest(256)) { continuation in + telemetry[id] = continuation + continuation.onTermination = { [weak self] _ in Task { await self?.removeTelemetry(id) } } + } + } + private func removeTelemetry(_ id: UUID) { telemetry.removeValue(forKey: id) } + + func prepare(channel: String, options: AHPSubscriptionOptions, + capabilities: AHPTransportCapabilities, + retained: AHPSubscription? = nil) async throws -> (AHPSubscription, ChannelReceiveLimits?, Bool) { + guard !closed else { throw AHPClientError.shutdown } + let offer: ChannelReceiveLimits? + if let window = options.window, capabilities.preservesOriginalJSONText { + let limits = ChannelReceiveLimits( + windowBytes: window.windowBytes ?? config.windowBytes, + maximumFrameBytes: window.maximumFrameBytes ?? config.maximumFrameBytes ?? capabilities.maximumFrameBytes + ) + try validate(limits) + guard config.maxBufferedSendBytes > 0 else { + throw TransportError.protocol("maxBufferedSendBytes must be positive") + } + offer = limits + } else { offer = nil } + if channel.hasPrefix("ahp-tcp:"), offer == nil { + throw TransportError.protocol("TCP requires a window-capable transport and a window offer") + } + let created = channels[channel] == nil + let record: Channel + if let existing = channels[channel] { record = existing } + else { + let subscription = retained ?? AHPSubscription(channel: channel, options: options) + record = Channel(subscription, options: options) + record.installed = retained != nil + record.hasConsumer = retained != nil + if retained != nil { record.fallback = await subscription.status().fallbackReason } + channels[channel] = record + if !order.contains(channel) { order.append(channel) } + } + if record.ready { await record.subscription.bind(self) } + await record.subscription.setPendingWindowOffer(offer != nil) + return (record.subscription, offer, created) + } + + func prepareGeneric(_ params: Data, capabilities: AHPTransportCapabilities) async throws { + let p = try JSONDecoder().decode(SubscribeParams.self, from: params) + if channels[p.channel] == nil, let offer = p.flowControl?.receive { + _ = try await prepare(channel: p.channel, options: .windowed( + windowBytes: offer.windowBytes, maximumFrameBytes: offer.maximumFrameBytes, + delivery: p.delivery, view: p.view), capabilities: capabilities) + } + channels[p.channel]?.pendingRequests += 1 + } + + func prepareReconnect(_ params: Data, capabilities: AHPTransportCapabilities) async throws { + let request = try JSONDecoder().decode(ReconnectParams.self, from: params) + guard let windows = request.windows else { return } + guard capabilities.preservesOriginalJSONText, + Set(windows.items.map(\.channel)).count == windows.items.count else { + throw TransportError.protocol("reconnect windows require raw text and unique channels") + } + for item in windows.items { + try validate(item.flowControl.receive) + guard request.subscriptions.contains(item.channel), !item.channel.hasPrefix("ahp-tcp:") else { + throw TransportError.protocol("invalid reconnect window channel") + } + if channels[item.channel] == nil { + _ = try await prepare(channel: item.channel, options: .windowed( + windowBytes: item.flowControl.receive.windowBytes, + maximumFrameBytes: item.flowControl.receive.maximumFrameBytes), capabilities: capabilities) + } + } + } + + func optionsForRequest(_ request: SubscribeParams) -> AHPSubscriptionOptions { + var options = channels[request.channel]?.options ?? .init() + options.delivery = request.delivery + options.view = request.view + options.window = request.flowControl.map { + .init(windowBytes: $0.receive.windowBytes, maximumFrameBytes: $0.receive.maximumFrameBytes, + receive: options.window?.receive ?? .auto, coalescing: options.window?.coalescing ?? true) + } + return options + } + + func pauseForOrdinarySubscribe(_ channel: String) async throws { + try Task.checkCancellation() + guard let r = channels[channel], r.limits != nil else { return } + r.pauseNewData = true + guard let entry = r.sends.first, entry.started || entry.emitting, entry.scalars != nil else { return } + let id = UUID() + try await withTaskCancellationHandler { + try await withCheckedThrowingContinuation { r.modeWaiters[id] = $0 } + } onCancel: { + Task { await self.cancelModeWait(channel, id: id) } + } + } + + private func cancelModeWait(_ channel: String, id: UUID) { + guard let r = channels[channel] else { return } + r.modeWaiters.removeValue(forKey: id)?.resume(throwing: CancellationError()) + r.pauseNewData = false + startScheduler() + } + + func abortSubscribe(_ channel: String, error: Error? = nil, expected: AHPSubscription? = nil) async { + guard let r = channels[channel] else { return } + if let expected, r.subscription !== expected { return } + r.pendingRequests = max(0, r.pendingRequests - 1) + r.pauseNewData = false + if r.pendingRequests == 0 { await r.subscription.setPendingWindowOffer(false) } + if !r.installed && r.pendingRequests == 0 { + await drop(channel, error: error ?? AHPClientError.shutdown) + } + startScheduler() + } + + func abortPrepared(_ subscription: AHPSubscription, error: Error) async { + guard let r = channels[subscription.channel], r.subscription === subscription, + !r.installed, r.pendingRequests == 0 else { return } + await drop(subscription.channel, error: error) + } + + func abandonSubscribe(_ subscription: AHPSubscription, created: Bool, error: Error) async { + guard created, let r = channels[subscription.channel], r.subscription === subscription, + !r.installed, r.pendingRequests <= 1 else { return } + await drop(subscription.channel, error: error) + } + + func cleanAbandoned(_ subscription: AHPSubscription, created: Bool) async { + guard !closed, created else { return } + if let r = channels[subscription.channel] { + guard r.subscription === subscription, r.pendingRequests == 0, !r.hasConsumer else { return } + await unsubscribeIfCurrent(subscription) + } else { + try? await write(AHPNotificationMessage(method: "unsubscribe", + params: UnsubscribeParams(channel: subscription.channel)).wire()) + } + } + + private func unsubscribeIfCurrent(_ subscription: AHPSubscription) async { + do { try await subscription.close() } + catch is CancellationError {} + catch AHPWindowingError.backpressure {} + catch { await fail(error) } + } + + func installSubscribe(request: SubscribeParams, options: AHPSubscriptionOptions?, + resultData: Data, expected: AHPSubscription? = nil, + retain: Bool = true) async throws { + let channel = request.channel + let result = try JSONDecoder().decode(SubscribeResult.self, from: resultData) + if let limits = result.flowControl { + guard result.snapshot == nil, let offer = request.flowControl?.receive else { + throw TransportError.protocol("unsolicited flow control or snapshot in accepted result") + } + try validateAccepted(limits, offer: offer) + } else if channel.hasPrefix("ahp-tcp:") { + throw TransportError.protocol("TCP window offer declined") + } + guard let r = channels[channel] else { return } + if let expected, r.subscription !== expected { return } + r.pendingRequests = max(0, r.pendingRequests - 1) + if (r.limits == nil) != (result.flowControl == nil) { + guard r.assembly == nil else { + throw TransportError.protocol("subscription mode changed during inbound fragments") + } + r.epoch &+= 1 + r.sent = 0; r.credited = 0; r.received = 0; r.released = 0; r.creditSent = 0 + r.assembly = nil; r.assemblyBytes = 0 + r.creditTask?.cancel(); r.creditTask = nil + } + r.limits = result.flowControl + r.installed = true + if retain { r.hasConsumer = true } + r.ready = true + r.pauseNewData = false + r.options = options ?? optionsForRequest(request) + r.fallback = r.options.window == nil ? nil : (request.flowControl == nil ? "transport" : "declined") + r.awaitingSnapshot = result.flowControl != nil + await r.subscription.bind(self) + await r.subscription.install(result, options: r.options) + if let snapshot = result.snapshot { + let message = try AHPNotificationMessage( + method: "channel/snapshot", params: ChannelSnapshotParams(channel: channel, snapshot: snapshot)) + try await enqueue(r, message: message, bytes: 0, releaseBytes: 0) + } + await publish(channel, kind: r.limits == nil ? "declined" : "negotiated") + startScheduler() + } + + func installReconnect(request: Data, resultData: Data) async throws { + let p = try JSONDecoder().decode(ReconnectParams.self, from: request) + let result = try JSONDecoder().decode(ReconnectResult.self, from: resultData) + for channel in p.subscriptions where !(p.windows?.items.contains(where: { $0.channel == channel }) ?? false) { + guard let r = channels[channel] else { continue } + r.ready = true + await r.subscription.bind(self) + await publish(channel, kind: "resumed") + } + switch result { + case .replay(let value): + for action in value.actions { + if let r = channels[action.channel], r.limits == nil, + !(p.windows?.items.contains(where: { $0.channel == action.channel }) ?? false) { + let message = try AHPNotificationMessage(method: "action", params: action) + try await enqueue(r, message: message, bytes: 0, releaseBytes: 0) + } + } + for channel in value.missing { + await drop(channel, error: AHPWindowingError.subscriptionDropped(channel)) + } + case .snapshot(let value): + for snapshot in value.snapshots { + if let r = channels[snapshot.resource], r.limits == nil, + !(p.windows?.items.contains(where: { $0.channel == snapshot.resource }) ?? false) { + let message = try AHPNotificationMessage(method: "channel/snapshot", + params: ChannelSnapshotParams(channel: snapshot.resource, snapshot: snapshot)) + try await enqueue(r, message: message, bytes: 0, releaseBytes: 0) + } + } + } + guard let requested = p.windows else { return } + let items: [ResumedChannelSubscription] + switch result { + case .replay(let value): items = value.windows?.items ?? [] + case .snapshot(let value): items = value.windows?.items ?? [] + } + guard Set(items.map(\.channel)).count == items.count else { + throw TransportError.protocol("duplicate reconnect window") + } + for item in items { + guard let offered = requested.items.first(where: { $0.channel == item.channel }), + p.subscriptions.contains(item.channel), let r = channels[item.channel] else { + throw TransportError.protocol("unrequested reconnect window") + } + try validateAccepted(item.flowControl, offer: offered.flowControl.receive) + r.limits = item.flowControl + r.ready = true + r.epoch &+= 1 + r.sent = 0; r.credited = 0; r.received = 0; r.released = 0; r.creditSent = 0 + r.assembly = nil; r.assemblyBytes = 0 + r.awaitingSnapshot = item.recovery == .snapshot + await r.subscription.bind(self) + await r.subscription.install(SubscribeResult(flowControl: item.flowControl), options: r.options) + await publish(item.channel, kind: "resumed") + } + for offered in requested.items where !items.contains(where: { $0.channel == offered.channel }) { + await drop(offered.channel, error: AHPWindowingError.subscriptionDropped(offered.channel)) + } + } + + func subscription(_ channel: String) -> AHPSubscription? { channels[channel]?.subscription } + func bookkeepingCount() -> Int { order.count } + func isWindowed(_ channel: String) -> Bool { channels[channel]?.limits != nil } + var hasWindows: Bool { channels.values.contains { $0.limits != nil } } + + func receive(_ message: AHPNotificationMessage, rawText: String?) async throws -> AHPNotificationMessage? { + guard !closed else { return nil } + guard let r = channels[message.channel] else { return message } + if !r.ready && message.method != "channel/credit" { + // An unsubscribe can race with already-sent data from the previous + // subscription. The next result is the ordered receive boundary; + // until then, that traffic does not belong to the new owner. + if ["channel/frame", "channel/snapshot"].contains(message.method) { return nil } + await r.subscription.observe(try message.payload()) + return message + } + if r.reset && message.method != "channel/credit" { return nil } + if message.method == "channel/credit" { + guard r.limits != nil else { return nil } + let credit = try message.decode(ChannelCreditParams.self).consumedBytes + guard credit >= 0, credit <= r.sent else { throw TransportError.protocol("credit exceeds sent bytes") } + if credit > r.credited { + r.credited = credit + await publish(message.channel, kind: "creditReceived") + startScheduler() + } + return nil + } + guard let limits = r.limits else { + guard !["channel/frame", "channel/snapshot"].contains(message.method) else { + throw TransportError.protocol("window delivery message on ordinary subscription") + } + try await enqueue(r, message: message, bytes: 0, releaseBytes: 0) + return message + } + let counted = countsAgainstWindow(message.method) + guard !counted || rawText != nil else { throw TransportError.protocol("missing original JSON text") } + let bytes = counted ? ahpUTF16Bytes(rawText!) : 0 + guard bytes <= limits.clientReceive.maximumFrameBytes else { + throw TransportError.protocol("received data exceeds maximumFrameBytes") + } + r.received = try adding(r.received, bytes) + if message.method == "tcp/data" { + guard !r.receivedEOF, !r.reset else { throw TransportError.protocol("TCP data after EOF/reset") } + try validateTCPData(message) + } else if message.method == "tcp/eof" { + guard !r.receivedEOF, !r.reset else { throw TransportError.protocol("duplicate TCP EOF/reset") } + r.receivedEOF = true + } else if message.method == "tcp/reset" { + r.reset = true + r.epoch &+= 1 + r.assembly = nil; r.assemblyBytes = 0 + await r.subscription.discardQueuedData() + } + if message.method == "channel/frame" { + let frame = try message.decode(ChannelFrameParams.self) + guard !frame.data.isEmpty else { throw TransportError.protocol("empty channel/frame") } + r.assembly = (r.assembly ?? "") + frame.data + r.assemblyBytes = try adding(r.assemblyBytes, bytes) + if frame.final != true { + r.released = try adding(r.released, bytes) + await flushCredit(message.channel) + await publish(message.channel, kind: message.method == "tcp/reset" ? "dropped" : nil) + return nil + } + let joined = r.assembly! + let total = r.assemblyBytes + r.assembly = nil; r.assemblyBytes = 0 + guard let dict = try JSONSerialization.jsonObject(with: Data(joined.utf8)) as? [String: Any], + dict["id"] == nil, let method = dict["method"] as? String, + method != "channel/frame", countsAgainstWindow(method), + let params = dict["params"] as? [String: Any] else { + throw TransportError.protocol("invalid joined channel/frame") + } + let joinedMessage = try AHPNotificationMessage( + method: method, paramsJSON: JSONSerialization.data(withJSONObject: params)) + guard joinedMessage.channel == message.channel else { + throw TransportError.protocol("fragment channel mismatch") + } + if method == "tcp/data" { + guard !r.receivedEOF, !r.reset else { throw TransportError.protocol("TCP data after EOF/reset") } + try validateTCPData(joinedMessage) + } else if method == "tcp/eof" { + guard !r.receivedEOF, !r.reset else { throw TransportError.protocol("duplicate TCP EOF/reset") } + r.receivedEOF = true + } + try checkSnapshot(r, message: joinedMessage) + try await enqueue(r, message: joinedMessage, bytes: total, releaseBytes: bytes) + await publish(message.channel, kind: "fragmentedMessageReceived") + return joinedMessage + } else { + guard !counted || r.assembly == nil else { + throw TransportError.protocol("data interleaved with unfinished fragments") + } + if counted { try checkSnapshot(r, message: message) } + try await enqueue(r, message: message, bytes: bytes, releaseBytes: bytes) + if message.method == "tcp/reset" { + failSends(r, error: AHPWindowingError.subscriptionDropped(message.channel)) + try await r.subscription.close() + } + await publish(message.channel) + return message + } + } + + private func checkSnapshot(_ r: Channel, message: AHPNotificationMessage) throws { + if r.awaitingSnapshot { + guard message.method == "channel/snapshot" else { + throw TransportError.protocol("first data is not channel/snapshot") + } + r.awaitingSnapshot = false + } + if message.method == "channel/snapshot" { + guard try message.decode(ChannelSnapshotParams.self).snapshot.resource == message.channel else { + throw TransportError.protocol("snapshot channel mismatch") + } + } + } + + private func enqueue(_ r: Channel, message: AHPNotificationMessage, bytes: Int, + releaseBytes: Int) async throws { + let epoch = r.epoch + let token = AHPReleaseToken { [weak self] in + await self?.release(channel: message.channel, epoch: epoch, bytes: releaseBytes) + } + let delivery = AHPDelivery(message: message, payload: try message.payload(), + receivedBytes: bytes, epoch: epoch, releaseBytes: releaseBytes, + releaseBody: { await token.release() }) + let combined = try await r.subscription.enqueue(delivery, coalescers: config.coalescers, + coalescing: r.limits != nil && r.options.window?.coalescing != false) + await publish(message.channel, kind: combined ? "coalesced" : nil) + } + + private func release(channel: String, epoch: UInt64, bytes: Int) async { + guard let r = channels[channel], r.epoch == epoch, r.limits != nil, !closed, bytes > 0 else { return } + let wasFull = r.received - r.released >= r.limits!.clientReceive.windowBytes + r.released += bytes + let threshold = min(16_384, max(1, r.limits!.clientReceive.windowBytes / 4)) + if wasFull || r.released - r.creditSent >= threshold { + await flushCredit(channel) + } else if r.creditTask == nil { + r.creditTask = Task { [weak self] in + try? await Task.sleep(for: .milliseconds(5)) + if !Task.isCancelled { await self?.flushCredit(channel) } + } + } + await publish(channel) + } + + private func flushCredit(_ channel: String) async { + guard let r = channels[channel], !closed, r.limits != nil, r.released > r.creditSent else { return } + r.creditTask?.cancel(); r.creditTask = nil + r.creditSent = r.released + do { + let message = try AHPNotificationMessage( + method: "channel/credit", params: ChannelCreditParams(channel: channel, consumedBytes: r.creditSent)) + let epoch = r.epoch + try await writeBody(message.wire(), { [weak self] in + guard let self else { throw CancellationError() } + try await self.validateCredit(channel, epoch: epoch) + }, false) + await publish(channel, kind: "creditSent") + } catch is CancellationError { + } catch AHPWindowingError.backpressure { + r.creditSent = 0 + r.creditTask = Task { [weak self] in + try? await Task.sleep(for: .milliseconds(5)) + if !Task.isCancelled { await self?.flushCredit(channel) } + } + } catch { await fail(error) } + } + + private func validateCredit(_ channel: String, epoch: UInt64) throws { + guard !closed, let r = channels[channel], r.epoch == epoch, r.limits != nil else { + throw CancellationError() + } + } + + func send(_ message: AHPNotificationMessage) async throws { + try Task.checkCancellation() + guard !closed else { throw AHPClientError.shutdown } + if let r = channels[message.channel] { + if message.method == "tcp/data" { + guard !r.reset, !r.sentEOF, !r.sends.contains(where: { $0.message.method == "tcp/eof" }) else { + throw TransportError.protocol("TCP data after EOF/reset") + } + try validateTCPData(message) + } else if message.method == "tcp/eof" { + guard !r.reset, !r.sentEOF, !r.sends.contains(where: { $0.message.method == "tcp/eof" }) else { + throw TransportError.protocol("duplicate TCP EOF/reset") + } + } else if message.method == "tcp/reset" { + r.reset = true + r.epoch &+= 1 + r.assembly = nil; r.assemblyBytes = 0 + failSends(r, error: AHPWindowingError.subscriptionDropped(message.channel)) + await r.subscription.discardQueuedData() + await publish(message.channel, kind: "dropped") + try await writeBody(message.wire(), nil, countsAgainstWindow(message.method)) + try await r.subscription.close() + return + } + } + guard let r = channels[message.channel], countsAgainstWindow(message.method), r.limits != nil else { + try await writeBody(message.wire(), nil, countsAgainstWindow(message.method)) + return + } + let id = UUID() + try await withTaskCancellationHandler { + try await withCheckedThrowingContinuation { continuation in + do { + let wire = String(decoding: try message.wire(), as: UTF8.self) + let entry = SendEntry(id: id, message: message, wire: wire, reply: continuation) + if r.options.window?.coalescing != false, let tail = r.sends.last, !tail.started, !tail.emitting, + let merged = try coalesce(tail.message, message, direction: "outbound", custom: config.coalescers) { + let mergedWire = String(decoding: try merged.wire(), as: UTF8.self) + let mergedBytes = ahpUTF16Bytes(mergedWire) + let behind = queuedBehindHead(r) + let queued = behind - tail.bytes + mergedBytes + guard queued <= config.maxBufferedSendBytes else { + throw AHPWindowingError.backpressure + } + r.buffered += mergedBytes - tail.bytes + tail.message = merged; tail.wire = mergedWire; tail.bytes = mergedBytes + tail.bufferedBytes = mergedBytes + tail.replies.append(continuation) + tail.ids.append(id) + tail.messages.append(message) + Task { await self.publish(message.channel, kind: "coalesced") } + } else { + guard r.sends.isEmpty || queuedBehindHead(r) + entry.bytes <= config.maxBufferedSendBytes else { + throw AHPWindowingError.backpressure + } + r.buffered += entry.bytes + r.sends.append(entry) + } + startScheduler() + } catch { continuation.resume(throwing: error) } + } + } onCancel: { + Task { await self.cancelSend(channel: message.channel, id: id) } + } + } + + func dispatch(_ action: StateAction, channel: String, clientSeq: Int?) async throws -> DispatchHandle { + try await dispatchBody(action, channel, clientSeq) + } + + func update(channel: String, options: AHPSubscriptionOptions) async throws -> AHPSubscription { + try await updateBody(channel, options) + } + + func unsubscribe(_ channel: String) async throws { + await drop(channel, error: nil) + try await write(AHPNotificationMessage(method: "unsubscribe", params: UnsubscribeParams(channel: channel)).wire()) + } + + func drop(_ channel: String, error: Error?) async { + guard let r = channels.removeValue(forKey: channel) else { return } + for index in order.indices.reversed() where order[index] == channel { + order.remove(at: index) + if index < cursor { cursor -= 1 } + } + cursor = order.isEmpty ? 0 : cursor % order.count + r.creditTask?.cancel() + failSends(r, error: error ?? AHPClientError.shutdown) + await r.subscription.removeManagedRegistration() + await r.subscription.finish(error) + var value = status(r) + value.writable = false + let event = AHPFlowControlEvent(channel: channel, kind: "dropped", status: value) + for listener in telemetry.values { listener.yield(event) } + } + + func disconnect(_ error: Error) async { + guard !closed else { return } + closed = true + scheduler?.cancel(); scheduler = nil + for (channel, r) in channels { + r.creditTask?.cancel() + r.assembly = nil + failSends(r, error: error) + await publish(channel, kind: channel.hasPrefix("ahp-tcp:") ? "dropped" : nil) + if await r.subscription.managed, !channel.hasPrefix("ahp-tcp:") { + await r.subscription.detach() + } else { + await r.subscription.removeManagedRegistration() + await r.subscription.finish(error) + } + } + for listener in telemetry.values { listener.finish() } + telemetry.removeAll() + } + + private func startScheduler() { + guard scheduler == nil, !closed else { return } + scheduler = Task { await self.drain() } + } + + private func drain() async { + for channel in order { await publish(channel) } + while !closed && !Task.isCancelled { + var selected: (String, Channel, SendEntry)? + for _ in 0..= limits.hostReceive.windowBytes { continue } + selected = (channel, r, entry) + break + } + guard let (channel, r, entry) = selected else { break } + let oldOffset = entry.offset + do { + let (wire, final, originalBytes) = try nextFrame(entry, channel: channel, record: r) + entry.emitting = true + let id = UUID() + entry.emissionID = id + let epoch = r.epoch + let bytes = ahpUTF16Bytes(wire) + try await writeBody(Data(wire.utf8), { [weak self] in + guard let self else { throw AHPClientError.shutdown } + try await self.commitSend(channel: channel, id: id, epoch: epoch, bytes: bytes) + }, false) + entry.emitting = false + if r.sends.first === entry { + r.buffered -= originalBytes + entry.bufferedBytes -= originalBytes + entry.headAllowance = max(0, entry.headAllowance - originalBytes) + } + if final { + if r.sends.first === entry { + r.sends.removeFirst() + if let head = r.sends.first { head.headAllowance = head.bufferedBytes } + } + let replies = entry.replies + entry.replies.removeAll() + for reply in replies { reply.resume() } + let waiters = r.modeWaiters.values + r.modeWaiters.removeAll() + for waiter in waiters { waiter.resume() } + await publish(channel, kind: entry.scalars == nil ? nil : "fragmentedMessageSent") + } else { + await publish(channel) + } + } catch is EmissionDeferred { + entry.offset = oldOffset + entry.emitting = false + if !entry.started { entry.scalars = nil } + } catch is CancellationError { + entry.emitting = false + } catch AHPWindowingError.backpressure { + failSends(r, error: AHPWindowingError.backpressure) + await publish(channel) + } catch { + failSends(r, error: error) + await fail(error) + break + } + } + scheduler = nil + } + + private struct EmissionDeferred: Error {} + + private func write(_ data: Data) async throws { try await writeBody(data, nil, false) } + + private func commitSend(channel: String, id: UUID, epoch: UInt64, bytes: Int) async throws { + guard !closed, let r = channels[channel], let entry = r.sends.first, entry.emissionID == id else { + throw CancellationError() + } + guard r.epoch == epoch else { throw EmissionDeferred() } + if let limits = r.limits { + guard r.sent - r.credited < limits.hostReceive.windowBytes, + bytes <= limits.hostReceive.maximumFrameBytes else { throw EmissionDeferred() } + r.sent = try adding(r.sent, bytes) + } + entry.started = true + if entry.message.method == "tcp/eof" { r.sentEOF = true } + await publish(channel) + } + + private func nextFrame(_ entry: SendEntry, channel: String, record: Channel) throws -> (String, Bool, Int) { + guard let limits = record.limits else { return (entry.wire, true, entry.bufferedBytes) } + let maximum = limits.hostReceive.maximumFrameBytes + if entry.scalars == nil && entry.bytes <= maximum { return (entry.wire, true, entry.bufferedBytes) } + if entry.scalars == nil { entry.scalars = Array(entry.wire.unicodeScalars) } + let scalars = entry.scalars! + func frame(_ count: Int) throws -> String { + let text = String(String.UnicodeScalarView(scalars[entry.offset..<(entry.offset + count)])) + let message = try AHPNotificationMessage(method: "channel/frame", params: ChannelFrameParams( + channel: channel, data: text, final: entry.offset + count == scalars.count ? true : nil)) + return String(decoding: try message.wire(), as: UTF8.self) + } + let smallest = try frame(1) + guard ahpUTF16Bytes(smallest) <= maximum else { + throw TransportError.protocol("frame envelope cannot carry a scalar") + } + let remaining = limits.hostReceive.windowBytes - (record.sent - record.credited) + let target = min(maximum, max(remaining, ahpUTF16Bytes(smallest))) + var low = 1 + var high = min(scalars.count - entry.offset, target / 2) + while low < high { + let middle = (low + high + 1) / 2 + if ahpUTF16Bytes(try frame(middle)) <= target { low = middle } else { high = middle - 1 } + } + let text = try frame(low) + let original = String(String.UnicodeScalarView(scalars[entry.offset..<(entry.offset + low)])) + entry.offset += low + return (text, entry.offset == scalars.count, ahpUTF16Bytes(original)) + } + + private func cancelSend(channel: String, id: UUID) async { + guard let r = channels[channel], let index = r.sends.firstIndex(where: { $0.ids.contains(id) }), + !r.sends[index].started else { return } + let entry = r.sends.remove(at: index) + r.buffered -= entry.bytes + let removed = entry.ids.firstIndex(of: id)! + entry.ids.remove(at: removed) + entry.messages.remove(at: removed) + entry.replies.remove(at: removed).resume(throwing: CancellationError()) + do { + var replacements: [SendEntry] = [] + for source in entry.messages.indices { + let message = entry.messages[source] + if let tail = replacements.last, + let combined = try coalesce(tail.message, message, direction: "outbound", custom: config.coalescers) { + tail.message = combined + tail.wire = String(decoding: try combined.wire(), as: UTF8.self) + tail.bytes = ahpUTF16Bytes(tail.wire) + tail.bufferedBytes = tail.bytes + tail.replies.append(entry.replies[source]) + tail.ids.append(entry.ids[source]) + tail.messages.append(message) + } else { + replacements.append(SendEntry(id: entry.ids[source], message: message, + wire: String(decoding: try message.wire(), as: UTF8.self), reply: entry.replies[source])) + } + } + let bytes = replacements.reduce(0) { $0 + $1.bytes } + let head = index == 0 ? replacements.first : r.sends.first + let headBytes = head.map { min($0.bufferedBytes, $0.headAllowance) } ?? 0 + guard r.buffered + bytes - headBytes <= config.maxBufferedSendBytes else { + throw AHPWindowingError.backpressure + } + r.sends.insert(contentsOf: replacements, at: index) + r.buffered += bytes + } catch { + for reply in entry.replies { reply.resume(throwing: error) } + } + entry.replies.removeAll() + if entry.scalars != nil { + let waiters = r.modeWaiters.values + r.modeWaiters.removeAll() + for waiter in waiters { waiter.resume() } + } + await publish(channel) + } + + private func failSends(_ r: Channel, error: Error) { + let modeWaiters = r.modeWaiters.values + r.modeWaiters.removeAll() + for waiter in modeWaiters { waiter.resume(throwing: error) } + let entries = r.sends + r.sends.removeAll(); r.buffered = 0 + for entry in entries { + for reply in entry.replies { reply.resume(throwing: error) } + entry.replies.removeAll() + } + } + + private func status(_ r: Channel) -> AHPFlowControlStatus { + var value = AHPFlowControlStatus() + value.mode = r.limits == nil ? "ordinary" : "windowed" + value.fallbackReason = r.limits == nil ? r.fallback : nil + value.limits = r.limits + value.outstandingSendBytes = r.sent - r.credited + value.bufferedSendBytes = r.buffered + value.bufferedReceiveBytes = r.received - r.released + value.writable = !closed && !r.reset && queuedBehindHead(r) < config.maxBufferedSendBytes && + (r.limits.map { r.sent - r.credited < $0.hostReceive.windowBytes } ?? true) + return value + } + + private func queuedBehindHead(_ r: Channel) -> Int { + r.buffered - (r.sends.first.map { min($0.bufferedBytes, $0.headAllowance) } ?? 0) + } + + private func publish(_ channel: String, kind: String? = nil) async { + guard let r = channels[channel] else { return } + let snapshot = status(r) + let transition = r.lastWritable == snapshot.writable ? nil : (snapshot.writable ? "unblocked" : "blocked") + r.lastWritable = snapshot.writable + r.statusRevision &+= 1 + guard let value = await r.subscription.updateStatus(snapshot, engine: self, revision: r.statusRevision) else { return } + for eventKind in [kind, transition].compactMap({ $0 }) { + for listener in telemetry.values { + listener.yield(.init(channel: channel, kind: eventKind, status: value)) + } + } + } + + private func validate(_ limits: ChannelReceiveLimits) throws { + guard limits.windowBytes > 0, limits.maximumFrameBytes >= 1024, + limits.windowBytes <= 9_007_199_254_740_991, + limits.maximumFrameBytes <= 9_007_199_254_740_991 else { + throw TransportError.protocol("invalid window limits") + } + } + private func validateTCPData(_ message: AHPNotificationMessage) throws { + let value = try message.decode(TcpDataParams.self).data + guard !value.isEmpty, let data = Data(base64Encoded: value), data.base64EncodedString() == value else { + throw TransportError.protocol("invalid padded TCP base64") + } + } + private func validateAccepted(_ limits: ChannelFlowControl, offer: ChannelReceiveLimits) throws { + try validate(limits.clientReceive); try validate(limits.hostReceive) + guard limits.clientReceive.windowBytes <= offer.windowBytes, + limits.clientReceive.maximumFrameBytes <= offer.maximumFrameBytes else { + throw TransportError.protocol("accepted limits exceed offer") + } + } + private func adding(_ a: Int, _ b: Int) throws -> Int { + let (value, overflow) = a.addingReportingOverflow(b) + guard !overflow, value <= 9_007_199_254_740_991 else { + throw TransportError.protocol("window counter overflow") + } + return value + } +} diff --git a/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/WindowingTransport.swift b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/WindowingTransport.swift new file mode 100644 index 000000000..801371d3c --- /dev/null +++ b/clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/WindowingTransport.swift @@ -0,0 +1,58 @@ +import Foundation + +internal typealias AHPBeforeWrite = @Sendable () async throws -> Void + +public struct AHPWireMessage: Sendable { + public let message: TransportMessage + public let originalJSONText: String? + public init(message: TransportMessage, originalJSONText: String? = nil) { + self.message = message + self.originalJSONText = originalJSONText + } +} + +public struct AHPTransportCapabilities: Sendable { + public let preservesOriginalJSONText: Bool + public let maximumReceiveMessageUTF8Bytes: Int? + public init(preservesOriginalJSONText: Bool = false, maximumReceiveMessageUTF8Bytes: Int? = nil) { + self.preservesOriginalJSONText = preservesOriginalJSONText + self.maximumReceiveMessageUTF8Bytes = maximumReceiveMessageUTF8Bytes + } + public var maximumFrameBytes: Int { + maximumReceiveMessageUTF8Bytes.map { 2 * ($0 / 3) } ?? 1_048_576 + } +} + +/// Optional refinement; legacy transports and TransportMessage remain unchanged. +public protocol AHPWindowingTransport: AHPTransport { + var windowingCapabilities: AHPTransportCapabilities { get async } + func recvWireMessage() async throws -> AHPWireMessage? +} + +extension AHPWindowingTransport { + public var windowingCapabilities: AHPTransportCapabilities { get async { .init() } } + public func recvWireMessage() async throws -> AHPWireMessage? { + guard let message = try await recv() else { return nil } + switch message { + case .text(let text): return .init(message: message, originalJSONText: text) + case .binary(let data): + guard let text = String(data: data, encoding: .utf8) else { + throw TransportError.protocol("binary JSON is not UTF-8") + } + return .init(message: message, originalJSONText: text) + case .parsed: return .init(message: message) + } + } +} + +extension InMemoryTransport: AHPWindowingTransport { + public var windowingCapabilities: AHPTransportCapabilities { + get async { .init(preservesOriginalJSONText: true) } + } +} + +extension NWConnectionWebSocketTransport: AHPWindowingTransport { + public var windowingCapabilities: AHPTransportCapabilities { + get async { .init(preservesOriginalJSONText: true) } + } +} diff --git a/clients/swift/AgentHostProtocol/Tests/AgentHostProtocolClientTests/NWConnectionWebSocketTransportTests.swift b/clients/swift/AgentHostProtocol/Tests/AgentHostProtocolClientTests/NWConnectionWebSocketTransportTests.swift index 74560c0ba..b4580e9f2 100644 --- a/clients/swift/AgentHostProtocol/Tests/AgentHostProtocolClientTests/NWConnectionWebSocketTransportTests.swift +++ b/clients/swift/AgentHostProtocol/Tests/AgentHostProtocolClientTests/NWConnectionWebSocketTransportTests.swift @@ -5,6 +5,22 @@ import XCTest @testable import AgentHostProtocolClient final class NWConnectionWebSocketTransportTests: XCTestCase { + func testNativeTransportDrainsBufferedWebSocketContinuations() async throws { + let server = try await LoopbackWebSocketServer.start() + let transport = NWConnectionWebSocketTransport(url: await server.url(path: "/fragments")) + let watchdog = Task { + try? await Task.sleep(for: .seconds(2)) + if !Task.isCancelled { try? await transport.close() } + } + try await transport.send(.text("connect")) + try await server.sendFragmentedText("one", "two") + let response = try await withTimeout { try await transport.recv() } + XCTAssertEqual(stringPayload(of: response), "onetwo") + watchdog.cancel() + try await transport.close() + await server.stop() + } + func testNativeTransportPerformsHandshakeAndRoundTripsText() async throws { let server = try await LoopbackWebSocketServer.start() let transport = NWConnectionWebSocketTransport( @@ -133,6 +149,14 @@ private actor LoopbackWebSocketServer { try await sendRaw(makeServerFrame(opcode: 0x1, payload: Data(text.utf8)), on: connection) } + func sendFragmentedText(_ first: String, _ second: String) async throws { + guard let connection else { throw TransportError.closed } + var frames = makeServerFrame(opcode: 0x1, payload: Data(first.utf8)) + frames[0] = 0x01 + frames.append(makeServerFrame(opcode: 0x0, payload: Data(second.utf8))) + try await sendRaw(frames, on: connection) + } + func stop() { connection?.cancel() listener.cancel() diff --git a/clients/swift/AgentHostProtocol/Tests/AgentHostProtocolClientTests/WindowedDeliveryCorpusTests.swift b/clients/swift/AgentHostProtocol/Tests/AgentHostProtocolClientTests/WindowedDeliveryCorpusTests.swift new file mode 100644 index 000000000..d60683b92 --- /dev/null +++ b/clients/swift/AgentHostProtocol/Tests/AgentHostProtocolClientTests/WindowedDeliveryCorpusTests.swift @@ -0,0 +1,302 @@ +import Foundation +import XCTest +import AgentHostProtocol +@testable import AgentHostProtocolClient + +final class WindowedDeliveryCorpusTests: XCTestCase { + func testRecorderChecksAvailableMessagesBeforeItsDeadline() async throws { + let recorder = CorpusRecorder(timeout: .zero) + await recorder.append("{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"subscribe\",\"params\":{\"channel\":\"ahp-terminal:/ready\"}}") + let message = try await recorder.take(["method": "subscribe"]).1 + let id = try await recorder.requestId("subscribe") + XCTAssertEqual(message["id"] as? Int, 1) + XCTAssertEqual(id, 1) + } + + private static let corpus = URL(fileURLWithPath: #filePath) + .deletingLastPathComponent().deletingLastPathComponent() + .deletingLastPathComponent().deletingLastPathComponent() + .deletingLastPathComponent().deletingLastPathComponent() + .appendingPathComponent("types/test-cases/windowed-delivery") + + func testEveryRawSize() throws { + let object = try JSONSerialization.jsonObject( + with: Data(contentsOf: Self.corpus.appendingPathComponent("sizes.json"))) as! [String: Any] + let cases = object["cases"] as! [[String: Any]] + XCTAssertFalse(cases.isEmpty) + for item in cases { + XCTAssertEqual(ahpUTF16Bytes(item["text"] as! String), item["utf16Bytes"] as? Int, + item["name"] as! String) + } + } + + func testEveryScenario() async throws { + let directory = Self.corpus.appendingPathComponent("scenarios") + let files = try FileManager.default.contentsOfDirectory(at: directory, includingPropertiesForKeys: nil) + .filter { $0.pathExtension == "json" }.sorted { $0.lastPathComponent < $1.lastPathComponent } + XCTAssertEqual(files.count, 24) + for file in files { + let scenario = try JSONSerialization.jsonObject(with: Data(contentsOf: file)) as! [String: Any] + do { try await run(scenario) } + catch { XCTFail("\(file.lastPathComponent): \(error)") } + } + } + + private func run(_ scenario: [String: Any]) async throws { + let config = scenario["windowing"] as? [String: Int] ?? [:] + let (transport, host) = InMemoryTransport.pair() + let client = AHPClient(transport: transport, config: .init(windowing: .init( + windowBytes: config["windowBytes"] ?? 262_144, + maximumFrameBytes: config["maximumFrameBytes"] ?? 16_384, + maxBufferedSendBytes: config["maxBufferedSendBytes"] ?? 1_048_576))) + let recorder = CorpusRecorder() + let pump = Task { + while let message = try? await host.recv() { + if case .text(let text) = message { await recorder.append(text) } + } + await recorder.finish() + } + try await client.connect() + var subscriptions: [String: Task] = [:] + var deliveries: [String: [AHPDelivery]] = [:] + var sends: [Task] = [] + do { + for step in scenario["steps"] as! [[String: Any]] { + let op = step["op"] as! String + let channel = step["channel"] as? String ?? "" + switch op { + case "subscribe": + let window = step["window"] as? [String: Int] + let options: AHPSubscriptionOptions = window.map { + .windowed(windowBytes: $0["windowBytes"], maximumFrameBytes: $0["maximumFrameBytes"], + receive: step["receive"] as? String == "manual" ? .manual : .auto, + coalescing: step["coalescing"] as? Bool ?? true) + } ?? .init() + subscriptions[channel] = Task { try await client.subscribe(channel, options: options) } + case "respond": + let id = try await recorder.requestId(step["to"] as! String) + let response: [String: Any] = ["jsonrpc": "2.0", "id": id, "result": step["result"]!] + try await host.send(.binary(JSONSerialization.data(withJSONObject: response))) + case "hostSends": + try await host.send(.text(step["text"] as! String)) + case "hostCreditsAll": + let bytes = await recorder.sentBytes(channel) + try await host.send(makeNotificationWire(method: "channel/credit", + params: ChannelCreditParams(channel: channel, consumedBytes: bytes))) + case "expectSent": + let match = step["match"] as! [String: Any] + let object: [String: Any] + if step["reassemble"] as? Bool == true { + var joined = "" + let targetChannel = (match["params"] as! [String: Any])["channel"] as! String + let sub = try await subscriptions[targetChannel]!.value + let cap = await sub.status().limits!.hostReceive.maximumFrameBytes + while true { + let (text, message) = try await recorder.take(["method": "channel/frame"]) + guard ahpUTF16Bytes(text) <= cap else { throw CorpusFailure("oversized outbound frame") } + let p = message["params"] as! [String: Any] + joined += p["data"] as! String + if p["final"] as? Bool == true { break } + } + object = try JSONSerialization.jsonObject(with: Data(joined.utf8)) as! [String: Any] + guard matches(object, match) else { throw CorpusFailure("joined data does not match") } + } else { object = try await recorder.take(match).1 } + for path in step["absent"] as? [String] ?? [] { + guard lookup(object, path) == nil else { throw CorpusFailure("unexpected \(path)") } + } + case "expectNotSent": + let duration = Duration.milliseconds(step["withinMs"] as! Int) + let deadline = ContinuousClock.now + duration + while ContinuousClock.now < deadline { + guard await !recorder.contains(step["match"]!) else { + throw CorpusFailure("unexpected sent message \(step["match"]!)") + } + try await Task.sleep(for: .milliseconds(2)) + } + case "next": + let sub = try await subscriptions[channel]!.value + guard let item = try await corpusTimeout({ try await sub.next() }) else { + throw CorpusFailure("subscription ended") + } + let p = try JSONSerialization.jsonObject(with: item.message.paramsJSON) + let object: [String: Any] = ["method": item.message.method, "params": p] + guard matches(object, step["expect"]!) else { throw CorpusFailure("delivery mismatch: \(object)") } + if let expected = step["expectBytes"] as? Int, expected != item.receivedBytes { + throw CorpusFailure("delivery size \(item.receivedBytes), expected \(expected)") + } + deliveries[channel, default: []].append(item) + case "advance": + let sub = try await subscriptions[channel]!.value + await sub.advance() + case "release": + await deliveries[channel]![step["index"] as! Int].release() + case "unsubscribe": + let sub = try await subscriptions[channel]!.value + try await sub.close() + case "send": + let sub = try await subscriptions[channel]!.value + let message = try AHPNotificationMessage( + method: step["method"] as! String, + paramsJSON: JSONSerialization.data(withJSONObject: step["params"]!)) + sends.append(Task { try await sub.send(message) }) + // Ensure admission order without waiting for blocked transport completion. + try await Task.sleep(for: .milliseconds(5)) + case "expectStatus": + let sub = try await subscriptions[channel]!.value + try await corpusTimeout { + while true { + let value = try JSONSerialization.jsonObject(with: JSONEncoder().encode(await sub.status())) + if matches(value, step["match"]!) { return } + try await Task.sleep(for: .milliseconds(1)) + } + } + case "expectClosed": + try await corpusTimeout { + while await client.connectionState != .disconnected { + try await Task.sleep(for: .milliseconds(1)) + } + } + default: throw CorpusFailure("unsupported corpus op \(op)") + } + } + } catch { + await client.shutdown(); try? await host.close(); pump.cancel() + for task in subscriptions.values { task.cancel() } + for task in sends { task.cancel(); _ = try? await task.value } + throw error + } + await client.shutdown(); try? await host.close(); pump.cancel() + for task in subscriptions.values { task.cancel() } + for task in sends { task.cancel(); _ = try? await task.value } + } +} + +actor CorpusRecorder { + private struct ChangeWaiter { + let continuation: CheckedContinuation + let timer: Task + } + private let timeout: Duration + private var changeWaiters: [UUID: ChangeWaiter] = [:] + private var finished = false + var messages: [(String, [String: Any])] = [] + var all: [(String, [String: Any])] = [] + var unanswered: [(Int, String)] = [] + + init(timeout: Duration = .seconds(1)) { self.timeout = timeout } + + func append(_ text: String) { + let object = try! JSONSerialization.jsonObject(with: Data(text.utf8)) as! [String: Any] + messages.append((text, object)); all.append((text, object)) + if let id = object["id"] as? Int, let method = object["method"] as? String { + unanswered.append((id, method)) + } + wakeAll() + } + + func finish() { + finished = true + wakeAll() + } + + private func wakeAll() { + let waiters = changeWaiters.values + changeWaiters.removeAll() + for waiter in waiters { + waiter.timer.cancel() + waiter.continuation.resume() + } + } + func requestId(_ method: String) async throws -> Int { + let deadline = ContinuousClock.now + timeout + while true { + try Task.checkCancellation() + if let i = unanswered.firstIndex(where: { $0.1 == method }) { + return unanswered.remove(at: i).0 + } + guard !finished else { throw CorpusFailure("transport closed waiting for unanswered \(method)") } + guard ContinuousClock.now < deadline else { throw CorpusFailure("no unanswered \(method)") } + try await waitForChange(until: deadline) + } + } + func take(_ match: [String: Any]) async throws -> (String, [String: Any]) { + let deadline = ContinuousClock.now + timeout + while true { + try Task.checkCancellation() + if !messages.isEmpty { + let item = messages.removeFirst() + if matches(item.1, match) { return item } + if item.1["method"] as? String == "channel/credit" { continue } + throw CorpusFailure("unexpected sent message \(item.1), expected \(match)") + } + guard !finished else { throw CorpusFailure("transport closed waiting for \(match)") } + guard ContinuousClock.now < deadline else { throw CorpusFailure("timed out waiting for \(match)") } + try await waitForChange(until: deadline) + } + } + + private func waitForChange(until deadline: ContinuousClock.Instant) async throws { + let id = UUID() + try await withTaskCancellationHandler { + try await withCheckedThrowingContinuation { continuation in + let timer = Task { [weak self] in + try? await Task.sleep(until: deadline, clock: .continuous) + if !Task.isCancelled { await self?.wake(id) } + } + changeWaiters[id] = ChangeWaiter(continuation: continuation, timer: timer) + } + } onCancel: { + Task { await self.cancelWait(id) } + } + } + + private func wake(_ id: UUID) { + guard let waiter = changeWaiters.removeValue(forKey: id) else { return } + waiter.timer.cancel() + waiter.continuation.resume() + } + + private func cancelWait(_ id: UUID) { + guard let waiter = changeWaiters.removeValue(forKey: id) else { return } + waiter.timer.cancel() + waiter.continuation.resume(throwing: CancellationError()) + } + func contains(_ match: Any) -> Bool { messages.contains { matches($0.1, match) } } + func allMethods() -> [String] { all.compactMap { $0.1["method"] as? String } } + func sentBytes(_ channel: String) -> Int { + all.reduce(0) { sum, item in + guard item.1["id"] == nil, let method = item.1["method"] as? String, + countsAgainstWindow(method), + (item.1["params"] as? [String: Any])?["channel"] as? String == channel else { return sum } + return sum + ahpUTF16Bytes(item.0) + } + } +} + +private func matches(_ value: Any, _ expected: Any) -> Bool { + if let e = expected as? [String: Any], let v = value as? [String: Any] { + return e.allSatisfy { key, value in v[key].map { matches($0, value) } ?? false } + } + return (value as? NSObject)?.isEqual(expected) ?? false +} +private func lookup(_ object: [String: Any], _ path: String) -> Any? { + var current: Any = object + for key in path.split(separator: ".") { + guard let next = (current as? [String: Any])?[String(key)] else { return nil } + current = next + } + return current +} +private struct CorpusFailure: Error, CustomStringConvertible { + let description: String + init(_ description: String) { self.description = description } +} +func corpusTimeout(timeout: Duration = .seconds(1), + _ body: @escaping @Sendable () async throws -> T) async throws -> T { + try await withThrowingTaskGroup(of: T.self) { group in + group.addTask { try await body() } + group.addTask { try await Task.sleep(for: timeout); throw CorpusFailure("step timed out") } + defer { group.cancelAll() } + return try await group.next()! + } +} diff --git a/clients/swift/AgentHostProtocol/Tests/AgentHostProtocolClientTests/WindowingLifecycleTests.swift b/clients/swift/AgentHostProtocol/Tests/AgentHostProtocolClientTests/WindowingLifecycleTests.swift new file mode 100644 index 000000000..c72439806 --- /dev/null +++ b/clients/swift/AgentHostProtocol/Tests/AgentHostProtocolClientTests/WindowingLifecycleTests.swift @@ -0,0 +1,394 @@ +import Foundation +import XCTest +import AgentHostProtocol +@testable import AgentHostProtocolClient + +final class WindowingLifecycleTests: XCTestCase { + func testRPCFailureRemovesUnacceptedChannelAndDoesNotBlockLaterSends() async throws { + let h = try await ReviewHarness.open() + let uri = "ahp-terminal:/denied" + let task = Task { try await h.client.subscribe(uri, options: .windowed()) } + let request = try await h.request("subscribe") + try await h.peer.send(.encoded(.errorResponse(id: request["id"] as! Int, + error: JsonRpcError(code: -32602, message: "denied")))) + do { _ = try await task.value; XCTFail("expected RPC failure") } catch AHPClientError.rpc {} + let missing = await h.client.subscription(for: uri) + XCTAssertNil(missing) + let count = await h.client._windowOrderCount() + XCTAssertEqual(count, 0) + try await h.client.send(method: "custom/data", params: ReviewData(channel: uri, data: "ordinary")) + _ = try await h.recorder.take(["method": "custom/data"]) + let state = await h.client.connectionState + XCTAssertEqual(state, .connected) + await h.close() + } + + func testDroppingChannelSuppressesAlreadyQueuedCredit() async throws { + let h = try await ReviewHarness.open() + let sub = try await h.subscribe("ahp-terminal:/credit", receive: .manual) + let item = try await sub.next() + await h.gate.block() + let blocker = Task { + try await h.client.send(method: "custom/data", params: ReviewData(channel: RootResourceURI, data: "block")) + } + try await reviewEventually { await h.client._callerWriteCount() == 1 } + await item?.release() + try await reviewEventually { await h.client._windowWriteCount() >= 2 } + let close = Task { try await sub.close() } + try await reviewEventually { await h.client.subscription(for: sub.channel) == nil } + await h.gate.unblock() + try await blocker.value + try await close.value + _ = try await h.recorder.take(["method": "custom/data"]) + _ = try await h.recorder.take(["method": "unsubscribe"]) + let methods = await h.recorder.allMethods() + XCTAssertFalse(methods.contains("channel/credit")) + await h.close() + } + + func testCallerCapCannotRejectCreditsFramesUnsubscribeOrReset() async throws { + let h = try await ReviewHarness.open() + let data = try await h.subscribe("ahp-terminal:/data", receive: .manual) + let closing = try await h.subscribe("ahp-terminal:/closing") + let tcp = try await h.subscribe("ahp-tcp:/tcp") + let snapshot = try await data.next() + await h.gate.block() + var callers: [Task] = [] + for i in 0..<256 { + callers.append(Task { + try await h.client.send(method: "custom/data", params: ReviewData(channel: RootResourceURI, data: "\(i)")) + }) + } + try await reviewEventually { await h.client._callerWriteCount() == 256 } + do { + try await h.client.send(method: "custom/data", params: ReviewData(channel: RootResourceURI, data: "overflow")) + XCTFail("caller admission should be bounded") + } catch AHPWindowingError.backpressure {} + await snapshot?.release() + let frame = Task { try await data.send(method: "custom/window", params: ReviewData(channel: data.channel, data: "x")) } + let unsubscribe = Task { try await closing.close() } + let reset = Task { try await tcp.send(method: "tcp/reset", params: TcpResetParams(channel: tcp.channel)) } + try await reviewEventually { await h.client._windowWriteCount() >= 260 } + let state = await h.client.connectionState + XCTAssertEqual(state, .connected) + await h.gate.unblock() + for caller in callers { try await caller.value } + try await frame.value + try await unsubscribe.value + try await reset.value + try await reviewEventually { + let methods = await h.recorder.allMethods() + return methods.contains("channel/credit") && methods.contains("custom/window") && + methods.contains("tcp/reset") && methods.filter { $0 == "unsubscribe" }.count >= 2 + } + await h.close() + } + + func testCancelledAndTimedOutSubscribeCleanUpLateAcceptanceWithoutDisconnect() async throws { + for cancel in [true, false] { + let h = try await ReviewHarness.open(timeout: cancel ? .seconds(30) : .milliseconds(20)) + let uri = "ahp-terminal:/abandoned" + let task = Task { try await h.client.subscribe(uri, options: .windowed()) } + let request = try await h.request("subscribe") + if cancel { task.cancel() } + do { _ = try await task.value; XCTFail("abandoned subscribe succeeded") } catch {} + let owner = await h.client.subscription(for: uri) + XCTAssertNil(owner) + let retained = await h.client._requestContextCount() + XCTAssertEqual(retained, 1, "late negotiation metadata must survive local abandonment") + let state = await h.client.connectionState + XCTAssertEqual(state, .connected) + try await h.accept(request) + _ = try await h.recorder.take(["method": "unsubscribe", "params": ["channel": uri]]) + try await reviewEventually { await h.client._requestContextCount() == 0 } + let after = await h.client.connectionState + XCTAssertEqual(after, .connected) + await h.close() + } + } + + func testCancelledResubscribeStillInstallsItsResponse() async throws { + let h = try await ReviewHarness.open() + let sub = try await h.subscribe("ahp-terminal:/existing") + let update = Task { try await sub.update(options: .init()) } + let request = try await h.request("subscribe") + update.cancel() + do { _ = try await update.value; XCTFail("expected cancellation") } catch is CancellationError {} + try await h.respond(request, result: SubscribeResult(snapshot: ReviewHarness.snapshot(sub.channel))) + try await reviewEventually { await sub.status().mode == "ordinary" } + let current = await h.client.subscription(for: sub.channel) + XCTAssertTrue(current === sub) + let state = await h.client.connectionState + XCTAssertEqual(state, .connected) + await h.close() + } + + func testLateReconnectResponseAfterTimeoutIsInstalled() async throws { + let h = try await ReviewHarness.open(timeout: .milliseconds(20)) + let uri = "ahp-terminal:/resume" + let limits = ReviewHarness.limits + let reconnect = Task { + try await h.client.reconnect(clientId: "client", lastSeenServerSeq: 1, subscriptions: [uri], + windows: .init(items: [.init(channel: uri, flowControl: .init(receive: limits.clientReceive), + lastAppliedServerSeq: 1)])) + } + let request = try await h.request("reconnect") + do { _ = try await reconnect.value; XCTFail("expected timeout") } catch AHPClientError.requestTimeout {} + try await h.respond(request, result: ReconnectReplayResult(type: .replay, actions: [], missing: [], + windows: .init(items: [.init(channel: uri, flowControl: limits, recovery: .snapshot)]))) + try await h.sendSnapshot(uri) + try await reviewEventually { await h.client._requestContextCount() == 0 } + let optional = await h.client.subscription(for: uri) + let sub = try XCTUnwrap(optional) + let delivery = try await sub.next() + XCTAssertEqual(delivery?.message.method, "channel/snapshot") + let state = await h.client.connectionState + XCTAssertEqual(state, .connected) + await h.close() + } + + func testUnsubscribeRacingSubscribeDoesNotResurrectAndOrderStaysBounded() async throws { + let h = try await ReviewHarness.open() + do { + let uri = "ahp-terminal:/race" + let first = Task { try await h.client.subscribe(uri, options: .windowed()) } + let request = try await h.request("subscribe") + try await h.client.unsubscribe(uri) + _ = try await h.recorder.take(["method": "unsubscribe"]) + try await h.accept(request) + _ = try await first.value + _ = try await h.recorder.take(["method": "unsubscribe"]) + let old = await h.client.subscription(for: uri) + XCTAssertNil(old) + for _ in 0..<25 { + let sub = try await h.subscribe(uri) + let snapshot = try await corpusTimeout(timeout: ReviewHarness.waitTimeout) { try await sub.next() } + XCTAssertEqual(snapshot?.message.method, "channel/snapshot") + let count = await h.client._windowOrderCount() + XCTAssertEqual(count, 1) + try await sub.close() + _ = try await h.recorder.take(["method": "unsubscribe"]) + let empty = await h.client._windowOrderCount() + XCTAssertEqual(empty, 0) + } + let state = await h.client.connectionState + XCTAssertEqual(state, .connected) + } catch { + await h.close() + throw error + } + await h.close() + } + + func testPendingResubscribeIgnoresPreviousSubscriptionsSnapshot() async throws { + let h = try await ReviewHarness.open() + do { + let uri = "ahp-terminal:/late-snapshot" + let first = Task { try await h.client.subscribe(uri, options: .windowed()) } + let firstRequest = try await h.request("subscribe") + try await h.accept(firstRequest) + let previous = try await first.value + try await previous.close() + _ = try await h.recorder.take(["method": "unsubscribe"]) + + let next = Task { try await h.client.subscribe(uri, options: .windowed(receive: .manual)) } + let nextRequest = try await h.request("subscribe") + // The old snapshot is queued before the next result, regardless of + // which task the client executor runs first. + try await h.sendSnapshot(uri) + try await h.accept(nextRequest) + let fresh = Snapshot(resource: uri, state: .empty(EmptyState()), fromSeq: 2) + try await h.peer.send(makeNotificationWire(method: "channel/snapshot", + params: ChannelSnapshotParams(channel: uri, snapshot: fresh))) + let subscription = try await next.value + let delivery = try await corpusTimeout(timeout: ReviewHarness.waitTimeout) { + try await subscription.next() + } + if case .snapshot(let snapshot) = delivery?.payload { + XCTAssertEqual(snapshot.fromSeq, 2) + } else { + XCTFail("the new generation's snapshot must be delivered first") + } + let status = await subscription.status() + XCTAssertEqual(status.bufferedReceiveBytes, delivery?.receivedBytes) + let state = await h.client.connectionState + XCTAssertEqual(state, .connected) + } catch { + await h.close() + throw error + } + await h.close() + } + + func testAlreadyCancelledManagedSubscribeLeavesNoReplayRegistration() async throws { + let host = WindowHost() + let multi = MultiHostClient() + _ = try await multi.add(HostConfig(id: "host", label: "Host", transportFactory: { _ in await host.open() }) + .withSessionSummaryRefreshOnConnect(false)) + try await reviewEventually { await multi.host("host")?.state.isConnected == true } + let subscribe = Task { + withUnsafeCurrentTask { $0?.cancel() } + return try await multi.subscribe(host: "host", uri: "ahp-terminal:/cancelled", options: .windowed()) + } + do { _ = try await subscribe.value; XCTFail("expected cancellation") } catch is CancellationError {} + let state = await multi.host("host") + XCTAssertFalse(state!.subscriptions.contains("ahp-terminal:/cancelled")) + XCTAssertTrue(state!.state.isConnected) + await multi.shutdown() + } + + func testFailedFallbackResubscribeOnlyDropsThatChannel() async throws { + let host = WindowHost() + let multi = MultiHostClient() + _ = try await multi.add(HostConfig(id: "host", label: "Host", transportFactory: { _ in await host.open() }) + .withSessionSummaryRefreshOnConnect(false).withReconnectPolicy(.immediateForever)) + try await reviewEventually { await multi.host("host")?.state.isConnected == true } + let bad = try await multi.subscribe(host: "host", uri: "ahp-terminal:/bad", options: .windowed()) + let good = try await multi.subscribe(host: "host", uri: "ahp-terminal:/good", options: .windowed()) + _ = try await bad.next(); _ = try await good.next() + let generation = await multi.host("host")!.generation + await host.setFallback(failing: bad.channel) + await host.disconnect() + try await reviewEventually { + guard let state = await multi.host("host") else { return false } + return state.generation > generation && state.state.isConnected + } + do { _ = try await bad.next(); XCTFail("failed resubscribe must close its owner") } catch {} + let state = await multi.host("host") + XCTAssertFalse(state!.subscriptions.contains(bad.channel)) + XCTAssertTrue(state!.subscriptions.contains(good.channel)) + let snapshot = try await good.next() + XCTAssertEqual(snapshot?.message.method, "channel/snapshot") + await multi.shutdown() + } + + func testShutdownDoesNotWaitForFallbackResubscribeTimeout() async throws { + let host = WindowHost() + let multi = MultiHostClient() + _ = try await multi.add(HostConfig(id: "host", label: "Host", transportFactory: { _ in await host.open() }) + .withSessionSummaryRefreshOnConnect(false).withReconnectPolicy(.immediateForever)) + try await reviewEventually { await multi.host("host")?.state.isConnected == true } + _ = try await multi.subscribe(host: "host", uri: "ahp-terminal:/held", options: .windowed()) + await host.setFallback(hold: true) + await host.disconnect() + try await reviewEventually { await host.subscribeCount() >= 2 } + let start = ContinuousClock.now + await multi.shutdown() + XCTAssertLessThan(start.duration(to: .now), .seconds(2)) + } + + func testOrdinaryFireAndForgetWritesDrainAndEventStreamsFinish() async throws { + let h = try await ReviewHarness.open() + let events = await h.client.events + let states = await h.client.stateChanges + try await h.client.notify(method: "custom/data", params: ReviewData(channel: RootResourceURI, data: "notify")) + _ = try await h.client.dispatch(.terminalInput(.init(type: .terminalInput, data: "dispatch")), + channel: "ahp-terminal:/ordinary") + try await h.client.unsubscribe("ahp-terminal:/ordinary") + await h.client.shutdown() + await h.reader.value + let methods = await h.recorder.allMethods() + XCTAssertEqual(methods, ["custom/data", "dispatchAction", "unsubscribe"]) + for await _ in events {} + for await _ in states {} + } +} + +private struct ReviewData: Codable, Sendable { let channel: String; let data: String } + +private struct ReviewHarness: Sendable { + static let waitTimeout: Duration = .seconds(30) + let client: AHPClient + let peer: InMemoryTransport + let gate: ReviewGate + let recorder: CorpusRecorder + let reader: Task + static let limits = ChannelFlowControl( + clientReceive: .init(windowBytes: 65_536, maximumFrameBytes: 1024), + hostReceive: .init(windowBytes: 65_536, maximumFrameBytes: 1024)) + + static func open(timeout: Duration = .seconds(30)) async throws -> Self { + let (transport, peer) = InMemoryTransport.pair() + let gate = ReviewGate() + let client = AHPClient(transport: ReviewTransport(base: transport, gate: gate), + config: .init(requestTimeout: timeout)) + let recorder = CorpusRecorder(timeout: waitTimeout) + let reader = Task { + while let message = try? await peer.recv() { + if case .text(let text) = message { await recorder.append(text) } + } + await recorder.finish() + } + try await client.connect() + return .init(client: client, peer: peer, gate: gate, recorder: recorder, reader: reader) + } + static func snapshot(_ channel: String) -> Snapshot { + .init(resource: channel, state: .empty(EmptyState()), fromSeq: 1) + } + func request(_ method: String) async throws -> [String: Any] { + let request = try await recorder.take(["method": method]).1 + _ = try await recorder.requestId(method) + return request + } + func respond(_ request: [String: Any], result: R) async throws { + try await peer.send(makeResponseWire(id: request["id"] as! Int, result: result)) + } + func accept(_ request: [String: Any]) async throws { + try await respond(request, result: SubscribeResult(flowControl: Self.limits)) + } + func sendSnapshot(_ uri: String) async throws { + try await peer.send(makeNotificationWire(method: "channel/snapshot", + params: ChannelSnapshotParams(channel: uri, snapshot: Self.snapshot(uri)))) + } + func subscribe(_ uri: String, receive: AHPReceiveMode = .auto) async throws -> AHPSubscription { + let task = Task { try await client.subscribe(uri, options: .windowed(receive: receive)) } + let request = try await request("subscribe") + try await accept(request) + try await sendSnapshot(uri) + return try await task.value + } + func close() async { + await gate.unblock() + await client.shutdown() + await reader.value + } +} + +private actor ReviewGate { + var blocked = false + var waiters: [CheckedContinuation] = [] + func block() { blocked = true } + func wait() async { + if blocked { await withCheckedContinuation { waiters.append($0) } } + } + func unblock() { + blocked = false + let pending = waiters + waiters.removeAll() + for waiter in pending { waiter.resume() } + } +} + +private final class ReviewTransport: AHPWindowingTransport, Sendable { + let base: InMemoryTransport + let gate: ReviewGate + init(base: InMemoryTransport, gate: ReviewGate) { self.base = base; self.gate = gate } + var windowingCapabilities: AHPTransportCapabilities { + get async { .init(preservesOriginalJSONText: true) } + } + func recv() async throws -> TransportMessage? { try await base.recv() } + func send(_ message: TransportMessage) async throws { + await gate.wait() + try await base.send(message) + } + func close() async throws { await gate.unblock(); try await base.close() } +} + +private func reviewEventually(_ predicate: @escaping @Sendable () async -> Bool) async throws { + let deadline = ContinuousClock.now + .seconds(2) + while ContinuousClock.now < deadline { + if await predicate() { return } + try await Task.sleep(for: .milliseconds(2)) + } + throw TransportError.protocol("review regression timed out") +} diff --git a/clients/swift/AgentHostProtocol/Tests/AgentHostProtocolClientTests/WindowingTests.swift b/clients/swift/AgentHostProtocol/Tests/AgentHostProtocolClientTests/WindowingTests.swift new file mode 100644 index 000000000..67cf363cc --- /dev/null +++ b/clients/swift/AgentHostProtocol/Tests/AgentHostProtocolClientTests/WindowingTests.swift @@ -0,0 +1,461 @@ +import Foundation +import XCTest +import AgentHostProtocol +@testable import AgentHostProtocolClient + +final class WindowingTests: XCTestCase { + func testModeChangeWaitsForStartedFragments() async throws { + let host = WindowHost() + await host.setHostLimits(window: 512, frame: 1024) + let client = AHPClient(transport: await host.open()) + try await client.connect() + let sub = try await client.subscribe("ahp-terminal:/t1", options: .windowed()) + let send = Task { try await sub.send(method: "custom/data", + params: WindowData(channel: sub.channel, data: String(repeating: "x", count: 2000))) } + try await windowEventually { await sub.status().outstandingSendBytes >= 512 } + let update = Task { try await sub.update(options: .init()) } + try await Task.sleep(for: .milliseconds(10)) + let before = await host.subscribeCount() + XCTAssertEqual(before, 1, "mode-changing RPC must not abandon an unfinished message") + let credits = Task { + while !Task.isCancelled { + try? await host.creditAll(sub.channel) + try? await Task.sleep(for: .milliseconds(2)) + } + } + _ = try await update.value + try await send.value + credits.cancel() + let status = await sub.status() + XCTAssertEqual(status.mode, "ordinary") + XCTAssertEqual(status.outstandingSendBytes, 0) + await client.shutdown() + } + + func testCoalescingDoesNotEquateBooleanAndNumberMetadata() async throws { + let host = WindowHost() + let client = AHPClient(transport: await host.open()) + try await client.connect() + let sub = try await client.subscribe("ahp-chat:/c1", options: .windowed(receive: .manual)) + _ = try await sub.next() + for (seq, value) in [(2, "1"), (3, "true")] { + try await host.sendRaw("{\"jsonrpc\":\"2.0\",\"method\":\"action\",\"params\":{\"channel\":\"ahp-chat:/c1\",\"serverSeq\":\(seq),\"action\":{\"type\":\"chat/delta\",\"turnId\":\"t1\",\"partId\":\"p1\",\"content\":\"a\",\"_meta\":{\"value\":\(value)}}}}") + } + try await windowEventually { await sub.lastAppliedServerSeq == 3 } + let a = try await windowNext(sub) + let b = try await windowNext(sub) + XCTAssertEqual(try a.message.decode(ActionEnvelope.self).serverSeq, 2) + XCTAssertEqual(try b.message.decode(ActionEnvelope.self).serverSeq, 3) + await client.shutdown() + } + + func testRawHandshakeAndParsedEnvelopeUseOriginalText() async throws { + let host = WindowHost() + let client = AHPClient(transport: ParsedEnvelopeTransport(await host.open())) + try await client.connect() + let params = SubscribeParams(channel: "ahp-terminal:/raw", + flowControl: .init(receive: .init(windowBytes: 4096, maximumFrameBytes: 1024))) + _ = try await client.requestRaw(method: "subscribe", paramsData: JSONEncoder().encode(params)) + let owner = await client.subscription(for: params.channel) + let sub = try XCTUnwrap(owner) + var iterator = sub.events.makeAsyncIterator() + _ = try await iterator.next() + let raw = "{ \n \"jsonrpc\" : \"2.0\", \"method\" : \"action\", \"params\" : {\"channel\":\"ahp-terminal:/raw\",\"serverSeq\":2,\"action\":{\"type\":\"terminal/data\",\"data\":\"漢é😀\"}} }" + try await host.sendRaw(raw) + let item = try await iterator.next() + XCTAssertEqual(item?.receivedBytes, ahpUTF16Bytes(raw)) + if case .event(.action(let envelope)) = item?.payload { + XCTAssertEqual(envelope.serverSeq, 2) + } else { XCTFail("original text should bypass lossy parsed transport payload") } + var another = sub.events.makeAsyncIterator() + do { _ = try await another.next(); XCTFail("multiple credit owners") } + catch AHPWindowingError.consumerAlreadyWaiting {} + await client.shutdown() + } + + func testConcurrentSubscribeOffersRemainRequestScoped() async throws { + let host = WindowHost() + let client = AHPClient(transport: await host.open()) + try await client.connect() + let first = Task { try await client.subscribe("ahp-terminal:/t1", options: .windowed(windowBytes: 4096)) } + let second = Task { try await client.subscribe("ahp-terminal:/t1", options: .windowed(windowBytes: 2048)) } + let a = try await first.value + let b = try await second.value + XCTAssertTrue(a === b) + let state = await client.connectionState + XCTAssertEqual(state, .connected) + await client.shutdown() + } + + func testCancellingCoalescedTCPSendAndEOFBarrier() async throws { + let host = WindowHost() + await host.setHostLimits(window: 1, frame: 1024) + let client = AHPClient(transport: await host.open()) + try await client.connect() + let sub = try await client.subscribe("ahp-tcp:/t1", options: .windowed()) + try await sub.send(method: "tcp/data", params: TcpDataParams(channel: sub.channel, data: "YQ==")) + let first = Task { try await sub.send(method: "tcp/data", params: TcpDataParams(channel: sub.channel, data: "Yg==")) } + try await windowEventually { await sub.status().bufferedSendBytes > 0 } + let second = Task { try await sub.send(method: "tcp/data", params: TcpDataParams(channel: sub.channel, data: "Yw==")) } + try await Task.sleep(for: .milliseconds(5)) + second.cancel() + do { try await second.value; XCTFail("cancelled merged send succeeded") } catch is CancellationError {} + try await host.creditAll(sub.channel) + try await first.value + let sent = await host.messages().filter { $0.method == "tcp/data" } + XCTAssertEqual(try sent.last?.decode(TcpDataParams.self).data, "Yg==") + try await host.creditAll(sub.channel) + try await sub.send(method: "tcp/eof", params: TcpEofParams(channel: sub.channel)) + do { + try await sub.send(method: "tcp/data", params: TcpDataParams(channel: sub.channel, data: "ZA==")) + XCTFail("data after EOF") + } catch is TransportError {} + await client.shutdown() + } + + func testBlockedManagedDispatchDoesNotBlockReconnectSupervisor() async throws { + let host = WindowHost() + await host.setHostLimits(window: 1, frame: 1024) + let multi = MultiHostClient() + _ = try await multi.add(HostConfig(id: "host", label: "Host", transportFactory: { _ in await host.open() }) + .withSessionSummaryRefreshOnConnect(false).withReconnectPolicy(.immediateForever)) + try await windowEventually { await multi.host("host")?.state.isConnected == true } + let sub = try await multi.subscribe(host: "host", uri: "ahp-terminal:/t1", options: .windowed()) + _ = try await multi.dispatch(host: "host", action: .terminalInput(.init(type: .terminalInput, data: "a")), + channel: sub.channel) + let pending = Task { try await multi.dispatch(host: "host", + action: .terminalInput(.init(type: .terminalInput, data: "b")), channel: sub.channel) } + try await windowEventually { await sub.status().bufferedSendBytes > 0 } + let generation = await multi.host("host")!.generation + try await multi.reconnect("host") + try await windowEventually { + guard let state = await multi.host("host") else { return false } + return state.state.isConnected && state.generation > generation + } + do { _ = try await pending.value; XCTFail("unsent data replayed automatically") } catch {} + await multi.shutdown() + } + + func testActualTransportDefaultsAndLegacyTransportFallback() async throws { + let url = URLSessionWebSocketTransport(url: URL(string: "ws://127.0.0.1:1")!) + let capabilities = await url.windowingCapabilities + XCTAssertEqual(capabilities.maximumFrameBytes, + 2 * (try XCTUnwrap(capabilities.maximumReceiveMessageUTF8Bytes) / 3)) + try await url.close() + let native = NWConnectionWebSocketTransport(url: URL(string: "ws://127.0.0.1:1")!) + let nativeCapabilities = await native.windowingCapabilities + XCTAssertEqual(nativeCapabilities.maximumFrameBytes, 1_048_576) + try await native.close() + + let host = WindowHost() + let transport = await host.open() + let client = AHPClient(transport: LegacyWindowTransport(transport)) + try await client.connect() + let sub = try await client.subscribe("ahp-terminal:/t1", options: .windowed()) + let status = await sub.status() + XCTAssertEqual(status.mode, "ordinary") + XCTAssertEqual(status.fallbackReason, "transport") + let first = try await sub.next() + XCTAssertEqual(first?.message.method, "channel/snapshot") + XCTAssertEqual(first?.receivedBytes, 0) + await client.shutdown() + } + + func testCustomCoalescerManualIdempotentCreditAndModeChanges() async throws { + let custom = AHPCoalescer(method: "action") { earlier, later, context in + guard context.direction == "inbound" else { return nil } + var a = try earlier.decode(ActionEnvelope.self) + let b = try later.decode(ActionEnvelope.self) + guard case .terminalData(let x) = a.action, case .terminalData(let y) = b.action else { return nil } + a.action = .terminalData(.init(type: .terminalData, data: (x.data + y.data).uppercased())) + a.serverSeq = b.serverSeq + return try .init(method: "action", params: a) + } + let host = WindowHost() + let client = AHPClient(transport: await host.open(), config: .init(windowing: .init(coalescers: [custom]))) + try await client.connect() + let sub = try await client.subscribe("ahp-terminal:/t1", options: .windowed(receive: .manual)) + let snapshot = try await windowNext(sub) + try await host.action(sub.channel, data: "a", seq: 2) + try await host.action(sub.channel, data: "b", seq: 3) + try await windowEventually { await sub.lastAppliedServerSeq == 3 } + let merged = try await windowNext(sub) + if case .event(.action(let envelope)) = merged.payload, + case .terminalData(let action) = envelope.action { + XCTAssertEqual(action.data, "AB") + XCTAssertEqual(envelope.serverSeq, 3) + } else { XCTFail("expected merged typed action") } + await merged.release(); await merged.release(); await snapshot.release() + try await windowEventually { await sub.status().bufferedReceiveBytes == 0 } + _ = try await sub.update(options: .init()) + let ordinary = await sub.status() + XCTAssertEqual(ordinary.mode, "ordinary") + let ordinarySnapshot = try await windowNext(sub) + XCTAssertEqual(ordinarySnapshot.receivedBytes, 0) + _ = try await sub.update(options: .windowed(receive: .manual)) + let newSnapshot = try await windowNext(sub) + await newSnapshot.release() + let status = await sub.status() + XCTAssertEqual(status.mode, "windowed") + XCTAssertEqual(status.bufferedReceiveBytes, 0) + await client.shutdown() + } + + func testBlockedChannelFairnessCancellationAndBackpressure() async throws { + let host = WindowHost() + await host.setHostLimits(window: 1, frame: 1024) + let client = AHPClient(transport: await host.open(), config: .init(windowing: .init( + maxBufferedSendBytes: 20_000))) + try await client.connect() + let a = try await client.subscribe("ahp-terminal:/a", options: .windowed()) + let b = try await client.subscribe("ahp-terminal:/b", options: .windowed()) + let large = Task { try await a.send(method: "custom/data", params: WindowData( + channel: a.channel, data: String(repeating: "漢😀", count: 500))) } + try await windowEventually { await a.status().outstandingSendBytes > 0 } + try await b.send(method: "custom/data", params: WindowData(channel: b.channel, data: "other")) + let messages = await host.messages() + XCTAssertTrue(messages.contains { $0.channel == b.channel && $0.method == "custom/data" }) + let waiting = Task { try await b.send(method: "custom/data", params: WindowData(channel: b.channel, data: "cancel")) } + try await windowEventually { await b.status().bufferedSendBytes > 0 } + waiting.cancel() + do { try await waiting.value; XCTFail("cancelled send succeeded") } catch is CancellationError {} + try await windowEventually { await b.status().bufferedSendBytes == 0 } + let head = Task { try await b.send(method: "custom/data", params: WindowData(channel: b.channel, data: "head")) } + try await windowEventually { await b.status().bufferedSendBytes > 0 } + do { + try await b.send(method: "custom/data", params: WindowData(channel: b.channel, + data: String(repeating: "x", count: 30_000))) + XCTFail("unbounded queue admission") + } catch AHPWindowingError.backpressure {} + head.cancel() + _ = try? await head.value + let writable = Task { try await b.waitUntilWritable() } + try await host.creditAll(b.channel) + try await writable.value + await client.shutdown() + _ = try? await large.value + } + + func testManagedReconnectSnapshotReplayLiveAndDropped() async throws { + for recovery in ["snapshot", "replay", "live", "dropped"] { + let host = WindowHost() + await host.setRecovery(recovery) + let config = HostConfig(id: "host", label: "Host", transportFactory: { _ in await host.open() }) + .withSessionSummaryRefreshOnConnect(false).withReconnectPolicy(.immediateForever) + let multi = MultiHostClient() + _ = try await multi.add(config) + try await windowEventually { await multi.host("host")?.state.isConnected == true } + let sub = try await multi.subscribe(host: "host", uri: "ahp-terminal:/t1", + options: .windowed(receive: .manual)) + let initial = try await windowNext(sub) + try await host.action(sub.channel, data: "retained", seq: 5) + try await windowEventually { await sub.lastAppliedServerSeq == 5 } + let before = await multi.host("host")!.generation + await host.disconnect() + try await windowEventually { + guard let handle = await multi.host("host") else { return false } + return handle.generation > before && handle.state.isConnected + } + let requests = await host.reconnectRequests() + XCTAssertEqual(requests.last?.windows?.items.first?.lastAppliedServerSeq, 5) + let retained = try await windowNext(sub) + await initial.release(); await retained.release() + try await Task.sleep(for: .milliseconds(15)) + let credits = await host.currentCredits() + XCTAssertEqual(credits, 0, "old epoch must not credit new connection") + if recovery == "dropped" { + do { _ = try await sub.next(); XCTFail("missing window stayed open") } + catch AHPWindowingError.subscriptionDropped {} + let handle = await multi.host("host") + XCTAssertFalse(handle!.subscriptions.contains(sub.channel)) + } else if recovery == "snapshot" { + let replacement = try await windowNext(sub) + XCTAssertEqual(replacement.message.method, "channel/snapshot") + await replacement.release() + try await windowEventually { await host.currentCredits() > 0 } + } else { + try await host.action(sub.channel, data: "new", seq: 6) + let item = try await windowNext(sub) + XCTAssertEqual(item.message.method, "action") + await item.release() + try await windowEventually { await host.currentCredits() > 0 } + } + await multi.shutdown() + } + } + + func testTCPIsNotResumedAndClosingOwnerRemovesManagedSubscription() async throws { + let host = WindowHost() + let multi = MultiHostClient() + _ = try await multi.add(HostConfig(id: "host", label: "Host", transportFactory: { _ in await host.open() }) + .withSessionSummaryRefreshOnConnect(false).withReconnectPolicy(.immediateForever)) + try await windowEventually { await multi.host("host")?.state.isConnected == true } + let tcp = try await multi.subscribe(host: "host", uri: "ahp-tcp:/t1", options: .windowed(receive: .manual)) + _ = try await tcp.next() + let terminal = try await multi.subscribe(host: "host", uri: "ahp-terminal:/t1", options: .windowed()) + try await terminal.close() + let before = await multi.host("host")!.generation + await host.disconnect() + try await windowEventually { + guard let value = await multi.host("host") else { return false } + return value.generation > before && value.state.isConnected + } + let request = await host.reconnectRequests().last + XCTAssertFalse(request?.subscriptions.contains(tcp.channel) ?? false) + XCTAssertFalse(request?.subscriptions.contains(terminal.channel) ?? false) + do { _ = try await tcp.next(); XCTFail("TCP survived disconnect") } catch {} + await multi.shutdown() + } +} + +private struct WindowData: Codable, Sendable { let channel: String; let data: String } +private final class LegacyWindowTransport: AHPTransport, Sendable { + let base: InMemoryTransport + init(_ base: InMemoryTransport) { self.base = base } + func send(_ message: TransportMessage) async throws { try await base.send(message) } + func recv() async throws -> TransportMessage? { try await base.recv() } + func close() async throws { try await base.close() } +} + +private final class ParsedEnvelopeTransport: AHPWindowingTransport, Sendable { + let base: InMemoryTransport + init(_ base: InMemoryTransport) { self.base = base } + var windowingCapabilities: AHPTransportCapabilities { + get async { .init(preservesOriginalJSONText: true) } + } + func send(_ message: TransportMessage) async throws { try await base.send(message) } + func recv() async throws -> TransportMessage? { try await base.recv() } + func close() async throws { try await base.close() } + func recvWireMessage() async throws -> AHPWireMessage? { + guard let message = try await base.recv(), case .text(let text) = message else { return nil } + return .init(message: .parsed(try message.intoParsed()), originalJSONText: text) + } +} + +actor WindowHost { + private var current: InMemoryTransport? + private var connection = 0 + private var records: [(Int, AHPNotificationMessage, Int)] = [] + private var reconnects: [ReconnectParams] = [] + private var recovery = "replay" + private var hostLimits = ChannelReceiveLimits(windowBytes: 65_536, maximumFrameBytes: 1024) + private var subscriptions = 0 + private var failReconnect = false + private var failedResubscribe: String? + private var holdResubscribe = false + + func setRecovery(_ value: String) { recovery = value } + func setFallback(failing channel: String? = nil, hold: Bool = false) { + failReconnect = true + failedResubscribe = channel + holdResubscribe = hold + } + func setHostLimits(window: Int, frame: Int) { hostLimits = .init(windowBytes: window, maximumFrameBytes: frame) } + func open() -> InMemoryTransport { + let (client, host) = InMemoryTransport.pair() + connection += 1 + let number = connection + current = host + Task { await self.drive(host, number: number) } + return client + } + func disconnect() async { try? await current?.close() } + func sendRaw(_ text: String) async throws { try await current!.send(.text(text)) } + func messages() -> [AHPNotificationMessage] { records.map { $0.1 } } + func reconnectRequests() -> [ReconnectParams] { reconnects } + func subscribeCount() -> Int { subscriptions } + func currentCredits() -> Int { + records.filter { $0.0 == connection && $0.1.method == "channel/credit" }.count + } + func action(_ channel: String, data: String, seq: Int) async throws { + let params = ActionEnvelope(channel: channel, action: .terminalData(.init(type: .terminalData, data: data)), + serverSeq: seq) + try await current!.send(makeNotificationWire(method: "action", params: params)) + } + func creditAll(_ channel: String) async throws { + let size = records.filter { $0.0 == connection && $0.1.channel == channel && + countsAgainstWindow($0.1.method) }.reduce(0) { $0 + $1.2 } + try await current!.send(makeNotificationWire(method: "channel/credit", + params: ChannelCreditParams(channel: channel, consumedBytes: size))) + } + private func snapshot(_ channel: String, transport: InMemoryTransport) async throws { + let snapshot = Snapshot(resource: channel, state: .empty(EmptyState()), fromSeq: 1) + try await transport.send(makeNotificationWire(method: "channel/snapshot", + params: ChannelSnapshotParams(channel: channel, snapshot: snapshot))) + } + private func drive(_ transport: InMemoryTransport, number: Int) async { + while let frame = try? await transport.recv() { + guard case .text(let text) = frame, + let object = try? JSONSerialization.jsonObject(with: Data(text.utf8)) as? [String: Any], + let method = object["method"] as? String, + let params = object["params"] as? [String: Any], + let bytes = try? JSONSerialization.data(withJSONObject: params) else { continue } + if let id = object["id"] as? Int { + do { + switch method { + case "initialize": + let result = InitializeResult(protocolVersion: "1.0.0", serverSeq: 1, snapshots: [ + Snapshot(resource: RootResourceURI, state: .root(RootState(agents: [])), fromSeq: 1)]) + try await transport.send(makeResponseWire(id: id, result: result)) + case "subscribe": + subscriptions += 1 + let request = try JSONDecoder().decode(SubscribeParams.self, from: bytes) + if number > 1, holdResubscribe { continue } + if number > 1, request.channel == failedResubscribe { + try await transport.send(.encoded(.errorResponse(id: id, + error: JsonRpcError(code: -32000, message: "subscription unavailable")))) + continue + } + if let offer = request.flowControl?.receive { + let limits = ChannelFlowControl(clientReceive: offer, hostReceive: hostLimits) + try await transport.send(makeResponseWire(id: id, result: SubscribeResult(flowControl: limits))) + try await snapshot(request.channel, transport: transport) + } else { + try await transport.send(makeResponseWire(id: id, result: SubscribeResult( + snapshot: Snapshot(resource: request.channel, state: .empty(EmptyState()), fromSeq: 1)))) + } + case "reconnect": + let request = try JSONDecoder().decode(ReconnectParams.self, from: bytes) + reconnects.append(request) + if failReconnect { + try await transport.send(.encoded(.errorResponse(id: id, + error: JsonRpcError(code: -32601, message: "reconnect unsupported")))) + continue + } + let items = recovery == "dropped" ? [] : (request.windows?.items ?? []).map { + ResumedChannelSubscription(channel: $0.channel, + flowControl: .init(clientReceive: $0.flowControl.receive, hostReceive: hostLimits), + recovery: recovery == "snapshot" ? .snapshot : recovery == "live" ? .live : .replay) + } + // A top-level snapshot response must not prune kept windowed channels. + try await transport.send(makeResponseWire(id: id, result: ReconnectSnapshotResult( + type: .snapshot, snapshots: [ + Snapshot(resource: RootResourceURI, state: .root(RootState(agents: [])), fromSeq: 1)], + windows: request.windows == nil ? nil : .init(items: items)))) + if recovery == "snapshot" { + for item in items { try await snapshot(item.channel, transport: transport) } + } + default: break + } + } catch { return } + } else if let notification = try? AHPNotificationMessage(method: method, paramsJSON: bytes) { + records.append((number, notification, ahpUTF16Bytes(text))) + } + } + } +} + +private func windowEventually(_ predicate: @escaping @Sendable () async -> Bool) async throws { + let deadline = ContinuousClock.now + .seconds(2) + while ContinuousClock.now < deadline { + if await predicate() { return } + try await Task.sleep(for: .milliseconds(2)) + } + + throw TransportError.protocol("test condition timed out") +} + +private func windowNext(_ subscription: AHPSubscription) async throws -> AHPDelivery { + let value = try await subscription.next() + return try XCTUnwrap(value) +} diff --git a/clients/swift/Examples/Package.swift b/clients/swift/Examples/Package.swift new file mode 100644 index 000000000..78a2a3dc5 --- /dev/null +++ b/clients/swift/Examples/Package.swift @@ -0,0 +1,37 @@ +// swift-tools-version: 6.0 + +import PackageDescription + +// Runnable examples for the AgentHostProtocol Swift client. +// +// This is a separate local package so that the published package manifest at +// the repository root does not gain example targets. It depends on the root +// package by path. +// +// Run from the repository root: +// swift run --package-path clients/swift/Examples WindowingBasic +// swift run --package-path clients/swift/Examples WindowingBackpressure + +let ahp: [Target.Dependency] = [ + .product(name: "AgentHostProtocol", package: "AgentHostProtocol"), + .product(name: "AgentHostProtocolClient", package: "AgentHostProtocol"), +] + +let package = Package( + name: "AgentHostProtocolExamples", + platforms: [ + .macOS(.v13), + ], + products: [ + .executable(name: "WindowingBasic", targets: ["WindowingBasic"]), + .executable(name: "WindowingBackpressure", targets: ["WindowingBackpressure"]), + ], + dependencies: [ + .package(name: "AgentHostProtocol", path: "../../.."), + ], + targets: [ + .target(name: "WindowingDemoHost", dependencies: ahp), + .executableTarget(name: "WindowingBasic", dependencies: ahp + ["WindowingDemoHost"]), + .executableTarget(name: "WindowingBackpressure", dependencies: ahp + ["WindowingDemoHost"]), + ] +) diff --git a/clients/swift/Examples/Sources/WindowingBackpressure/main.swift b/clients/swift/Examples/Sources/WindowingBackpressure/main.swift new file mode 100644 index 000000000..c9cfeec5e --- /dev/null +++ b/clients/swift/Examples/Sources/WindowingBackpressure/main.swift @@ -0,0 +1,163 @@ +// Windowed delivery with backpressure and a custom coalescer. +// +// Part 1, receiving: the client renders slowly. While it is busy, the host +// fills the window and the client queues what arrives. Adjacent queued +// messages are combined before the renderer sees them: +// +// - The built-in coalescer joins adjacent `terminal/data` output. +// - A custom coalescer keeps only the last of adjacent `terminal/resized` +// actions, because only the final size matters. +// +// Part 2, sending: the client pastes a large block of text into the terminal. +// The host processes input slowly, so the client's send window fills. The +// client waits until the subscription is writable before it sends each chunk, +// instead of queueing the whole paste in memory. +// +// Run: swift run --package-path clients/swift/Examples WindowingBackpressure + +import AgentHostProtocol +import AgentHostProtocolClient +import Foundation +import WindowingDemoHost + +let terminal = "ahp-terminal:/demo" +func log(_ message: String) { print(" client | \(message)") } + +// A custom coalescer runs on two adjacent messages that are still queued. It +// returns the combined message, or nil to leave them separate. The client +// never passes it messages that the protocol requires to stay separate, such +// as actions with an origin. +let latestResize = AHPCoalescer(method: "action") { earlier, later, context in + guard context.direction == "inbound" else { return nil } + let first = try earlier.decode(ActionEnvelope.self).action + let second = try later.decode(ActionEnvelope.self).action + if case .terminalResized = first, case .terminalResized = second { + return later + } + return nil +} + +let host = DemoHost(DemoHostOptions(clientWindowBytes: 2048, hostWindowBytes: 1024, inputDelay: .milliseconds(15))) +let client = AHPClient(transport: host.clientTransport, config: AHPClientConfig(windowing: AHPWindowingConfig( + // Allow at most 2048 bytes of unsent data to wait behind the message being sent. + maxBufferedSendBytes: 2048, + coalescers: [latestResize] +))) +try await client.connect() +try await client.initialize(clientId: "windowing-example", protocolVersions: SUPPORTED_PROTOCOL_VERSIONS) + +// Count what the client did. This stream is for monitoring only and can drop +// events; it never affects credit. +actor EventCounts { + private(set) var counts: [String: Int] = [:] + func add(_ kind: String) { counts[kind, default: 0] += 1 } +} +let counts = EventCounts() +let flowControlEvents = await client.flowControlEvents +let monitor = Task { + for await event in flowControlEvents { + await counts.add(event.kind) + } +} + +let subscription = try await client.subscribe(terminal, options: .windowed(windowBytes: 4096)) +let accepted = await subscription.status().limits.map { String($0.clientReceive.windowBytes) } ?? "none" +log("subscribed: accepted window=\(accepted) bytes") + +// ─── Part 1: a slow renderer ──────────────────────────────────────────────── + +print("\nPart 1: slow renderer") +let progress: [StateAction] = [10, 20, 30, 40, 50, 60, 70, 80, 90, 100].map { percent in + let padded = String(repeating: " ", count: max(0, 3 - String(percent).count)) + String(percent) + return .terminalData(TerminalDataAction( + type: .terminalData, data: "[\(padded)%] compiled \(percent * 120 / 100) modules\r\n")) +} +let hostActions: [StateAction] = [ + .terminalData(TerminalDataAction(type: .terminalData, data: "Building 120 modules\r\n")), + // The user drags the corner of the terminal window. + .terminalResized(TerminalResizedAction(type: .terminalResized, cols: 100, rows: 30)), + .terminalResized(TerminalResizedAction(type: .terminalResized, cols: 110, rows: 32)), + .terminalResized(TerminalResizedAction(type: .terminalResized, cols: 120, rows: 36)), + .terminalResized(TerminalResizedAction(type: .terminalResized, cols: 132, rows: 40)), +] + progress + [ + .terminalData(TerminalDataAction(type: .terminalData, data: "Build finished\r\n")), +] +let writer = Task { + await host.whenSubscribed() + for action in hostActions { await host.sendAction(action) } +} + +// With automatic release, taking the next delivery releases the previous one. +// Breaking out of the loop keeps the subscription open for part 2. +var renders = 0 +render: for try await delivery in subscription.events { + renders += 1 + switch delivery.payload { + case .snapshot: + log("render snapshot (\(delivery.receivedBytes) bytes)") + case .event(.action(let envelope)): + switch envelope.action { + case .terminalResized(let resize): + log("render resize to \(resize.cols)x\(resize.rows) (\(delivery.receivedBytes) bytes)") + case .terminalData(let output): + let lines = output.data.components(separatedBy: "\r\n").filter { !$0.isEmpty } + log("render \(lines.count) line(s) of output (\(delivery.receivedBytes) bytes): " + + String(reflecting: lines.last ?? "")) + if output.data.contains("Build finished") { break render } + default: + break + } + default: + break + } + try await Task.sleep(for: .milliseconds(60)) // The renderer is slow; the host fills the window meanwhile. +} +log("rendered \(renders) deliveries for \(hostActions.count + 1) host messages") +await writer.value + +// ─── Part 2: pasting into a slow terminal ─────────────────────────────────── + +print("\nPart 2: pasting into a slow terminal") +let paste = (1...8).map { "echo \"line \($0) of the pasted script\"\n" } +var waits = 0 +for (index, chunk) in paste.enumerated() { + // Wait for room in the host window and the local send queue, rather than + // queueing the whole paste. With one producer, this keeps sends within the limit. + let status = await subscription.status() + if !status.writable { + log("not writable (\(status.outstandingSendBytes) bytes awaiting credit, " + + "\(status.bufferedSendBytes) bytes queued); waiting") + waits += 1 + try await subscription.waitUntilWritable() + } + log("send chunk \(index + 1) of \(paste.count)") + let input = StateAction.terminalInput(TerminalInputAction(type: .terminalInput, data: chunk)) + do { + // Returns when the message is written to the transport. + try await client.dispatch(input, channel: terminal) + } catch AHPWindowingError.backpressure { + // If another producer filled the queue first, the send fails instead of + // growing the queue. Wait, then send the chunk again. + try await subscription.waitUntilWritable() + try await client.dispatch(input, channel: terminal) + } +} +// Wait until the host has processed, and credited, every chunk. +for await status in await subscription.statusChanges +where status.outstandingSendBytes == 0 && status.bufferedSendBytes == 0 { + break +} +let intact = await host.receivedInput == paste.joined() +log("sent \(paste.count) chunks, waiting \(waits) times; host received the paste intact: \(intact)") + +// Monitoring events are delivered asynchronously; give the last ones a moment to arrive. +try await Task.sleep(for: .milliseconds(20)) +print("\nFlow-control events:") +for (kind, count) in await counts.counts.sorted(by: { $0.key < $1.key }) { + print(" \(kind): \(count)") +} + +monitor.cancel() +try await subscription.close() +await host.close() +await client.shutdown() diff --git a/clients/swift/Examples/Sources/WindowingBasic/main.swift b/clients/swift/Examples/Sources/WindowingBasic/main.swift new file mode 100644 index 000000000..aa74094d9 --- /dev/null +++ b/clients/swift/Examples/Sources/WindowingBasic/main.swift @@ -0,0 +1,79 @@ +// Basic windowed delivery. +// +// The client offers windowed delivery on a terminal subscription. The host +// accepts a smaller window, then sends a snapshot and terminal output. When +// the window is full, the host waits; each time the client takes the next +// delivery, it releases the previous one and returns credit, so the host can +// continue. +// +// Run: swift run --package-path clients/swift/Examples WindowingBasic + +import AgentHostProtocol +import AgentHostProtocolClient +import WindowingDemoHost + +let terminal = "ahp-terminal:/demo" +let output = [ + "$ npm test", + "> demo@1.0.0 test", + "> node --test", + "ok 1 - parses the configuration file", + "ok 2 - rejects an unknown option", + "ok 3 - reads input from stdin", + "ok 4 - writes the report", + "# pass 4, fail 0", +] + +func log(_ message: String) { print(" client | \(message)") } + +// The host accepts at most a 1024-byte window, so it fills after a few messages. +let host = DemoHost(DemoHostOptions(clientWindowBytes: 1024, hostWindowBytes: 4096)) +let client = AHPClient(transport: host.clientTransport) +try await client.connect() +try await client.initialize(clientId: "windowing-example", protocolVersions: SUPPORTED_PROTOCOL_VERSIONS) + +// Windowing is off by default. Opt in for this subscription and offer a +// 4096-byte window. Coalescing is turned off so that every host message is a +// separate delivery; see WindowingBackpressure for coalescing. +let subscription = try await client.subscribe( + terminal, options: .windowed(windowBytes: 4096, coalescing: false)) +let status = await subscription.status() +log("subscribed: mode=\(status.mode), " + + "accepted window=\(status.limits.map { String($0.clientReceive.windowBytes) } ?? "none") bytes") + +// The host writes terminal output as fast as the window allows. +let writer = Task { + await host.whenSubscribed() + for line in output { + await host.sendAction(.terminalData(TerminalDataAction(type: .terminalData, data: "\(line)\r\n"))) + } +} + +// With automatic release (the default for `events`), taking the next delivery +// releases the previous one. The client then returns credit for it. +var lines = 0 +for try await delivery in subscription.events { + let buffered = await subscription.status().bufferedReceiveBytes + switch delivery.payload { + case .snapshot: + log("snapshot (\(delivery.receivedBytes) bytes; \(buffered) bytes not yet released)") + case .event(.action(let envelope)): + if case .terminalData(let action) = envelope.action { + let text = String(action.data.reversed().drop(while: \.isWhitespace).reversed()) + log("output \(String(reflecting: text)) (\(delivery.receivedBytes) bytes; \(buffered) bytes not yet released)") + lines += 1 + } + default: + break + } + try await Task.sleep(for: .milliseconds(20)) // Simulate rendering the delivery. + if lines == output.count { break } +} + +log("received all \(lines) lines") +await writer.value +// Leaving the loop does not unsubscribe or release the last delivery; closing +// the subscription does both. +try await subscription.close() +await host.close() +await client.shutdown() diff --git a/clients/swift/Examples/Sources/WindowingDemoHost/DemoHost.swift b/clients/swift/Examples/Sources/WindowingDemoHost/DemoHost.swift new file mode 100644 index 000000000..28529ae0a --- /dev/null +++ b/clients/swift/Examples/Sources/WindowingDemoHost/DemoHost.swift @@ -0,0 +1,226 @@ +// A small in-memory AHP host for the windowing examples. +// +// It implements only what the examples need: `initialize`, a windowed +// `subscribe` to one terminal, and the host side of the window rules: +// +// - It counts every data message it sends, in UTF-16 bytes, and stops while +// `sent - consumed` is at or above the client's window. +// - It counts every data message the client sends, and returns +// `channel/credit` after it has processed each one. +// +// A real host applies the same rules to every windowed subscription. + +import AgentHostProtocol +import AgentHostProtocolClient +import Foundation + +public struct DemoHostOptions: Sendable { + /// The largest window the host accepts for data that it sends. + public var clientWindowBytes: Int + /// The window the host offers for data that the client sends. + public var hostWindowBytes: Int + /// How long the host takes to process one message from the client. + public var inputDelay: Duration + + public init(clientWindowBytes: Int, hostWindowBytes: Int, inputDelay: Duration = .zero) { + self.clientWindowBytes = clientWindowBytes + self.hostWindowBytes = hostWindowBytes + self.inputDelay = inputDelay + } +} + +private let maximumFrameBytes = 16_384 +private func utf16Bytes(_ text: String) -> Int { text.utf16.count * 2 } +private func log(_ message: String) { print(" host | \(message)") } + +private struct Notification: Encodable { + var jsonrpc = "2.0" + var method: String + var params: Params +} + +private struct Response: Encodable { + var jsonrpc = "2.0" + var id: Int + var result: Result +} + +private struct PendingInput: Sendable { + let bytes: Int + let action: StateAction +} + +public actor DemoHost { + /// The half of the transport to give to `AHPClient`. + public nonisolated let clientTransport: InMemoryTransport + /// Terminal input that the host has processed, in order. + public private(set) var receivedInput = "" + + private let options: DemoHostOptions + private let transport: InMemoryTransport + private let encoder: JSONEncoder + private var limits: ChannelFlowControl? + private var channel = "" + private var serverSeq = 0 + private var subscribed = false + private var subscribedWaiters: [CheckedContinuation] = [] + + // Data that the host sends to the client. + private var sentBytes = 0 + private var consumedBytes = 0 + private var creditWaiter: CheckedContinuation? + + // Data that the client sends to the host. + private var receivedBytes = 0 + private var processedBytes = 0 + private let inputs: AsyncStream.Continuation + + public init(_ options: DemoHostOptions) { + self.options = options + (clientTransport, transport) = InMemoryTransport.pair() + // Byte counts depend only on the serialized text, so a stable key order + // and unescaped slashes match what other AHP implementations send. + encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] + let (pending, inputs) = AsyncStream.makeStream(of: PendingInput.self) + self.inputs = inputs + + Task { await self.run() } + // Process input one message at a time, then return credit for it. + Task { + for await input in pending { + try? await Task.sleep(for: options.inputDelay) + await self.processed(input) + } + } + } + + /// Returns after the host has accepted the subscription and sent the snapshot. + public func whenSubscribed() async { + if subscribed { return } + await withCheckedContinuation { subscribedWaiters.append($0) } + } + + /// Send one terminal action, first waiting for credit if the client's window is full. + public func sendAction(_ action: StateAction) async { + guard let limits else { preconditionFailure("sendAction before the subscription was accepted") } + let window = limits.clientReceive.windowBytes + if sentBytes - consumedBytes >= window { + log("window full (\(sentBytes - consumedBytes) of \(window) bytes outstanding); pausing output") + while sentBytes - consumedBytes >= window { + await withCheckedContinuation { creditWaiter = $0 } + } + log("credit received (consumedBytes=\(consumedBytes)); resuming output") + } + serverSeq += 1 + await sendData(Notification(method: "action", params: ActionEnvelope( + channel: channel, action: action, serverSeq: serverSeq))) + } + + public func close() async { + inputs.finish() + try? await transport.close() + } + + private func serialize(_ message: some Encodable) -> String { + // The host only encodes generated protocol models, which always encode. + String(decoding: try! encoder.encode(message), as: UTF8.self) + } + + private func sendData(_ message: some Encodable) async { + let text = serialize(message) + sentBytes += utf16Bytes(text) + try? await transport.send(.text(text)) + } + + private func sendControl(_ message: some Encodable) async { + try? await transport.send(.text(serialize(message))) + } + + private func run() async { + while let frame = try? await transport.recv() { + guard case .text(let text) = frame, + let message = try? JSONSerialization.jsonObject(with: Data(text.utf8)) as? [String: Any], + let method = message["method"] as? String, + let params = message["params"] as? [String: Any], + let paramsJSON = try? JSONSerialization.data(withJSONObject: params) + else { continue } + switch method { + case "initialize": + guard let id = message["id"] as? Int else { continue } + await sendControl(Response(id: id, result: InitializeResult( + protocolVersion: SUPPORTED_PROTOCOL_VERSIONS[0], serverSeq: 0, snapshots: []))) + case "subscribe": + guard let id = message["id"] as? Int, + let request = try? JSONDecoder().decode(SubscribeParams.self, from: paramsJSON) + else { continue } + await acceptSubscription(id: id, request: request) + case "channel/credit": + // Credit is a running total, so ignore values that are not larger. + if let consumed = params["consumedBytes"] as? Int, consumed > consumedBytes { + consumedBytes = consumed + creditWaiter?.resume() + creditWaiter = nil + } + case "dispatchAction": + if let request = try? JSONDecoder().decode(DispatchActionParams.self, from: paramsJSON) { + receiveInput(text, action: request.action) + } + default: + break + } + } + } + + private func acceptSubscription(id: Int, request: SubscribeParams) async { + guard let offered = request.flowControl?.receive else { + preconditionFailure("the examples expect a windowed subscribe") + } + channel = request.channel + // A host can lower the limits that the client offered, but never raise them. + let limits = ChannelFlowControl( + clientReceive: ChannelReceiveLimits( + windowBytes: min(offered.windowBytes, options.clientWindowBytes), + maximumFrameBytes: min(offered.maximumFrameBytes, maximumFrameBytes)), + hostReceive: ChannelReceiveLimits( + windowBytes: options.hostWindowBytes, maximumFrameBytes: maximumFrameBytes)) + self.limits = limits + log("accepted windowed delivery: client window \(limits.clientReceive.windowBytes) bytes " + + "(offered \(offered.windowBytes)), host window \(limits.hostReceive.windowBytes) bytes") + await sendControl(Response(id: id, result: SubscribeResult(flowControl: limits))) + + // The first data message on a windowed subscription is always the snapshot. + let state = TerminalState( + title: "demo shell", + content: [], + lifecycle: .running(TerminalRunningLifecycleState(status: .running)), + claim: .client(TerminalClientClaim(kind: .client, clientId: "windowing-example"))) + await sendData(Notification(method: "channel/snapshot", params: ChannelSnapshotParams( + channel: channel, + snapshot: Snapshot(resource: channel, state: .terminal(state), fromSeq: serverSeq)))) + + subscribed = true + for waiter in subscribedWaiters { waiter.resume() } + subscribedWaiters.removeAll() + } + + private func receiveInput(_ text: String, action: StateAction) { + let bytes = utf16Bytes(text) + receivedBytes += bytes + if let limits = limits?.hostReceive, + // The client can go past the window only with the message that fills it. + receivedBytes - processedBytes > limits.windowBytes + limits.maximumFrameBytes { + log("client went past the host window") + } + inputs.yield(PendingInput(bytes: bytes, action: action)) + } + + private func processed(_ input: PendingInput) async { + if case .terminalInput(let terminalInput) = input.action { + receivedInput += terminalInput.data + } + processedBytes += input.bytes + await sendControl(Notification(method: "channel/credit", params: ChannelCreditParams( + channel: channel, consumedBytes: processedBytes))) + } +} diff --git a/clients/typescript/README.md b/clients/typescript/README.md index e7ea7371f..e2aafe3ff 100644 --- a/clients/typescript/README.md +++ b/clients/typescript/README.md @@ -118,6 +118,235 @@ Acquire the lease before awaiting `ready`: its event iterator is attached before the `subscribe` request is sent, so actions delivered during the snapshot round-trip remain ordered behind that snapshot instead of being lost. +## Windowed delivery and backpressure + +Windowing is **off by default**. Opt in on each subscription: + +```ts +const { subscription } = await client.subscribe(terminalUri, { windowed: true }); +for await (const event of subscription) { + if (event.type === 'snapshot') { + mirror.applySnapshot(event.params.snapshot); + } else if (event.type === 'action') { + await renderTerminal(event.params); + } +} +``` + +The first event is the snapshot, whether the host accepts windowing or declines +and returns an ordinary snapshot. With automatic release (the default), asking +for the next event releases the previous one. Await work inside the loop to +slow down the host. Closing a windowed owner (including an iterator's `return()`) +unsubscribes it and stops credit, including pending timers and old receipts. +Ordinary local observer iterators retain their existing close behavior. + +The extended `DeliveryEvent` type includes snapshots, TCP data/EOF/reset, and +unknown notifications, without extending the legacy `SubscriptionEvent` union. +`subscribe(uri, { windowed: true })` returns `Subscription`; +existing ordinary calls keep their existing type and buffering behavior. +An ordinary empty-state result without a snapshot has no synthetic snapshot. + +### Examples + +Two runnable examples use an in-memory demo host that follows the window rules, +so they need no server: + +- [`examples/windowing-basic.ts`](examples/windowing-basic.ts) opts a terminal + subscription into windowed delivery and shows the host pausing when the + window is full and resuming as the client returns credit. +- [`examples/windowing-backpressure.ts`](examples/windowing-backpressure.ts) + adds a slow renderer, a custom coalescer that keeps only the latest terminal + resize, and a paste that waits for `waitWritable()` instead of queueing. + +```sh +npx tsx examples/windowing-basic.ts +npx tsx examples/windowing-backpressure.ts +npm run typecheck:examples +``` + +The examples import from `../src`; applications import from +`@microsoft/agent-host-protocol/client` instead. + +### Defaults and transport support + +```ts +const transport = await WebSocketTransport.connect(url, { + maximumFrameBytes: 1_048_576, // UTF-16 bytes, including the JSON envelope +}); +const client = new AhpClient(transport, { + windowing: { + windowBytes: 262_144, + // maximumFrameBytes: override the transport-derived default if needed + maxBufferedSendBytes: 1_048_576, + coalescers: [], + }, +}); +const { subscription } = await client.subscribe(terminalUri, { + windowed: { windowBytes: 65_536, maximumFrameBytes: 16_384 }, + // Existing delivery and view preferences remain independent. +}); +``` + +Connection configuration supplies defaults; it does not enable windowing. +Global WebSocket exposes no native maximum message size, so this adapter's +default `maximumFrameBytes` is **1,048,576**. `fromSocket(socket, options)` +accepts the same override. There are no new runtime dependencies. +When windowing is active, the scheduler also uses the adapter's cancelable +`waitWritable()` to keep the physical WebSocket send buffer near its internal +1 MiB high-water / 256 KiB low-water marks. `bufferedAmount` remains available +in transport bytes, separately from the subscription's UTF-16 accounting. + +Custom transports opt in by exposing `supportsWindowedDelivery = true`. They +must return exact text frames, UTF-8 binary frames, or parsed frames with +`rawText` containing the original JSON, and must send pre-serialized strings +verbatim. A parsed-only transport can keep its existing implementation; the +subscription reports `fallbackReason: 'transport'` and does not offer windowing. +Host rejection reports `'declined'`. There is no protocol-version gating. +TCP requires accepted windowing and never falls back. +Before using the low-level `createTcpConnection` request, check that the host +advertised `InitializeResult.tcpConnections`; raw request APIs otherwise remain +capability-agnostic. + +### Manual receive release + +```ts +const { subscription } = await client.subscribe(tcpUri, { + windowed: { receive: 'manual' }, +}); +for await (const delivery of subscription.deliveries()) { + if (delivery.event.type === 'tcpData') { + const bytes = Uint8Array.from(atob(delivery.event.params.data), c => c.charCodeAt(0)); + await sink.write(bytes); + } + delivery.release(); // idempotent; releases only this delivery's remaining credit +} +``` + +Use one owner to consume the reliable queue; do not mix `next()` and +`deliveries()` on that handle. `deliveries()` also works in auto mode and +automatically releases the previous receipt on advance. `delivery.bytes` is the +sum of the **original raw wire sizes**, not a re-serialization of its payload. + +Fragments are special: non-final fragments are credited when stored so that a +message larger than the window can finish. The final fragment is credited on +consumer release. The receipt includes every original fragment's size, but +`release()` never credits a non-final fragment twice. Windowing is not a limit +on complete-message size or total retained memory; a consumer must still avoid +retaining unbounded application data. + +### Send-side backpressure and diagnostics + +```ts +for await (const chunk of boundedSource) { + await subscription.waitWritable(); + const binary = Array.from(chunk, byte => String.fromCharCode(byte)).join(''); + await client.notifyAsync('tcp/data', { channel: tcpUri, data: btoa(binary) }); +} +await client.notifyAsync('tcp/eof', { channel: tcpUri }); +``` + +Use bounded source chunks and await each send instead of constructing an +unbounded set of promises. `notifyAsync(method, params, { signal? })` and +`dispatchAsync(channel, action, { clientSeq?, signal? })` complete when written +to the transport, not when echoed, processed, or flushed by a WebSocket. +The window, fragmentation, and a fair per-channel scheduler apply to both +async sends and the existing synchronous `notify`/`dispatch`. +The first logical send is accepted even if it exceeds `maxBufferedSendBytes`; +the window and fragmentation pace it. The cap bounds data queued behind that +head message, and coalescing does not bypass it. Queue overflow throws/rejects +with `BackpressureError` without disconnecting the client; awaited sends also +fail on overflow rather than retaining unbounded admission waiters. `waitWritable` +is advisory, not a reservation when multiple producers share a subscription. + +Cancellation removes a queued message before transmission begins; a fragmented +message that has started is completed to preserve framing. Unsent messages fail +on connection loss and are never automatically replayed. TCP EOF is ordered and +counts against the window; TCP reset bypasses the window and discards pending data. + +`subscription.flowControl` (or `client.flowControl(uri)`) reports mode, negotiated +limits, outstanding sends, queued sends, unreleased receive bytes, and writability. +`flowControlChanges()` observes changes without releasing credit. +`client.waitWritable(uri, { signal? })` is the connection-level convenience API. +`client.flowControlEvents()` provides a lossy observer stream for blocked/unblocked, +credit, fragments, coalescing, resumed, and dropped events. +Credit batching and exempt-message queue/size/rate bounds are internal policies, +not additional public configuration. +Malformed frames on accepted windowed connections and protocol violations close +the connection with a `TransportError` of kind `protocol`. + +### Coalescing and custom behavior + +Coalescing defaults to on for accepted windowing and can be disabled with +`windowed: { coalescing: false }`. It merges only adjacent queued items, adds no +delay, and never changes an item already handed to its consumer. +Built-ins concatenate received terminal output, concatenate chat deltas for the +same turn/part/metadata, and concatenate **decoded** outgoing TCP bytes. +Origins/rejections, distinct dispatch sequence numbers, snapshots, EOF/reset, +unknown methods, and different connection epochs are not merged by the built-ins. +Receive coalescing is local and independent of `delivery.maxLatencyMs`. + +Register pure custom coalescers per connection, before the built-ins: + +```ts +import type { ActionEnvelope } from '@microsoft/agent-host-protocol'; +import type { ChannelCoalescer } from '@microsoft/agent-host-protocol/client'; + +const terminalCoalescer: ChannelCoalescer = { + method: 'action', + merge(earlier, later, { direction }) { + if (direction !== 'receive') return; + const a = earlier.params as ActionEnvelope; + const b = later.params as ActionEnvelope; + if (a.action.type !== 'terminal/data' || b.action.type !== 'terminal/data') return; + return { + ...later, + params: { ...b, action: { ...b.action, data: a.action.data + b.action.data } }, + }; + }, +}; +const client = new AhpClient(transport, { windowing: { coalescers: [terminalCoalescer] } }); +``` + +A custom merge must preserve channel semantics and required events. The client +enforces channel/method identity and excludes controls, snapshots, dispatches, +origins, and cross-epoch merges. Returning `undefined` delegates to the next +custom coalescer or built-in. + +### Shared ownership and reconnect + +`ManagedSubscriptionManager.acquire(uri, owner, { windowed: ... })` exposes a +single reliable owner through `lease.deliveries()` (also +`lease.subscription.deliveries()`). Choose one lease to consume that stream; +the existing independent `lease.events` iterators remain lossy observers. +They never release the owner queue's credit. Managed status also exposes +`flowControl`, `flowControlChanges()` and `waitWritable()`. +Active windowed leases on a multi-host client's raw connection automatically +rebind their owner queue and observer fan-out before the next handshake; a dropped +subscription closes the managed entry with an observable error. +A failed per-channel re-subscribe drops only that channel, not the host. +Closing the owner removes it from future reconnect requests. + +`subscribe(uri, { windowed: true, signal })` supports local cancellation. +Canceled/timed-out requests clean up their provisional state, and late server +acceptances are accounted for or unsubscribed rather than closing the connection. + +`MultiHostClient.subscribe(hostId, uri, { windowed: true })` returns +`{ result, subscription }`; its existing two-argument call still returns the +ordinary `SubscribeResult`. `HostConfig.clientConfig.windowing` supplies defaults. +Host-aware async sends, `flowControl`, `waitWritable`, and `flowControlEvents` +are available; telemetry carries `hostId`. `HostClientHandle` checks generation +before and after awaited sends. + +The multi-host supervisor keeps the owner handle, queued deliveries and options +across reconnect. It sends `reconnect.windows` using each queue's retained +sequence, resets wire counters, and appends snapshot/replay/live recovery data +after retained items. Releasing an old-connection receipt sends no credit on the +new connection. Dropped channels end with a visible error. TCP never resumes. +Low-level `client.reconnect` also accepts the generated `windows` parameter. +Global `events()` and `deliveryEvents()` are observe-only; use owner deliveries +with `AhpStateMirror.applyDelivery` or +`MultiHostStateMirror.applyDelivery(hostId, delivery)` for reliable state feeding. + ## Reducers and state mirror The reducer functions (`rootReducer`, `sessionReducer`, diff --git a/clients/typescript/examples/tsconfig.json b/clients/typescript/examples/tsconfig.json new file mode 100644 index 000000000..510aa3ca1 --- /dev/null +++ b/clients/typescript/examples/tsconfig.json @@ -0,0 +1,8 @@ +{ + "extends": "../tsconfig.json", + "compilerOptions": { + "noEmit": true, + "rootDir": ".." + }, + "include": ["./*.ts", "../src/**/*.ts"] +} diff --git a/clients/typescript/examples/windowing-backpressure.ts b/clients/typescript/examples/windowing-backpressure.ts new file mode 100644 index 000000000..f37f8cc46 --- /dev/null +++ b/clients/typescript/examples/windowing-backpressure.ts @@ -0,0 +1,159 @@ +/** + * Windowed delivery with backpressure and a custom coalescer. + * + * Part 1, receiving: the client renders slowly. While it is busy, the host + * fills the window and the client queues what arrives. Adjacent queued + * messages are combined before the renderer sees them: + * + * - The built-in coalescer joins adjacent `terminal/data` output. + * - A custom coalescer keeps only the last of adjacent `terminal/resized` + * actions, because only the final size matters. + * + * Part 2, sending: the client pastes a large block of text into the terminal. + * The host processes input slowly, so the client's send window fills. The + * client waits until the subscription is writable before it sends each chunk, + * instead of queueing the whole paste in memory. + * + * Run: npx tsx examples/windowing-backpressure.ts + */ + +import { + AhpClient, + BackpressureError, + type ChannelCoalescer, + type FlowControlEvent, +} from '../src/client/index.js'; +import { + ActionType, + SUPPORTED_PROTOCOL_VERSIONS, + type ActionEnvelope, + type StateAction, +} from '../src/types/index.js'; +import { DemoHost } from './windowing-demo-host.js'; + +const TERMINAL = 'ahp-terminal:/demo'; +const sleep = (ms: number) => new Promise(resolve => setTimeout(resolve, ms)); +const log = (message: string) => console.log(` client | ${message}`); + +// A custom coalescer runs on two adjacent messages that are still queued. It +// returns the combined message, or undefined to leave them separate. The +// client never passes it messages that the protocol requires to stay separate, +// such as actions with an origin. +const latestResize: ChannelCoalescer = { + method: 'action', + merge(earlier, later, { direction }) { + if (direction !== 'receive') return undefined; + const first = (earlier.params as ActionEnvelope).action; + const second = (later.params as ActionEnvelope).action; + if (first.type === ActionType.TerminalResized && second.type === ActionType.TerminalResized) { + return later; + } + return undefined; + }, +}; + +const host = new DemoHost({ clientWindowBytes: 2048, hostWindowBytes: 1024, inputDelayMs: 15 }); +const client = new AhpClient(host.clientTransport, { + windowing: { + coalescers: [latestResize], + // Allow at most 2048 bytes of unsent data to wait behind the message being sent. + maxBufferedSendBytes: 2048, + }, +}); +client.connect(); +await client.initialize({ clientId: 'windowing-example', protocolVersions: SUPPORTED_PROTOCOL_VERSIONS }); + +// Count what the client did. This stream is for monitoring only and can drop +// events; it never affects credit. +const counts = new Map(); +void (async () => { + for await (const event of client.flowControlEvents()) { + counts.set(event.type, (counts.get(event.type) ?? 0) + 1); + } +})(); + +const { subscription } = await client.subscribe(TERMINAL, { windowed: { windowBytes: 4096 } }); +log(`subscribed: accepted window=${subscription.flowControl.limits?.clientReceive.windowBytes} bytes`); + +// ─── Part 1: a slow renderer ──────────────────────────────────────────────── + +console.log('\nPart 1: slow renderer'); +const hostActions: StateAction[] = [ + { type: ActionType.TerminalData, data: 'Building 120 modules\r\n' }, + // The user drags the corner of the terminal window. + { type: ActionType.TerminalResized, cols: 100, rows: 30 }, + { type: ActionType.TerminalResized, cols: 110, rows: 32 }, + { type: ActionType.TerminalResized, cols: 120, rows: 36 }, + { type: ActionType.TerminalResized, cols: 132, rows: 40 }, + ...[10, 20, 30, 40, 50, 60, 70, 80, 90, 100].map(percent => ({ + type: ActionType.TerminalData, + data: `[${String(percent).padStart(3)}%] compiled ${(percent * 120) / 100} modules\r\n`, + } satisfies StateAction)), + { type: ActionType.TerminalData, data: 'Build finished\r\n' }, +]; +void (async () => { + await host.whenSubscribed(); + for (const action of hostActions) await host.sendAction(action); +})(); + +const deliveries = subscription.deliveries(); +let renders = 0; +for (;;) { + const { value: delivery, done } = await deliveries.next(); + if (done) break; + const { event } = delivery; + renders++; + if (event.type === 'snapshot') { + log(`render snapshot (${delivery.bytes} bytes)`); + } else if (event.type === 'action') { + const action = event.params.action; + if (action.type === ActionType.TerminalResized) { + log(`render resize to ${action.cols}x${action.rows} (${delivery.bytes} bytes)`); + } else if (action.type === ActionType.TerminalData) { + const lines = action.data.split('\r\n').filter(Boolean); + log(`render ${lines.length} line(s) of output (${delivery.bytes} bytes): ${JSON.stringify(lines.at(-1))}`); + if (action.data.includes('Build finished')) break; + } + } + await sleep(60); // The renderer is slow; the host fills the window meanwhile. +} +log(`rendered ${renders} deliveries for ${hostActions.length + 1} host messages`); + +// ─── Part 2: pasting into a slow terminal ─────────────────────────────────── + +console.log('\nPart 2: pasting into a slow terminal'); +const paste = Array.from({ length: 8 }, (_, i) => `echo "line ${i + 1} of the pasted script"\n`); +let waits = 0; +for (const chunk of paste) { + // Wait for room in the host window and the local send queue, rather than + // queueing the whole paste. With one producer, this keeps sends within the limit. + if (!subscription.flowControl.writable) { + const { outstandingSendBytes, bufferedSendBytes } = subscription.flowControl; + log(`not writable (${outstandingSendBytes} bytes awaiting credit, ${bufferedSendBytes} bytes queued); waiting`); + waits++; + await subscription.waitWritable(); + } + log(`send chunk ${paste.indexOf(chunk) + 1} of ${paste.length}`); + try { + // Resolves when the message is written to the transport. + await client.dispatchAsync(TERMINAL, { type: ActionType.TerminalInput, data: chunk }); + } catch (error) { + // If another producer filled the queue first, the send fails instead of + // growing the queue. Wait, then send the chunk again. + if (!(error instanceof BackpressureError)) throw error; + await subscription.waitWritable(); + await client.dispatchAsync(TERMINAL, { type: ActionType.TerminalInput, data: chunk }); + } +} +// Wait until the host has processed, and credited, every chunk. +for await (const status of subscription.flowControlChanges()) { + if (status.outstandingSendBytes === 0 && status.bufferedSendBytes === 0) break; +} +log(`sent ${paste.length} chunks, waiting ${waits} times; host received the paste intact: ${host.receivedInput === paste.join('')}`); + +console.log('\nFlow-control events:'); +for (const [type, count] of [...counts].sort()) console.log(` ${type}: ${count}`); + +await subscription.close(); +host.close(); +await client.shutdown(); diff --git a/clients/typescript/examples/windowing-basic.ts b/clients/typescript/examples/windowing-basic.ts new file mode 100644 index 000000000..872e17870 --- /dev/null +++ b/clients/typescript/examples/windowing-basic.ts @@ -0,0 +1,74 @@ +/** + * Basic windowed delivery. + * + * The client offers windowed delivery on a terminal subscription. The host + * accepts a smaller window, then sends a snapshot and terminal output. When + * the window is full, the host waits; each time the client takes the next + * delivery, it releases the previous one and returns credit, so the host can + * continue. + * + * Run: npx tsx examples/windowing-basic.ts + */ + +import { AhpClient } from '../src/client/index.js'; +import { ActionType, SUPPORTED_PROTOCOL_VERSIONS } from '../src/types/index.js'; +import { DemoHost } from './windowing-demo-host.js'; + +const TERMINAL = 'ahp-terminal:/demo'; +const OUTPUT = [ + '$ npm test', + '> demo@1.0.0 test', + '> node --test', + 'ok 1 - parses the configuration file', + 'ok 2 - rejects an unknown option', + 'ok 3 - reads input from stdin', + 'ok 4 - writes the report', + '# pass 4, fail 0', +]; + +const sleep = (ms: number) => new Promise(resolve => setTimeout(resolve, ms)); +const log = (message: string) => console.log(` client | ${message}`); + +// The host accepts at most a 1024-byte window, so it fills after a few messages. +const host = new DemoHost({ clientWindowBytes: 1024, hostWindowBytes: 4096 }); +const client = new AhpClient(host.clientTransport); +client.connect(); +await client.initialize({ clientId: 'windowing-example', protocolVersions: SUPPORTED_PROTOCOL_VERSIONS }); + +// Windowing is off by default. Opt in for this subscription and offer a +// 4096-byte window. Coalescing is turned off so that every host message is a +// separate delivery; see windowing-backpressure.ts for coalescing. +const { subscription } = await client.subscribe(TERMINAL, { + windowed: { windowBytes: 4096, coalescing: false }, +}); +const { mode, limits } = subscription.flowControl; +log(`subscribed: mode=${mode}, accepted window=${limits?.clientReceive.windowBytes} bytes`); + +// The host writes terminal output as fast as the window allows. +void (async () => { + await host.whenSubscribed(); + for (const line of OUTPUT) { + await host.sendAction({ type: ActionType.TerminalData, data: `${line}\r\n` }); + } +})(); + +// With automatic release (the default), taking the next delivery releases the +// previous one. The client then returns credit for it. +let lines = 0; +for await (const delivery of subscription.deliveries()) { + const { event } = delivery; + const buffered = subscription.flowControl.bufferedReceiveBytes; + if (event.type === 'snapshot') { + log(`snapshot (${delivery.bytes} bytes; ${buffered} bytes not yet released)`); + } else if (event.type === 'action' && event.params.action.type === ActionType.TerminalData) { + const text = event.params.action.data.trimEnd(); + log(`output ${JSON.stringify(text)} (${delivery.bytes} bytes; ${buffered} bytes not yet released)`); + lines++; + } + await sleep(20); // Simulate rendering the delivery. + if (lines === OUTPUT.length) break; // Leaving the loop unsubscribes. +} + +log(`received all ${lines} lines`); +host.close(); +await client.shutdown(); diff --git a/clients/typescript/examples/windowing-demo-host.ts b/clients/typescript/examples/windowing-demo-host.ts new file mode 100644 index 000000000..4c8443cba --- /dev/null +++ b/clients/typescript/examples/windowing-demo-host.ts @@ -0,0 +1,187 @@ +/** + * A small in-memory AHP host for the windowing examples. + * + * It implements only what the examples need: `initialize`, a windowed + * `subscribe` to one terminal, and the host side of the window rules: + * + * - It counts every data message it sends, in UTF-16 bytes, and stops while + * `sent - consumed` is at or above the client's window. + * - It counts every data message the client sends, and returns + * `channel/credit` after it has processed each one. + * + * A real host applies the same rules to every windowed subscription. + */ + +import { InMemoryTransport, type AhpTransport } from '../src/client/index.js'; +import { ActionType, SUPPORTED_PROTOCOL_VERSIONS, type StateAction } from '../src/types/index.js'; +import type { ChannelFlowControl } from '../src/types/common/state.js'; + +export interface DemoHostOptions { + /** The largest window the host accepts for data that it sends. */ + readonly clientWindowBytes: number; + /** The window the host offers for data that the client sends. */ + readonly hostWindowBytes: number; + /** How long the host takes to process one message from the client. */ + readonly inputDelayMs?: number; +} + +const MAXIMUM_FRAME_BYTES = 16_384; +const sleep = (ms: number) => new Promise(resolve => setTimeout(resolve, ms)); +const utf16Bytes = (text: string) => text.length * 2; +const log = (message: string) => console.log(` host | ${message}`); + +export class DemoHost { + /** The half of the transport to give to `AhpClient`. */ + readonly clientTransport: AhpTransport; + /** Terminal input that the host has processed, in order. */ + receivedInput = ''; + + private readonly transport: AhpTransport; + private readonly subscribed: Promise; + private resolveSubscribed!: () => void; + private limits?: ChannelFlowControl; + private channel = ''; + private serverSeq = 0; + + // Data that the host sends to the client. + private sentBytes = 0; + private consumedBytes = 0; + private creditWaiter?: () => void; + + // Data that the client sends to the host. + private receivedBytes = 0; + private processedBytes = 0; + private inputQueue: Promise = Promise.resolve(); + + constructor(private readonly options: DemoHostOptions) { + [this.clientTransport, this.transport] = InMemoryTransport.pair(); + this.subscribed = new Promise(resolve => { this.resolveSubscribed = resolve; }); + void this.run(); + } + + /** Resolves after the host has accepted the subscription and sent the snapshot. */ + whenSubscribed(): Promise { + return this.subscribed; + } + + /** Send one terminal action, first waiting for credit if the client's window is full. */ + async sendAction(action: StateAction): Promise { + const window = this.limits!.clientReceive.windowBytes; + if (this.sentBytes - this.consumedBytes >= window) { + log(`window full (${this.sentBytes - this.consumedBytes} of ${window} bytes outstanding); pausing output`); + while (this.sentBytes - this.consumedBytes >= window) { + await new Promise(resolve => { this.creditWaiter = resolve; }); + } + log(`credit received (consumedBytes=${this.consumedBytes}); resuming output`); + } + this.sendData({ + jsonrpc: '2.0', + method: 'action', + params: { channel: this.channel, action, serverSeq: ++this.serverSeq }, + }); + } + + close(): void { + void this.transport.close(); + } + + private sendData(message: object): void { + const text = JSON.stringify(message); + this.sentBytes += utf16Bytes(text); + void this.transport.send(text); + } + + private sendControl(message: object): void { + void this.transport.send(JSON.stringify(message)); + } + + private async run(): Promise { + for (;;) { + const frame = await this.transport.recv().catch(() => null); + if (!frame) return; + if (frame.kind !== 'text') continue; + const message = JSON.parse(frame.text); + switch (message.method) { + case 'initialize': + this.sendControl({ + jsonrpc: '2.0', + id: message.id, + result: { protocolVersion: SUPPORTED_PROTOCOL_VERSIONS[0], serverSeq: 0, snapshots: [] }, + }); + break; + case 'subscribe': + this.acceptSubscription(message); + break; + case 'channel/credit': + // Credit is a running total, so ignore values that are not larger. + if (message.params.consumedBytes > this.consumedBytes) { + this.consumedBytes = message.params.consumedBytes; + this.creditWaiter?.(); + this.creditWaiter = undefined; + } + break; + case 'dispatchAction': + this.receiveInput(frame.text, message.params.action); + break; + } + } + } + + private acceptSubscription(message: any): void { + const offered = message.params.flowControl?.receive; + if (!offered) throw new Error('the examples expect a windowed subscribe'); + this.channel = message.params.channel; + // A host can lower the limits that the client offered, but never raise them. + this.limits = { + clientReceive: { + windowBytes: Math.min(offered.windowBytes, this.options.clientWindowBytes), + maximumFrameBytes: Math.min(offered.maximumFrameBytes, MAXIMUM_FRAME_BYTES), + }, + hostReceive: { windowBytes: this.options.hostWindowBytes, maximumFrameBytes: MAXIMUM_FRAME_BYTES }, + }; + log(`accepted windowed delivery: client window ${this.limits.clientReceive.windowBytes} bytes ` + + `(offered ${offered.windowBytes}), host window ${this.limits.hostReceive.windowBytes} bytes`); + this.sendControl({ jsonrpc: '2.0', id: message.id, result: { flowControl: this.limits } }); + + // The first data message on a windowed subscription is always the snapshot. + this.sendData({ + jsonrpc: '2.0', + method: 'channel/snapshot', + params: { + channel: this.channel, + snapshot: { + resource: this.channel, + state: { + title: 'demo shell', + content: [], + lifecycle: { status: 'running' }, + claim: { kind: 'client', clientId: 'windowing-example' }, + }, + fromSeq: this.serverSeq, + }, + }, + }); + this.resolveSubscribed(); + } + + private receiveInput(text: string, action: StateAction): void { + const bytes = utf16Bytes(text); + this.receivedBytes += bytes; + const limits = this.limits!.hostReceive; + // The client can go past the window only with the message that fills it. + if (this.receivedBytes - this.processedBytes > limits.windowBytes + limits.maximumFrameBytes) { + log('client went past the host window'); + } + // Process input one message at a time, then return credit for it. + this.inputQueue = this.inputQueue.then(async () => { + await sleep(this.options.inputDelayMs ?? 0); + if (action.type === ActionType.TerminalInput) this.receivedInput += action.data; + this.processedBytes += bytes; + this.sendControl({ + jsonrpc: '2.0', + method: 'channel/credit', + params: { channel: this.channel, consumedBytes: this.processedBytes }, + }); + }); + } +} diff --git a/clients/typescript/package.json b/clients/typescript/package.json index 46a64c56e..c7a11149e 100644 --- a/clients/typescript/package.json +++ b/clients/typescript/package.json @@ -49,6 +49,7 @@ "scripts": { "build": "tsc -p tsconfig.json", "typecheck": "tsc -p tsconfig.json --noEmit", + "typecheck:examples": "tsc -p examples/tsconfig.json", "test": "node --test --import tsx test/*.test.ts", "clean": "rm -rf dist" }, diff --git a/clients/typescript/src/client/client.ts b/clients/typescript/src/client/client.ts index 36066b583..35af2ecd4 100644 --- a/clients/typescript/src/client/client.ts +++ b/clients/typescript/src/client/client.ts @@ -19,6 +19,7 @@ import type { PingParams, ReconnectParams, ReconnectResult, + ChannelDeliveryResumeOptions, ResourceReadParams, ResourceReadResult, ResourceWriteParams, @@ -75,7 +76,10 @@ import type { AuthRequiredParams } from '../types/common/notifications.js'; import type { URI } from '../types/common/state.js'; import { JsonRpcErrorCodes } from '../types/common/errors.js'; import { AsyncBroadcastQueue } from './async-queue.js'; -import type { ClientEvent, ConnectionState, SubscriptionEvent } from './events.js'; +import type { ClientEvent, ClientDeliveryEvent, ConnectionState, SubscriptionEvent, DeliveryEvent, ReceivedDelivery } from './events.js'; +import { deliveryEvent } from './events.js'; +import { WindowEngine, type WindowChannel, type WindowingConfig, type WindowOptions, + type FlowControlStatus, type FlowControlEvent } from './flow-control.js'; import { ClientClosedError, RpcError, @@ -93,6 +97,8 @@ import { /** Configuration for an {@link AhpClient}. */ export interface AhpClientConfig { + /** Defaults for opt-in subscriptions. This alone does not enable windowing. */ + windowing?: WindowingConfig; /** * Maximum time in milliseconds to wait for a request to resolve before * failing with {@link RpcTimeoutError}. Default `30000`. Set to `0` or @@ -109,6 +115,10 @@ export interface AhpClientConfig { /** Optional preferences for a `subscribe` request. */ export interface SubscribeOptions { + /** Cancel the local subscribe; a late server acceptance is unsubscribed. */ + signal?: AbortSignal; + /** Offer consumer-driven windowed delivery. Default false. */ + windowed?: boolean | WindowOptions; /** Advisory delivery preferences for this subscription. */ delivery?: SubscriptionDeliveryOptions; /** Optional client-requested shape for the returned snapshot. */ @@ -126,37 +136,97 @@ function noop(): void { /* intentionally empty */ } /** * Handle to a single resource subscription. Iterate to receive * {@link SubscriptionEvent}s. Call {@link Subscription.close} to terminate - * this consumer's iterator (the server-side subscription is released only - * when {@link AhpClient.unsubscribe} is called for this URI). + * this consumer's iterator. A windowed owner also unsubscribes; ordinary + * local observers require {@link AhpClient.unsubscribe} for server release. */ -export class Subscription implements AsyncIterableIterator { +export class Subscription implements AsyncIterableIterator { /** Channel URI this subscription is bound to. */ readonly uri: URI; /** @internal */ - private readonly inner: AsyncIterableIterator; + private readonly inner: AsyncIterableIterator; + private readMode?: 'events' | 'deliveries'; + private detached = false; /** @internal */ - constructor(uri: URI, inner: AsyncIterableIterator) { + constructor(uri: URI, inner: AsyncIterableIterator, readonly windowChannel?: WindowChannel) { this.uri = uri; this.inner = inner; } - next(): Promise> { + next(): Promise> { + if (this.detached) return Promise.resolve({ value: undefined, done: true }); + if (this.windowChannel) { + if (this.windowChannel.options.receive === 'manual') { + return Promise.reject(new TypeError('manual receive requires deliveries() and release()')); + } + if (this.readMode === 'deliveries') return Promise.reject(new TypeError('cannot mix events and deliveries')); + this.readMode = 'events'; + return this.windowChannel.next().then(result => result.done + ? { value: undefined, done: true } + : { value: result.value.event as TEvent, done: false }); + } return this.inner.next(); } - return(): Promise> { - return this.inner.return ? this.inner.return() : Promise.resolve({ value: undefined as unknown as SubscriptionEvent, done: true }); + return(): Promise> { + this.detached = true; + this.windowChannel?.close(); + return this.inner.return ? this.inner.return() : Promise.resolve({ value: undefined, done: true }); } [Symbol.asyncIterator](): this { return this; } - /** Terminate this consumer's iterator. Does not unsubscribe server-side. */ + /** Terminate this iterator; a windowed owner also releases the server subscription. */ async close(): Promise { await this.return(); } + + /** Full typed delivery surface, including snapshots, TCP and unknown methods. */ + deliveries(): AsyncIterableIterator { + if (this.readMode === 'events') throw new TypeError('cannot mix events and deliveries'); + this.readMode = 'deliveries'; + const subscription = this; + return { + [Symbol.asyncIterator]() { return this; }, + next: () => { + if (subscription.detached) return Promise.resolve({ value: undefined, done: true }); + if (subscription.windowChannel) return subscription.windowChannel.next( + subscription.windowChannel.options.receive === 'manual'); + return subscription.inner.next().then(value => value.done ? { value: undefined, done: true } : { + value: { event: value.value, notification: eventNotification(value.value), bytes: 0, release() {} }, + done: false, + }); + }, + return: async () => { await subscription.close(); return { value: undefined, done: true }; }, + }; + } + + /** Current send/receive accounting. Observer streams never release credit. */ + get flowControl(): FlowControlStatus { return this.windowChannel?.status ?? ordinaryStatus; } + /** @internal */ + get isClosed(): boolean { return this.detached; } + flowControlChanges(): AsyncIterableIterator { + if (this.windowChannel) return this.windowChannel.changes.reader(); + const ordinary = new AsyncBroadcastQueue(1); + ordinary.close(); + return ordinary.reader(); + } + waitWritable(options: { signal?: AbortSignal } = {}): Promise { + return this.windowChannel?.waitWritable(options.signal) ?? Promise.resolve(); + } +} + +const ordinaryStatus: FlowControlStatus = Object.freeze({ + mode: 'ordinary', outstandingSendBytes: 0, bufferedSendBytes: 0, bufferedReceiveBytes: 0, writable: true, +}); +function eventNotification(event: DeliveryEvent): JsonRpcNotification { + if (event.type === 'notification') return event.notification; + const methods = { action: 'action', sessionAdded: 'root/sessionAdded', sessionRemoved: 'root/sessionRemoved', + sessionSummaryChanged: 'root/sessionSummaryChanged', authRequired: 'auth/required', snapshot: 'channel/snapshot', + tcpData: 'tcp/data', tcpEof: 'tcp/eof', tcpReset: 'tcp/reset' }; + return { jsonrpc: '2.0', method: methods[event.type], params: event.params }; } /** Result of {@link AhpClient.dispatch}. */ @@ -222,6 +292,11 @@ interface PendingRequest { reject(error: Error): void; method: string; timer: ReturnType | null; + onResult?: (result: unknown) => void; + channels?: readonly URI[]; + channel?: WindowChannel; + cleanup?: () => void; + subscribe?: SubscribeParams; } // ─── Client ────────────────────────────────────────────────────────────────── @@ -235,6 +310,11 @@ interface PendingRequest { * down. */ export class AhpClient { + private readonly windowEngine: WindowEngine; + private readonly windowOwners = new Map>(); + private readonly subscribePreferences = new WeakMap(); + private readonly unsubscribedWindowed = new Set(); + private readonly lateRequests = new Map(); private readonly transport: AhpTransport; private readonly requestTimeoutMs: number; private readonly subscriptionBuffer: number; @@ -248,6 +328,7 @@ export class AhpClient { private nextClientSeq = 1; private state: ConnectionState = { status: 'idle' }; private receiveLoop: Promise | null = null; + private tornDown = false; private serverRequestHandler: ServerRequestHandler | null = null; constructor(transport: AhpTransport, config: AhpClientConfig = {}) { @@ -258,6 +339,11 @@ export class AhpClient { // Clamp to >= 1 so the queue can deliver at least one value to // already-parked readers without immediately being trimmed. this.subscriptionBuffer = buffer >= 1 ? Math.floor(buffer) : 1; + this.windowEngine = new WindowEngine(transport, config.windowing ?? {}, this.subscriptionBuffer, + error => { + this.tearDown({ type: 'transport', error }); + void Promise.resolve(this.transport.close()).catch(noop); + }, notification => this.handleNotification(notification), this); } /** Current connection state. */ @@ -320,7 +406,7 @@ export class AhpClient { * subscription and event streams. */ async shutdown(): Promise { - if (this.state.status === 'closing' || this.state.status === 'closed') return; + if (this.tornDown) return; this.setState({ status: 'closing' }); // Tear down first so pending requests reject with ClientClosedError // rather than racing against the receive loop seeing the transport @@ -369,14 +455,17 @@ export class AhpClient { clientId: string; lastSeenServerSeq: number; subscriptions: readonly URI[]; + windows?: ChannelDeliveryResumeOptions; + signal?: AbortSignal; }): Promise { const params: ReconnectParams = { channel: 'ahp-root://', clientId: args.clientId, lastSeenServerSeq: args.lastSeenServerSeq, subscriptions: [...args.subscriptions], + ...(args.windows ? { windows: args.windows } : {}), }; - return this.request('reconnect', params); + return this.request('reconnect', params, { signal: args.signal }); } /** @@ -385,25 +474,130 @@ export class AhpClient { * before the `subscribe` request is sent, so no events delivered during * the round-trip are missed. */ + subscribe(uri: URI, options: SubscribeOptions & { windowed: true | WindowOptions }): + Promise<{ result: SubscribeResult; subscription: Subscription }>; + subscribe(uri: URI, options?: SubscribeOptions): + Promise<{ result: SubscribeResult; subscription: Subscription }>; async subscribe( uri: URI, options: SubscribeOptions = {}, - ): Promise<{ result: SubscribeResult; subscription: Subscription }> { - const subscription = this.attachSubscription(uri); + ): Promise<{ result: SubscribeResult; subscription: Subscription }> { + const windowed = !!options.windowed; + const windowOptions = typeof options.windowed === 'object' ? { ...options.windowed } : {}; + const existing = this.windowEngine.channels.get(uri); + const subscription = windowed || existing + ? this.windowOwner(this.windowEngine.ensure(uri, windowOptions)) + : this.attachSubscription(uri); try { + const receive = windowed ? this.windowEngine.offer(windowOptions) : undefined; + if (uri.startsWith('ahp-tcp:') && (!windowed || !this.transport.supportsWindowedDelivery)) { + throw new TypeError('TCP requires a raw-text transport and windowed subscription'); + } const params: SubscribeParams = { channel: uri, ...(options.delivery ? { delivery: options.delivery } : {}), ...(options.view ? { view: options.view } : {}), + ...(receive && this.transport.supportsWindowedDelivery ? { flowControl: { receive } } : {}), }; - const result = await this.request('subscribe', params); + this.subscribePreferences.set(params, windowOptions); + if (windowed && !this.transport.supportsWindowedDelivery) { + subscription.windowChannel!.fallbackReason = 'transport'; + } + const result = await this.request('subscribe', params, { signal: options.signal }); + if (subscription.windowChannel) subscription.windowChannel.options = windowOptions; return { result, subscription }; } catch (err) { - await subscription.close(); + if (!existing && !subscription.windowChannel) { + await subscription.close(); + } else if (!existing && subscription.windowChannel && !subscription.windowChannel.hasInstalled + && ![...this.pending.values()].some(request => request.channel === subscription.windowChannel)) { + subscription.windowChannel.close(err as Error, false); + } throw err; } } + private windowOwner(channel: WindowChannel): Subscription { + let owner = channel.owner ?? this.windowOwners.get(channel.uri); + if (!owner || owner.isClosed || owner.windowChannel !== channel) { + if (owner?.isClosed && owner.windowChannel === channel) channel.prepareNewOwner(); + const empty = new AsyncBroadcastQueue(); + owner = new Subscription(channel.uri, empty.reader(), channel); + channel.owner = owner; + this.windowOwners.set(channel.uri, owner); + } + return owner; + } + + /** Borrow the reliable owner; re-subscribing returns this same handle. */ + windowedSubscription(uri: URI): Subscription | undefined { + const channel = this.windowEngine.channels.get(uri); + return channel ? this.windowOwner(channel) : undefined; + } + + flowControl(uri: URI): FlowControlStatus { return this.windowEngine.channels.get(uri)?.status ?? ordinaryStatus; } + flowControlEvents(): AsyncIterableIterator { return this.windowEngine.events.reader(); } + /** Extended observe-only stream; reliable consumption belongs to the subscription owner. */ + deliveryEvents(): AsyncIterableIterator { return this.windowEngine.observations.reader(); } + waitWritable(uri: URI, options: { signal?: AbortSignal } = {}): Promise { + this.assertOpen(); + return this.windowEngine.channels.get(uri)?.waitWritable(options.signal) ?? Promise.resolve(); + } + /** @internal Host-owned queues survive a connection, but its sends and fragments do not. */ + retainWindowedSubscriptions(): Map { return this.windowEngine.retainChannels(); } + /** @internal */ + restoreWindowedSubscriptions(channels: ReadonlyMap): void { + this.windowEngine.restore(channels); + } + /** @internal */ + windowedResumeOptions(): ChannelDeliveryResumeOptions { return this.windowEngine.resumeOptions(); } + + /** @internal Complete a channel lifecycle without retaining dead scheduler entries. */ + retireWindowedSubscription(channel: WindowChannel, error?: Error, notify = true): void { + if (this.windowOwners.get(channel.uri)?.windowChannel === channel) this.windowOwners.delete(channel.uri); + this.ignoreWindowedChannel(channel.uri); + let hadPending = false; + for (const [id, pending] of this.pending) { + if (pending.channel !== channel) continue; + hadPending = true; + this.abandonRequest(id, error ?? new ClientClosedError('subscription closed'), false); + } + if (notify && !this.isClosed() && (channel.hasInstalled || hadPending)) { + this.sendMessage({ jsonrpc: '2.0', method: 'unsubscribe', params: { channel: channel.uri } }).catch(noop); + } + } + + private ignoreWindowedChannel(uri: URI): void { + this.unsubscribedWindowed.delete(uri); + this.unsubscribedWindowed.add(uri); + if (this.unsubscribedWindowed.size > 4096) { + this.unsubscribedWindowed.delete(this.unsubscribedWindowed.values().next().value!); + } + } + + private abandonRequest(id: number, error: Error, notify = true): void { + const pending = this.pending.get(id); + if (!pending) return; + this.pending.delete(id); + if (pending.timer) clearTimeout(pending.timer); + pending.cleanup?.(); + const unsent = this.windowEngine.cancelRequest(id, error); + if (pending.channels && !unsent) { + this.lateRequests.set(id, { channels: pending.channels, subscribe: pending.subscribe }); + if (this.lateRequests.size > 4096) this.lateRequests.delete(this.lateRequests.keys().next().value!); + } + for (const uri of pending.channels ?? []) { + this.subscriptions.get(uri)?.close(); + this.subscriptions.delete(uri); + this.ignoreWindowedChannel(uri); + this.windowEngine.channels.get(uri)?.close(error, false); + if (notify && !unsent && !this.isClosed()) { + this.sendMessage({ jsonrpc: '2.0', method: 'unsubscribe', params: { channel: uri } }).catch(noop); + } + } + pending.reject(error); + } + /** * Attach a new local {@link Subscription} without sending a `subscribe` * request — use this when the URI was included in `initialSubscriptions` @@ -435,6 +629,7 @@ export class AhpClient { this.subscriptions.delete(uri); queue.close(); } + this.forgetWindowedSubscription(uri); const params: UnsubscribeParams = { channel: uri }; this.notify('unsubscribe', params); return Promise.resolve(); @@ -458,11 +653,21 @@ export class AhpClient { } else { this.nextClientSeq = seq + 1; } + const params: DispatchActionParams = { channel, clientSeq: seq, action }; this.notify('dispatchAction', params); return { clientSeq: seq }; } + async dispatchAsync(channel: URI, action: StateAction, + options: { clientSeq?: number; signal?: AbortSignal } = {}): Promise { + this.assertOpen(); + const seq = options.clientSeq ?? this.nextClientSeq; + this.nextClientSeq = Math.max(this.nextClientSeq, seq + 1); + await this.notifyAsync('dispatchAction', { channel, clientSeq: seq, action }, options); + return { clientSeq: seq }; + } + /** * Protocol-level liveness ping. Useful in browsers, which cannot send * WebSocket ping frames directly. @@ -593,7 +798,31 @@ export class AhpClient { async request( method: M, params: CommandMap[M]['params'], + options: { signal?: AbortSignal } = {}, ): Promise { + this.assertOpen(); + if (options.signal?.aborted) throw options.signal.reason ?? new DOMException('Aborted', 'AbortError'); + let pausedChannel: WindowChannel | undefined; + if (method === 'initialize' && (params as InitializeParams).initialSubscriptions?.some(uri => uri.startsWith('ahp-tcp:'))) { + throw new TypeError('TCP cannot be an initial subscription'); + } + if (method === 'subscribe' && (params as SubscribeParams).channel.startsWith('ahp-tcp:') + && !(params as SubscribeParams).flowControl) throw new TypeError('TCP requires a windowed subscription'); + if (method === 'reconnect' && (params as ReconnectParams).subscriptions.some(uri => uri.startsWith('ahp-tcp:'))) { + throw new TypeError('TCP subscriptions cannot reconnect'); + } + if (method === 'subscribe' && !(params as SubscribeParams).flowControl) { + pausedChannel = this.windowEngine.channels.get((params as SubscribeParams).channel); + if (pausedChannel?.limits) { + pausedChannel.pauseNewMessages = true; + pausedChannel.changed(); + try { await pausedChannel.waitForFragment(); } + catch (error) { + pausedChannel.pauseNewMessages = false; pausedChannel.changed(); this.windowEngine.schedule(); + throw error; + } + } + } this.assertOpen(); const id = this.nextRequestId++; const msg: JsonRpcRequest = { @@ -602,28 +831,87 @@ export class AhpClient { method: method as string, params, }; + let onResult: PendingRequest['onResult']; + let channels: readonly URI[] | undefined; + let requestChannel: WindowChannel | undefined; + if (method === 'subscribe') { + const p = params as SubscribeParams; + if (p.flowControl) { + if (!this.transport.supportsWindowedDelivery) throw new TypeError('transport must retain raw text for windowing'); + this.windowEngine.offer(p.flowControl.receive); + this.windowEngine.ensure(p.channel); + } + channels = [p.channel]; + requestChannel = this.windowEngine.channels.get(p.channel); + onResult = result => { + this.unsubscribedWindowed.delete(p.channel); + const channel = this.windowEngine.channels.get(p.channel); + const preferences = this.subscribePreferences.get(p); + if (channel && preferences) channel.options = preferences; + this.windowEngine.installSubscribe(p, result as SubscribeResult, + channel?.fallbackReason === 'transport'); + }; + } else if (method === 'reconnect') { + const p = params as ReconnectParams; + channels = [...p.subscriptions]; + if (p.windows) { + if (!this.transport.supportsWindowedDelivery) throw new TypeError('transport must retain raw text for reconnect windows'); + const channels = new Set(); + for (const item of p.windows.items) { + if (channels.has(item.channel) || !p.subscriptions.includes(item.channel) || item.channel.startsWith('ahp-tcp:')) { + throw new TypeError('invalid reconnect window channel'); + } + channels.add(item.channel); + this.windowEngine.offer(item.flowControl.receive); + } + } + onResult = result => { + const resumed = result as ReconnectResult; + if (p.windows) this.windowEngine.installReconnect(p.windows, resumed); + else if (resumed.windows !== undefined) throw new TransportError('protocol', 'unsolicited reconnect windows'); + this.windowEngine.installOrdinaryReconnect(p, resumed); + }; + } return new Promise((resolve, reject) => { const pending: PendingRequest = { resolve: value => resolve(value as CommandMap[M]['result']), - reject, + reject: error => { + if (pausedChannel) { + pausedChannel.pauseNewMessages = false; + pausedChannel.changed(); this.windowEngine.schedule(); + } + if (requestChannel && !requestChannel.hasInstalled + && ![...this.pending.values()].some(request => request.channel === requestChannel)) { + requestChannel.close(error, false); + } + reject(error); + }, method: method as string, timer: null, + onResult, + channels, + channel: requestChannel, + subscribe: method === 'subscribe' ? params as SubscribeParams : undefined, }; this.pending.set(id, pending); if (this.requestTimeoutMs > 0) { pending.timer = setTimeout(() => { - if (this.pending.delete(id)) { - reject(new RpcTimeoutError(method as string, this.requestTimeoutMs)); - } + this.abandonRequest(id, new RpcTimeoutError(method as string, this.requestTimeoutMs)); }, this.requestTimeoutMs); } + if (options.signal) { + const abort = () => this.abandonRequest(id, options.signal!.reason ?? new DOMException('Aborted', 'AbortError')); + options.signal.addEventListener('abort', abort, { once: true }); + pending.cleanup = () => options.signal!.removeEventListener('abort', abort); + } this.sendMessage(msg).catch(err => { if (this.pending.delete(id)) { if (pending.timer) clearTimeout(pending.timer); - reject(err); + pending.cleanup?.(); + pending.reject(err); } }); }); @@ -642,6 +930,7 @@ export class AhpClient { params: ClientNotificationMap[M]['params'], ): void { this.assertOpen(); + if (method === 'unsubscribe') this.forgetWindowedSubscription((params as UnsubscribeParams).channel); const msg: JsonRpcNotification = { jsonrpc: '2.0', method: method as string, @@ -653,26 +942,35 @@ export class AhpClient { this.sendMessage(msg).catch(noop); } + async notifyAsync( + method: M, params: ClientNotificationMap[M]['params'], options: { signal?: AbortSignal } = {}, + ): Promise { + this.assertOpen(); + if (method === 'unsubscribe') this.forgetWindowedSubscription((params as UnsubscribeParams).channel); + await this.windowEngine.send({ jsonrpc: '2.0', method: method as string, params }, options.signal); + } + // ─── Internals ───────────────────────────────────────────────────────────── private isClosed(): boolean { return this.state.status === 'closing' || this.state.status === 'closed'; } + private forgetWindowedSubscription(uri: URI): void { + for (const [id, pending] of this.pending) { + if (pending.channels?.includes(uri)) this.abandonRequest(id, new ClientClosedError('subscription unsubscribed'), false); + } + const channel = this.windowEngine.channels.get(uri); + channel?.close(undefined, false); + this.ignoreWindowedChannel(uri); + } + private assertOpen(): void { if (this.isClosed()) throw new ClientClosedError(); } - private async sendMessage(msg: JsonRpcMessage): Promise { - try { - await this.transport.send(msg); - } catch (err) { - const te = err instanceof TransportError - ? err - : new TransportError('io', `transport send failed: ${(err as Error).message}`, { cause: err }); - this.tearDown({ type: 'transport', error: te }); - throw te; - } + private sendMessage(msg: JsonRpcMessage): Promise { + return this.windowEngine.send(msg); } private setState(next: ConnectionState): void { @@ -681,13 +979,19 @@ export class AhpClient { } private tearDown(reason: { type: 'shutdown' } | { type: 'transport'; error: TransportError }): void { - if (this.state.status === 'closed') return; + if (this.tornDown) return; + this.tornDown = true; this.setState({ status: 'closed', reason }); // Fail every pending request. const failure = reason.type === 'shutdown' ? new ClientClosedError() : reason.error; + this.windowEngine.stop(failure); + if (reason.type === 'transport') { + try { void Promise.resolve(this.transport.close()).catch(noop); } catch { /* best effort */ } + } for (const [id, pending] of this.pending) { if (pending.timer) clearTimeout(pending.timer); + pending.cleanup?.(); try { pending.reject(failure); } catch { @@ -698,6 +1002,8 @@ export class AhpClient { for (const queue of this.subscriptions.values()) queue.close(); this.subscriptions.clear(); + this.lateRequests.clear(); + this.unsubscribedWindowed.clear(); this.allEvents.close(); this.stateQueue.close(); } @@ -706,6 +1012,7 @@ export class AhpClient { try { while (this.state.status === 'connected') { const frame = await this.transport.recv(); + if (this.state.status !== 'connected') return; if (frame === null) { this.tearDown({ type: 'transport', error: new TransportError('closed', 'transport closed') }); return; @@ -722,21 +1029,28 @@ export class AhpClient { private handleFrame(frame: TransportFrame): void { let message: ProtocolMessage; + let rawText: string | undefined; try { switch (frame.kind) { case 'parsed': message = frame.message; + rawText = frame.rawText; break; case 'text': + rawText = frame.text; message = decodeMessage(frame.text); break; case 'binary': { const text = new TextDecoder('utf-8').decode(frame.data); + rawText = text; message = decodeMessage(text); break; } } } catch (err) { + if ([...this.windowEngine.channels.values()].some(channel => channel.limits)) { + throw new TransportError('protocol', 'malformed JSON on windowed connection', { cause: err }); + } // A single malformed frame doesn't tear down the channel — a // well-behaved server should not send them, and a transient bad // frame from a peer that recovers shouldn't kill in-flight @@ -747,10 +1061,10 @@ export class AhpClient { return; } - this.dispatchInbound(message); + this.dispatchInbound(message, rawText === undefined ? undefined : rawText.length * 2); } - private dispatchInbound(message: ProtocolMessage): void { + private dispatchInbound(message: ProtocolMessage, size?: number): void { // Response: { id, result } or { id, error } if ('id' in message && 'result' in message) { const success = message as JsonRpcSuccessResponse; @@ -758,7 +1072,28 @@ export class AhpClient { if (pending) { this.pending.delete(success.id); if (pending.timer) clearTimeout(pending.timer); - pending.resolve(success.result); + pending.cleanup?.(); + try { pending.onResult?.(success.result); pending.resolve(success.result); } + catch (error) { pending.reject(error as Error); throw error; } + } else { + const late = this.lateRequests.get(success.id); + this.lateRequests.delete(success.id); + for (const channel of late?.channels ?? []) { + const replacement = [...this.pending.values()].some(request => request.method === 'subscribe' + && request.channels?.includes(channel)); + const result = success.result as SubscribeResult; + if (late?.subscribe && replacement) { + if (result.flowControl) { + const current = this.windowEngine.ensure(channel); + current.install(result, late.subscribe.flowControl?.receive); + this.unsubscribedWindowed.delete(channel); + } + continue; + } + if (late?.subscribe && !result.flowControl && this.subscriptions.has(channel) + && !this.windowEngine.channels.get(channel)?.limits) continue; + void this.unsubscribe(channel); + } } return; } @@ -768,8 +1103,10 @@ export class AhpClient { if (pending) { this.pending.delete(failure.id); if (pending.timer) clearTimeout(pending.timer); + pending.cleanup?.(); pending.reject(new RpcError(failure.error.code, failure.error.message, failure.error.data)); } + this.lateRequests.delete(failure.id); return; } // Server-initiated request: { id, method } @@ -779,7 +1116,19 @@ export class AhpClient { } // Notification: { method } no id if ('method' in message) { - this.handleNotification(message as JsonRpcNotification); + const n = message as JsonRpcNotification; + const channel = (n.params as { channel?: URI })?.channel; + if (channel && this.unsubscribedWindowed.has(channel)) return; + const controller = channel && this.windowEngine.channels.get(channel); + if (controller) controller.receive(n, size); + else if (n.method === 'channel/frame' || n.method === 'channel/credit' || n.method === 'channel/snapshot') { + if (channel && !this.subscriptions.has(channel)) return; + throw new TransportError('protocol', 'windowed notification has no accepted subscription'); + } + if (!controller) { + if (channel) this.windowEngine.observations.publish({ channel, event: deliveryEvent(n) }); + this.handleNotification(n); + } return; } } diff --git a/clients/typescript/src/client/error.ts b/clients/typescript/src/client/error.ts index d0cc2b509..dffdc70a6 100644 --- a/clients/typescript/src/client/error.ts +++ b/clients/typescript/src/client/error.ts @@ -86,4 +86,13 @@ export class ClientClosedError extends AhpClientError { super(message); this.name = 'ClientClosedError'; } + +} + +/** A windowed subscription's bounded unsent queue is full. */ +export class BackpressureError extends AhpClientError { + constructor(readonly channel: string, readonly maximumBufferedBytes: number) { + super(`Send queue for "${channel}" exceeds ${maximumBufferedBytes} UTF-16 bytes; await waitWritable()`); + this.name = 'BackpressureError'; + } } diff --git a/clients/typescript/src/client/events.ts b/clients/typescript/src/client/events.ts index d005a8b39..03aa03226 100644 --- a/clients/typescript/src/client/events.ts +++ b/clients/typescript/src/client/events.ts @@ -17,7 +17,9 @@ import type { SessionRemovedParams, SessionSummaryChangedParams, } from '../types/channels-root/notifications.js'; -import type { AuthRequiredParams } from '../types/common/notifications.js'; +import type { AuthRequiredParams, ChannelSnapshotParams } from '../types/common/notifications.js'; +import type { TcpDataParams, TcpEofParams, TcpResetParams } from '../types/channels-tcp/notifications.js'; +import type { JsonRpcNotification } from '../types/common/messages.js'; import type { URI } from '../types/common/state.js'; import type { TransportError } from './error.js'; @@ -29,6 +31,45 @@ export type SubscriptionEvent = | { readonly type: 'sessionSummaryChanged'; readonly params: SessionSummaryChangedParams } | { readonly type: 'authRequired'; readonly params: AuthRequiredParams }; +/** Extended delivery surface, leaving legacy exhaustive SubscriptionEvent switches intact. */ +export type DeliveryEvent = SubscriptionEvent + | { readonly type: 'snapshot'; readonly params: ChannelSnapshotParams } + | { readonly type: 'tcpData'; readonly params: TcpDataParams } + | { readonly type: 'tcpEof'; readonly params: TcpEofParams } + | { readonly type: 'tcpReset'; readonly params: TcpResetParams } + | { readonly type: 'notification'; readonly notification: JsonRpcNotification }; + +export interface ReceivedDelivery { + readonly event: DeliveryEvent; + readonly notification: JsonRpcNotification; + /** Sum of original wire sizes, including fragments already credited when stored. */ + readonly bytes: number; + /** Release remaining credit; idempotent and safe after reconnect. */ + release(): void; +} + +/** Observe-only extended fan-in. It never owns a receipt or releases credit. */ +export interface ClientDeliveryEvent { + readonly channel: URI; + readonly event: DeliveryEvent; +} + +/** @internal */ +export function deliveryEvent(n: JsonRpcNotification): DeliveryEvent { + switch (n.method) { + case 'action': return { type: 'action', params: n.params as ActionEnvelope }; + case 'root/sessionAdded': return { type: 'sessionAdded', params: n.params as SessionAddedParams }; + case 'root/sessionRemoved': return { type: 'sessionRemoved', params: n.params as SessionRemovedParams }; + case 'root/sessionSummaryChanged': return { type: 'sessionSummaryChanged', params: n.params as SessionSummaryChangedParams }; + case 'auth/required': return { type: 'authRequired', params: n.params as AuthRequiredParams }; + case 'channel/snapshot': return { type: 'snapshot', params: n.params as ChannelSnapshotParams }; + case 'tcp/data': return { type: 'tcpData', params: n.params as TcpDataParams }; + case 'tcp/eof': return { type: 'tcpEof', params: n.params as TcpEofParams }; + case 'tcp/reset': return { type: 'tcpReset', params: n.params as TcpResetParams }; + default: return { type: 'notification', notification: n }; + } +} + /** * A single event delivered to the top-level {@link AhpClient.events} stream, * tagged with the channel URI it was scoped to. diff --git a/clients/typescript/src/client/flow-control.ts b/clients/typescript/src/client/flow-control.ts new file mode 100644 index 000000000..9003c9aad --- /dev/null +++ b/clients/typescript/src/client/flow-control.ts @@ -0,0 +1,1028 @@ +import type { ChannelFlowControl, ChannelReceiveLimits, URI } from '../types/common/state.js'; +import type { ChannelDeliveryResumeOptions, SubscribeParams, SubscribeResult, ReconnectResult, ReconnectParams } from '../types/common/commands.js'; +import type { JsonRpcNotification } from '../types/common/messages.js'; +import { countsAgainstWindow } from '../types/window-exempt.generated.js'; +import { AsyncBroadcastQueue } from './async-queue.js'; +import { BackpressureError, ClientClosedError, TransportError } from './error.js'; +import type { ClientDeliveryEvent, ReceivedDelivery } from './events.js'; +import type { AhpClient, Subscription } from './client.js'; +import type { DeliveryEvent } from './events.js'; +import { deliveryEvent } from './events.js'; +import type { AhpTransport, JsonRpcMessage } from './transport.js'; + +export interface WindowOptions { + windowBytes?: number; + maximumFrameBytes?: number; + receive?: 'auto' | 'manual'; + coalescing?: boolean; +} + +export interface CoalescerContext { + readonly channel: URI; + readonly direction: 'receive' | 'send'; +} + +/** Custom coalescers run before built-ins, on adjacent, still queued messages. */ +export interface ChannelCoalescer { + readonly method: string; + merge(earlier: JsonRpcNotification, later: JsonRpcNotification, + context: CoalescerContext): JsonRpcNotification | undefined; +} + +export interface WindowingConfig { + /** Default receive window, in UTF-16 bytes. Default 262144. */ + windowBytes?: number; + /** Defaults to the transport's limit, or 1048576 if unknown. */ + maximumFrameBytes?: number; + /** Limit for data queued behind the head send, in UTF-16 bytes. Default 1048576. */ + maxBufferedSendBytes?: number; + coalescers?: readonly ChannelCoalescer[]; +} + +export interface FlowControlStatus { + readonly mode: 'ordinary' | 'windowed'; + readonly fallbackReason?: 'declined' | 'transport'; + readonly limits?: ChannelFlowControl; + readonly outstandingSendBytes: number; + readonly bufferedSendBytes: number; + readonly bufferedReceiveBytes: number; + readonly writable: boolean; +} + +export interface FlowControlEvent { + readonly channel: URI; + readonly type: 'blocked' | 'unblocked' | 'creditSent' | 'creditReceived' + | 'fragmentedMessageReceived' | 'fragmentedMessageSent' | 'coalesced' + | 'resumed' | 'dropped'; + readonly status: FlowControlStatus; + readonly bytes?: number; +} + +export const DEFAULT_MAXIMUM_FRAME_BYTES = 1_048_576; +const MAX_CONTROL_BYTES = 1_048_576; +const CONTROL_INTERVAL_MS = 10; +const CONTROL_BURST = 256; +export function utf16Bytes(text: string): number { return text.length * 2; } + +function immutableLimits(limits: ChannelFlowControl): ChannelFlowControl { + return Object.freeze({ clientReceive: Object.freeze({ ...limits.clientReceive }), + hostReceive: Object.freeze({ ...limits.hostReceive }) }); +} + +function protocol(message: string): never { throw new TransportError('protocol', message); } +function validLimits(limits: ChannelReceiveLimits): void { + if (!limits || !Number.isSafeInteger(limits.windowBytes) || limits.windowBytes < 1 + || !Number.isSafeInteger(limits.maximumFrameBytes) || limits.maximumFrameBytes < 1024) { + protocol('invalid channel flow-control limits'); + } +} + +function params(n: JsonRpcNotification): Record { + return n.params as Record; +} + +function sameMetadata(a: Record, b: Record, omitted: string[]): boolean { + const rest = (v: Record) => Object.fromEntries( + Object.entries(v).filter(([k]) => !omitted.includes(k))); + const canonical = (v: unknown): string => { + if (Array.isArray(v)) return `[${v.map(canonical).join(',')}]`; + if (v && typeof v === 'object') return `{${Object.entries(v).sort(([a], [b]) => a.localeCompare(b)) + .map(([k, x]) => `${JSON.stringify(k)}:${canonical(x)}`).join(',')}}`; + return JSON.stringify(v) ?? 'undefined'; + }; + return canonical(rest(a)) === canonical(rest(b)); +} + +function decodeBase64(data: string): Uint8Array { + if (!/^(?:[A-Za-z0-9+/]{4})*(?:[A-Za-z0-9+/]{2}==|[A-Za-z0-9+/]{3}=)?$/.test(data) || !data) { + protocol('invalid tcp/data base64'); + } + const binary = atob(data); + const bytes = Uint8Array.from(binary, c => c.charCodeAt(0)); + if (encodeBase64(bytes) !== data) protocol('noncanonical tcp/data base64'); + return bytes; +} + +function encodeBase64(bytes: Uint8Array): string { + let binary = ''; + for (const byte of bytes) binary += String.fromCharCode(byte); + return btoa(binary); +} + +function mergeNotifications(a: JsonRpcNotification, b: JsonRpcNotification, + context: CoalescerContext, customs: readonly ChannelCoalescer[]): JsonRpcNotification | undefined { + if (a.method !== b.method || !countsAgainstWindow(a.method) + || ['channel/frame', 'channel/snapshot', 'tcp/eof', 'dispatchAction'].includes(a.method)) return; + const ap = params(a), bp = params(b); + if (ap.channel !== bp.channel || ap.origin !== undefined || bp.origin !== undefined + || ap.rejectionReason !== undefined || bp.rejectionReason !== undefined) return; + for (const custom of customs) { + if (custom.method !== a.method) continue; + const merged = custom.merge(a, b, context); + if (merged) { + if (merged.method !== a.method || params(merged).channel !== ap.channel + || !countsAgainstWindow(merged.method)) protocol('coalescer changed method or channel'); + if (context.direction === 'receive' && a.method === 'action') { + return { ...merged, params: { ...params(merged), serverSeq: bp.serverSeq } }; + } + return merged; + } + } + if (context.direction === 'send' && a.method === 'tcp/data') { + const first = decodeBase64(ap.data as string), second = decodeBase64(bp.data as string); + const bytes = new Uint8Array(first.length + second.length); + bytes.set(first); bytes.set(second, first.length); + return { ...b, params: { ...bp, data: encodeBase64(bytes) } }; + } + if (context.direction !== 'receive' || a.method !== 'action') return; + const aa = ap.action as Record, ba = bp.action as Record; + if (!aa || !ba || aa.type !== ba.type + || !sameMetadata(ap, bp, ['action', 'serverSeq'])) return; + const field = aa.type === 'terminal/data' ? 'data' : aa.type === 'chat/delta' ? 'content' : undefined; + if (!field || typeof aa[field] !== 'string' || typeof ba[field] !== 'string' + || !sameMetadata(aa, ba, [field])) return; + return { ...b, params: { ...bp, action: { ...ba, [field]: aa[field] + ba[field] } } }; +} + +interface Entry { + notification: JsonRpcNotification; + bytes: number; + debit: number; + epoch: number; +} +interface Send { + notification: JsonRpcNotification; + serialized: string; + text?: string; + offset: number; + bytes: number; + completions: Array<{ resolve(): void; reject(e: Error): void }>; + cleanup: Array<() => void>; + cancelable: boolean; +} + +/** Persistent owner queue; migrated by the host runtime, never by observer streams. @internal */ +export class WindowChannel { + owner?: Subscription; + readonly rebound = new Set<(client: AhpClient) => void>(); + readonly closed = new Set<(error?: Error) => void>(); + readonly changes = new AsyncBroadcastQueue(1); + options: WindowOptions; + offer?: ChannelReceiveLimits; + limits?: ChannelFlowControl; + fallbackReason?: 'declined' | 'transport'; + epoch = 0; + sent = 0; + remoteConsumed = 0; + received = 0; + released = 0; + advertised = 0; + lastAppliedServerSeq = 0; + expectingSnapshot = false; + connected = true; + pauseNewMessages = false; + fragment = ''; + fragmentBytes = 0; + fragmenting = false; + sends: Send[] = []; + private queue: Entry[] = []; + private waiter?: { resolve(v: IteratorResult): void; reject(e: Error): void }; + private current?: ReceivedDelivery; + private terminal = false; + private error?: Error; + private timer?: ReturnType; + private installed = false; + private ownerHasRead = false; + private waitingForOwnerSnapshot = false; + private lastStatus = ''; + private wasWritable = true; + private readonly retainedDebits = new Map(); + private tcpReset = false; + private tcpReceivedEof = false; + tcpSentEof = false; + engine: WindowEngine; + + constructor(readonly uri: URI, options: WindowOptions, engine: WindowEngine, + private readonly ordinaryBuffer: number) { + this.options = options; + this.engine = engine; + } + + get status(): FlowControlStatus { + return { + mode: this.limits ? 'windowed' : 'ordinary', + ...(this.fallbackReason ? { fallbackReason: this.fallbackReason } : {}), + ...(this.limits ? { limits: this.limits } : {}), + outstandingSendBytes: this.sent - this.remoteConsumed, + bufferedSendBytes: this.sends.reduce((sum, s) => sum + s.bytes, 0), + bufferedReceiveBytes: this.received - this.released + + [...this.retainedDebits.values()].reduce((sum, bytes) => sum + bytes, 0), + writable: this.connected && !this.terminal && !this.tcpReset && !this.tcpSentEof && !this.pauseNewMessages + && (!this.limits || this.sent - this.remoteConsumed < this.limits.hostReceive.windowBytes) + && this.sends.reduce((sum, s) => sum + s.bytes, 0) < this.engine.maxBufferedSendBytes, + }; + } + + get hasInstalled(): boolean { return this.installed; } + + changed(): void { + const status = this.status, key = JSON.stringify(status); + if (key === this.lastStatus) return; + this.lastStatus = key; + this.changes.publish(status); + if (status.writable !== this.wasWritable) { + this.wasWritable = status.writable; + this.engine.emit(this, status.writable ? 'unblocked' : 'blocked'); + } + } + + async waitWritable(signal?: AbortSignal): Promise { + if (signal?.aborted) throw signal.reason ?? new DOMException('Aborted', 'AbortError'); + if (this.status.writable) return; + if (!this.connected || this.terminal || this.tcpReset || this.tcpSentEof) throw this.error ?? new ClientClosedError(); + const reader = this.changes.reader(); + let abort: (() => void) | undefined; + const cancellation = new Promise((_, reject) => { + abort = () => reject(signal?.reason ?? new DOMException('Aborted', 'AbortError')); + signal?.addEventListener('abort', abort, { once: true }); + }); + try { + while (!this.status.writable) { + if (!this.connected || this.terminal || this.tcpReset || this.tcpSentEof) throw this.error ?? new ClientClosedError(); + await Promise.race([reader.next(), cancellation]); + } + } finally { + if (abort) signal?.removeEventListener('abort', abort); + await reader.return?.(); + } + } + + install(result: SubscribeResult, offered?: ChannelReceiveLimits, transportFallback = false): void { + const oldWindowed = !!this.limits; + if (result.flowControl) { + if (!offered || result.snapshot) protocol('unsolicited flow control or snapshot in accepted result'); + this.engine.validateAcceptance(result.flowControl, offered, this.uri); + if (!oldWindowed) this.resetCounters(); + this.limits = immutableLimits(result.flowControl); + this.fallbackReason = undefined; + this.expectingSnapshot = true; + } else { + if (this.uri.startsWith('ahp-tcp:')) protocol('TCP subscription requires windowed delivery'); + if (this.sends.some(send => send.offset > 0)) protocol('host declined while a fragmented send is active'); + this.limits = undefined; + this.fallbackReason = transportFallback ? 'transport' : offered ? 'declined' : undefined; + this.expectingSnapshot = false; + if (oldWindowed) this.resetCounters(); + if (result.snapshot) { + const notification = { jsonrpc: '2.0' as const, method: 'channel/snapshot', + params: { channel: this.uri, snapshot: result.snapshot } }; + this.enqueue({ notification, bytes: 0, debit: 0, epoch: this.epoch }, !this.installed); + this.engine.observe(this.uri, notification); + } + } + this.installed = true; + this.pauseNewMessages = false; + this.connected = true; + this.changed(); + this.engine.schedule(); + } + + resetCounters(): void { + const held = this.received - this.released; + if (held > 0) this.retainedDebits.set(this.epoch, held); + this.epoch++; + this.sent = this.remoteConsumed = this.received = this.released = this.advertised = 0; + this.fragment = ''; this.fragmentBytes = 0; this.fragmenting = false; + if (this.timer) clearTimeout(this.timer); + this.timer = undefined; + } + + async waitForFragment(): Promise { + while (this.sends[0]?.offset) { + const send = this.sends[0]; + await new Promise((resolve, reject) => send.completions.push({ resolve, reject })); + } + } + + receive(notification: JsonRpcNotification, size: number | undefined): void { + if (!this.connected) protocol('channel traffic before reconnect result'); + if (this.limits && !countsAgainstWindow(notification.method)) this.engine.checkControl(size); + if (this.tcpReset && notification.method !== 'tcp/reset') return; + if (notification.method === 'channel/credit') { + if (!this.limits) protocol('credit on ordinary subscription'); + const consumed = params(notification).consumedBytes as number; + if (!Number.isSafeInteger(consumed) || consumed < 0 || consumed > this.sent) { + protocol('channel credit exceeds bytes sent'); + } + if (consumed > this.remoteConsumed) { + this.remoteConsumed = consumed; + this.engine.emit(this, 'creditReceived', consumed); + this.changed(); this.engine.schedule(); + } + return; + } + const counted = !!this.limits && countsAgainstWindow(notification.method); + if (counted) { + if (size === undefined) protocol('windowed transport did not preserve raw JSON text'); + if (size > this.limits!.clientReceive.maximumFrameBytes) protocol('oversized channel data message'); + this.received = this.add(this.received, size); + } + const debit = counted ? size! : 0; + if (notification.method === 'channel/snapshot' && !this.limits) protocol('snapshot notification on ordinary subscription'); + if (notification.method === 'channel/frame') { + if (!this.limits) protocol('fragment on ordinary subscription'); + const p = params(notification); + if (typeof p.data !== 'string' || !p.data || (p.final !== undefined && typeof p.final !== 'boolean')) { + protocol('invalid channel fragment'); + } + const first = p.data.charCodeAt(0), last = p.data.charCodeAt(p.data.length - 1); + if ((first >= 0xdc00 && first <= 0xdfff) || (last >= 0xd800 && last <= 0xdbff)) { + protocol('fragment divides UTF-16 surrogate pair'); + } + this.fragmenting = true; + this.fragment += p.data; + this.fragmentBytes = this.add(this.fragmentBytes, debit); + if (!p.final) { this.releaseBytes(debit, this.epoch); this.changed(); return; } + let joined: JsonRpcNotification; + try { joined = JSON.parse(this.fragment) as JsonRpcNotification; } + catch { protocol('invalid JSON in joined channel fragments'); } + const bytes = this.fragmentBytes; + this.fragment = ''; this.fragmentBytes = 0; this.fragmenting = false; + if (!joined || joined.jsonrpc !== '2.0' || 'id' in joined || typeof joined.method !== 'string' + || !joined.params || params(joined).channel !== this.uri || !countsAgainstWindow(joined.method) + || joined.method === 'channel/frame') protocol('invalid joined notification or channel mismatch'); + this.engine.emit(this, 'fragmentedMessageReceived', bytes); + this.acceptData(joined, bytes, debit); + } else { + if (counted && this.fragmenting) protocol('data interleaved with unfinished fragments'); + if (notification.method === 'tcp/reset') this.resetTcp(); + this.acceptData(notification, debit, debit); + } + this.changed(); + } + + private acceptData(notification: JsonRpcNotification, bytes: number, debit: number, observe = true): void { + if (this.expectingSnapshot && countsAgainstWindow(notification.method)) { + if (notification.method !== 'channel/snapshot') protocol('first channel data must be a snapshot'); + this.expectingSnapshot = false; + } + if (notification.method === 'channel/snapshot') { + const snapshot = params(notification).snapshot as { resource?: string; fromSeq?: number }; + if (!snapshot || snapshot.resource !== this.uri) protocol('snapshot channel mismatch'); + } + if (notification.method === 'tcp/data') { + if (this.tcpReceivedEof) protocol('tcp/data after tcp/eof'); + decodeBase64(params(notification).data as string); + } else if (notification.method === 'tcp/eof') this.tcpReceivedEof = true; + const p = params(notification); + const seq = notification.method === 'channel/snapshot' + ? (p.snapshot as { fromSeq: number }).fromSeq : p.serverSeq; + if ((notification.method === 'action' || notification.method === 'channel/snapshot') + && (!Number.isSafeInteger(seq) || (seq as number) < 0)) protocol('invalid channel sequence number'); + if (typeof seq === 'number') this.lastAppliedServerSeq = Math.max(this.lastAppliedServerSeq, seq); + this.enqueue({ notification, bytes, debit, epoch: this.epoch }); + if (observe) this.engine.observe(this.uri, notification); + } + + seedOrdinary(notification: JsonRpcNotification): void { + this.connected = true; + this.acceptData(notification, 0, 0, false); + this.changed(); + } + + private enqueue(entry: Entry, first = false): void { + if (this.terminal) return; + if (this.waitingForOwnerSnapshot) { + if (entry.notification.method !== 'channel/snapshot') { + this.releaseBytes(entry.debit, entry.epoch); + return; + } + this.waitingForOwnerSnapshot = false; + } + if (first && this.queue.length) this.queue.unshift(entry); + else { + const tail = this.queue[this.queue.length - 1]; + const merged = this.limits && this.options.coalescing !== false && tail?.epoch === entry.epoch + ? mergeNotifications(tail.notification, entry.notification, + { channel: this.uri, direction: 'receive' }, this.engine.coalescers) : undefined; + if (merged) { + tail.notification = merged; tail.bytes += entry.bytes; tail.debit += entry.debit; + this.engine.emit(this, 'coalesced', entry.bytes); + } else this.queue.push(entry); + } + if (!this.limits && this.queue.length > this.ordinaryBuffer) { + const pinned = !this.ownerHasRead && this.queue[0].notification.method === 'channel/snapshot'; + this.queue.splice(pinned ? 1 : 0, 1); + } + if (this.waiter) { + const waiter = this.waiter; this.waiter = undefined; + waiter.resolve({ value: this.take(), done: false }); + } + } + + private take(): ReceivedDelivery { + const entry = this.queue.shift()!; + this.ownerHasRead = true; + let released = false; + const delivery: ReceivedDelivery = { + event: deliveryEvent(entry.notification), notification: entry.notification, bytes: entry.bytes, + release: () => { + if (released) return; + released = true; + this.releaseBytes(entry.debit, entry.epoch); + }, + }; + return delivery; + } + + next(manual = false): Promise> { + if (!manual) this.releaseCurrent(); + const finish = (value: IteratorResult) => { + if (!manual && !value.done) this.current = value.value; + return value; + }; + if (this.queue.length) return Promise.resolve(finish({ value: this.take(), done: false })); + if (this.terminal) { + if (this.error) return Promise.reject(this.error); + return Promise.resolve({ value: undefined, done: true }); + } + if (this.waiter) return Promise.reject(new TypeError('only one owner read may be pending')); + return new Promise>((resolve, reject) => { + this.waiter = { resolve: v => resolve(finish(v)), reject }; + }); + } + + releaseCurrent(): void { this.current?.release(); this.current = undefined; } + + abandonReader(): void { + this.releaseCurrent(); + for (const entry of this.queue) this.releaseBytes(entry.debit, entry.epoch); + this.queue = []; + if (this.waiter) { + const waiter = this.waiter; this.waiter = undefined; + waiter.resolve({ value: undefined, done: true }); + } + } + + prepareNewOwner(): void { + this.abandonReader(); + this.ownerHasRead = false; + this.waitingForOwnerSnapshot = true; + } + + private add(total: number, bytes: number): number { + if (!Number.isSafeInteger(total + bytes)) protocol('channel byte counter overflow'); + return total + bytes; + } + + private releaseBytes(bytes: number, epoch: number): void { + if (!bytes || this.terminal) return; + if (epoch !== this.epoch) { + const held = this.retainedDebits.get(epoch); + if (held !== undefined) { + if (held <= bytes) this.retainedDebits.delete(epoch); + else this.retainedDebits.set(epoch, held - bytes); + this.changed(); + } + return; + } + if (!this.limits) return; + this.released = this.add(this.released, bytes); + this.changed(); + if (!this.connected || this.tcpReset) return; + const threshold = Math.max(1, Math.min(16_384, Math.floor(this.limits.clientReceive.windowBytes / 4))); + if (this.released - this.advertised >= threshold + || this.received - this.advertised >= this.limits.clientReceive.windowBytes) this.flushCredit(); + else if (!this.timer) this.timer = setTimeout(() => this.flushCredit(), 5); + } + + private flushCredit(): void { + if (this.timer) clearTimeout(this.timer); + this.timer = undefined; + if (!this.connected || this.terminal || !this.limits || this.released <= this.advertised) return; + this.advertised = this.released; + this.engine.control({ jsonrpc: '2.0', method: 'channel/credit', + params: { channel: this.uri, consumedBytes: this.advertised } }, this).catch(() => undefined); + this.engine.emit(this, 'creditSent', this.advertised); + } + + disconnect(error: Error, retain: boolean): void { + this.connected = false; + this.fragment = ''; this.fragmentBytes = 0; this.fragmenting = false; + if (this.timer) clearTimeout(this.timer); + this.timer = undefined; + this.cancelSends(error); + if (!retain) this.close(error); + this.changed(); + } + + close(error?: Error, notify = true): void { + if (this.terminal) return; + this.connected = false; + this.terminal = true; + this.epoch++; + this.limits = undefined; + this.fragment = ''; this.fragmentBytes = 0; this.fragmenting = false; + this.retainedDebits.clear(); + if (this.timer) clearTimeout(this.timer); + this.timer = undefined; + this.releaseCurrent(); + if (!error) { + for (const entry of this.queue) this.releaseBytes(entry.debit, entry.epoch); + this.queue = []; + } + this.error = error; + this.cancelSends(error ?? new ClientClosedError()); + if (this.waiter) { + const w = this.waiter; this.waiter = undefined; + if (error) w.reject(error); else w.resolve({ value: undefined, done: true }); + } + this.changed(); + this.changes.close(); + this.engine.retire(this, error, notify); + for (const callback of this.closed) callback(error); + this.closed.clear(); + this.rebound.clear(); + } + + resetTcp(): void { + this.tcpReset = true; + if (this.timer) clearTimeout(this.timer); + this.timer = undefined; + this.current?.release(); + for (const entry of this.queue) this.releaseBytes(entry.debit, entry.epoch); + this.queue = []; + this.fragment = ''; this.fragmentBytes = 0; this.fragmenting = false; + this.cancelSends(new ClientClosedError('TCP reset')); + } + get isTcpReset(): boolean { return this.tcpReset; } + + cancelSends(error: Error): void { + for (const send of this.sends) { + for (const cleanup of send.cleanup) cleanup(); + for (const completion of send.completions) completion.reject(error); + } + this.sends = []; + this.changed(); + } + + get isClosed(): boolean { return this.terminal; } +} + +interface Control { + message: JsonRpcMessage | string; + resolve(): void; + reject(error: Error): void; + exempt: boolean; + creditOwner?: WindowChannel; + started?: boolean; + canceled?: Error; +} + +/** Negotiation, accounting and one fair, serialized writer per connection. @internal */ +export class WindowEngine { + readonly channels = new Map(); + readonly events = new AsyncBroadcastQueue(1024); + readonly observations = new AsyncBroadcastQueue(4096); + readonly coalescers: readonly ChannelCoalescer[]; + readonly maxBufferedSendBytes: number; + private controls: Control[] = []; + private scheduled = false; + private stopped = false; + private cursor = 0; + private retain = false; + private controlPeriod = 0; + private controlCount = 0; + private receivedControlPeriod = 0; + private receivedControlCount = 0; + private readonly sendAbort = new AbortController(); + private currentControl?: Control; + private driving = false; + + constructor(readonly transport: AhpTransport, readonly config: WindowingConfig, + private readonly buffer: number, private readonly fail: (error: TransportError) => void, + private readonly legacyObserve: (notification: JsonRpcNotification) => void, + readonly client: AhpClient) { + this.coalescers = [...(config.coalescers ?? [])]; + this.maxBufferedSendBytes = config.maxBufferedSendBytes ?? 1_048_576; + if (!Number.isSafeInteger(this.maxBufferedSendBytes) || this.maxBufferedSendBytes < 1) { + throw new TypeError('maxBufferedSendBytes must be a positive safe integer'); + } + } + + observe(channel: URI, notification: JsonRpcNotification): void { + this.observations.publish({ channel, event: deliveryEvent(notification) }); + this.legacyObserve(notification); + } + + ensure(uri: URI, options: WindowOptions = {}): WindowChannel { + let channel = this.channels.get(uri); + if (!channel || channel.isClosed) { + channel = new WindowChannel(uri, options, this, this.buffer); + this.channels.set(uri, channel); + } + return channel; + } + + retire(channel: WindowChannel, error?: Error, notify = true): void { + if (this.channels.get(channel.uri) !== channel) return; + this.channels.delete(channel.uri); + this.controls = this.controls.filter(control => { + if (control.creditOwner !== channel) return true; + control.resolve(); return false; + }); + this.client.retireWindowedSubscription(channel, error, notify && !this.stopped); + } + + cancelRequest(id: number, error: Error): boolean { + const matches = (control: Control) => typeof control.message !== 'string' + && 'id' in control.message && control.message.id === id && 'method' in control.message; + const index = this.controls.findIndex(matches); + if (index >= 0) { + this.controls.splice(index, 1)[0].reject(error); + return true; + } + if (this.currentControl && matches(this.currentControl) && !this.currentControl.started) { + this.currentControl.canceled = error; + this.currentControl.reject(error); + return true; + } + return false; + } + + offer(options: WindowOptions): ChannelReceiveLimits { + const limits = { + windowBytes: options.windowBytes ?? this.config.windowBytes ?? 262_144, + maximumFrameBytes: options.maximumFrameBytes ?? this.config.maximumFrameBytes + ?? this.transport.maximumFrameBytes ?? DEFAULT_MAXIMUM_FRAME_BYTES, + }; + validLimits(limits); + return limits; + } + + validateAcceptance(flow: ChannelFlowControl, offer: ChannelReceiveLimits, uri: URI): void { + validLimits(flow.clientReceive); validLimits(flow.hostReceive); + if (flow.clientReceive.windowBytes > offer.windowBytes + || flow.clientReceive.maximumFrameBytes > offer.maximumFrameBytes) protocol('accepted limits exceed offer'); + const smallest = utf16Bytes(JSON.stringify({ jsonrpc: '2.0', method: 'channel/frame', + params: { channel: uri, data: 'x', final: true } })); + if (smallest > flow.clientReceive.maximumFrameBytes || smallest > flow.hostReceive.maximumFrameBytes) { + protocol('channel URI cannot fit in a fragment'); + } + } + + installSubscribe(p: SubscribeParams, result: SubscribeResult, fallback = false): void { + const channel = this.channels.get(p.channel); + if (channel) channel.install(result, p.flowControl?.receive, fallback); + else if (result.flowControl) protocol('flow-control result has no local subscription'); + } + + resumeOptions(): ChannelDeliveryResumeOptions { + return { items: [...this.channels.values()].filter(c => !!c.limits && !c.isClosed + && !c.uri.startsWith('ahp-tcp:')).map(c => ({ + channel: c.uri, flowControl: { receive: this.offer(c.options) }, + lastAppliedServerSeq: c.lastAppliedServerSeq, + })) }; + } + + retainChannels(): Map { + this.retain = true; + return new Map([...this.channels].filter(([uri, c]) => !uri.startsWith('ahp-tcp:') && !c.isClosed)); + } + + restore(channels: ReadonlyMap): void { + for (const [uri, channel] of channels) { + if (channel.connected) throw new TypeError('close the old connection before restoring subscriptions'); + channel.engine = this; + channel.connected = false; + channel.resetCounters(); + this.channels.set(uri, channel); + for (const callback of channel.rebound) callback(this.client); + } + } + + installReconnect(windows: ChannelDeliveryResumeOptions, result: ReconnectResult): void { + const kept = new Map(result.windows?.items.map(i => [i.channel, i]) ?? []); + if (kept.size !== (result.windows?.items.length ?? 0)) protocol('duplicate reconnect window'); + for (const uri of kept.keys()) if (!windows.items.some(i => i.channel === uri)) protocol('unsolicited reconnect window'); + for (const item of windows.items) { + const channel = this.ensure(item.channel); + const resumed = kept.get(item.channel); + if (!resumed) { + channel.disconnect(new ClientClosedError(`subscription dropped on reconnect: ${item.channel}`), false); + this.emit(channel, 'dropped'); continue; + } + this.validateAcceptance(resumed.flowControl, item.flowControl.receive, item.channel); + if (!['snapshot', 'replay', 'live'].includes(resumed.recovery)) protocol('unknown reconnect recovery kind'); + channel.resetCounters(); channel.limits = immutableLimits(resumed.flowControl); + channel.expectingSnapshot = resumed.recovery === 'snapshot'; + channel.connected = true; channel.changed(); + this.emit(channel, 'resumed'); + } + this.schedule(); + } + + installOrdinaryReconnect(request: ReconnectParams, result: ReconnectResult): void { + for (const uri of request.subscriptions) { + const channel = this.channels.get(uri); + if (channel && !channel.limits && !channel.isClosed) { + channel.connected = true; + channel.changed(); + } + } + if (result.type === 'snapshot') { + for (const snapshot of result.snapshots) { + const channel = this.channels.get(snapshot.resource); + if (channel && !channel.limits && !channel.isClosed) channel.seedOrdinary({ + jsonrpc: '2.0', method: 'channel/snapshot', params: { channel: snapshot.resource, snapshot }, + }); + } + } else { + for (const uri of result.missing) { + const channel = this.channels.get(uri); + if (channel && !channel.limits) { + channel.disconnect(new ClientClosedError(`subscription dropped on reconnect: ${uri}`), false); + this.emit(channel, 'dropped'); + } + } + for (const action of result.actions) { + const channel = this.channels.get(action.channel); + if (channel && !channel.limits && !channel.isClosed) channel.seedOrdinary({ + jsonrpc: '2.0', method: 'action', params: action, + }); + } + } + } + + emit(channel: WindowChannel, type: FlowControlEvent['type'], bytes?: number): void { + this.events.publish({ channel: channel.uri, type, status: channel.status, + ...(bytes !== undefined ? { bytes } : {}) }); + } + + send(message: JsonRpcMessage, signal?: AbortSignal): Promise { + if ('id' in message || !('method' in message)) return this.control(message); + const notification = message as JsonRpcNotification; + const uri = params(notification)?.channel as URI; + const channel = this.channels.get(uri); + if (notification.method === 'channel/frame') throw new TypeError('channel/frame is managed by the client'); + if (notification.method === 'channel/credit') throw new TypeError('channel/credit is managed by delivery.release()'); + if (channel && notification.method === 'tcp/reset') channel.resetTcp(); + if (['tcp/data', 'tcp/eof'].includes(notification.method)) { + if (!channel?.limits) throw new TypeError('TCP sends require an accepted windowed subscription'); + if (channel.isTcpReset || channel.tcpSentEof) throw new ClientClosedError('TCP send direction closed'); + if (notification.method === 'tcp/data') decodeBase64(params(notification).data as string); + } + if (!channel || !countsAgainstWindow(notification.method) + || (!channel.limits && !channel.sends.length)) return this.control(message); + if (signal?.aborted) return Promise.reject(signal.reason ?? new DOMException('Aborted', 'AbortError')); + if (!channel.connected || channel.isClosed) throw new ClientClosedError(); + const serialized = JSON.stringify(notification); + const frozen = JSON.parse(serialized) as JsonRpcNotification; + const bytes = utf16Bytes(serialized); + const tail = channel.sends[channel.sends.length - 1]; + let merged = channel.limits && channel.options.coalescing !== false && !signal && tail && !tail.cancelable && !tail.text + ? mergeNotifications(tail.notification, frozen, + { channel: uri, direction: 'send' }, this.coalescers) : undefined; + if (merged && tail === channel.sends[0] && utf16Bytes(JSON.stringify(merged)) > this.maxBufferedSendBytes) { + merged = undefined; + } + const mergedSerialized = merged ? JSON.stringify(merged) : serialized; + const mergedBytes = utf16Bytes(mergedSerialized); + const behind = channel.sends.slice(1).reduce((sum, send) => sum + send.bytes, 0); + const queuedBytes = merged + ? behind + (tail === channel.sends[0] ? 0 : mergedBytes - tail!.bytes) + : behind + (channel.sends.length ? bytes : 0); + if (queuedBytes > this.maxBufferedSendBytes) { + throw new BackpressureError(uri, this.maxBufferedSendBytes); + } + const send: Send = merged ? tail! : { + notification: frozen, serialized, bytes, offset: 0, completions: [], cleanup: [], cancelable: !!signal, + }; + if (merged) { + send.notification = JSON.parse(mergedSerialized) as JsonRpcNotification; + send.serialized = mergedSerialized; send.bytes = mergedBytes; this.emit(channel, 'coalesced', bytes); + } + else channel.sends.push(send); + if (notification.method === 'tcp/eof') channel.tcpSentEof = true; + const promise = new Promise((resolve, reject) => send.completions.push({ resolve, reject })); + if (signal) { + const abort = () => { + if (send.offset > 0 || send.text !== undefined) return; + const index = channel.sends.indexOf(send); + if (index < 0) return; + // Independently cancelable sends must not be merged. + channel.sends.splice(index, 1); + if (send.notification.method === 'tcp/eof') channel.tcpSentEof = false; + for (const completion of send.completions) completion.reject(signal.reason ?? new DOMException('Aborted', 'AbortError')); + for (const cleanup of send.cleanup) cleanup(); + channel.changed(); + }; + signal.addEventListener('abort', abort, { once: true }); + send.cleanup.push(() => signal.removeEventListener('abort', abort)); + } + channel.changed(); this.schedule(); + return promise; + } + + control(message: JsonRpcMessage | string, creditOwner?: WindowChannel): Promise { + if (this.stopped) return Promise.reject(new ClientClosedError()); + const windowed = [...this.channels.values()].some(channel => !!channel.limits); + if (creditOwner) { + const queued = this.controls.find(control => control.creditOwner === creditOwner); + if (queued) { queued.message = message; return Promise.resolve(); } + } + if (!creditOwner && windowed && this.controls.length >= 4096) { + return Promise.reject(new BackpressureError('ahp-root://', MAX_CONTROL_BYTES)); + } + const shape = typeof message === 'string' ? JSON.parse(message) as JsonRpcMessage : message; + const exempt = !('id' in shape) && 'method' in shape && !countsAgainstWindow(shape.method); + if (windowed && exempt && utf16Bytes(typeof message === 'string' ? message : JSON.stringify(message)) > MAX_CONTROL_BYTES) { + return Promise.reject(new BackpressureError('ahp-root://', MAX_CONTROL_BYTES)); + } + let control!: Control; + const promise = new Promise((resolve, reject) => { + control = { message, resolve, reject, exempt, creditOwner }; + }); + if (!windowed && !this.driving && !this.currentControl && !this.controls.length + && ![...this.channels.values()].some(channel => channel.sends.length)) { + this.currentControl = control; + const finish = (error?: unknown) => { + if (this.currentControl === control) this.currentControl = undefined; + if (error) { + const failure = error instanceof BackpressureError || error instanceof TransportError ? error + : new TransportError('io', `transport send failed: ${String(error)}`, { cause: error }); + control.reject(failure); + if (!(failure instanceof BackpressureError)) this.fail(failure); + } else control.resolve(); + if (this.controls.length || [...this.channels.values()].some(channel => channel.sends.length)) this.schedule(); + }; + try { + control.started = true; + const sent = this.transport.send(message); + if (sent) void Promise.resolve(sent).then(() => finish(), finish); + else finish(); + } catch (error) { finish(error); } + return promise; + } + this.controls.push(control); + this.schedule(); return promise; + } + + checkControl(size: number | undefined): void { + if (size === undefined) protocol('control frame did not preserve raw text'); + if (size > MAX_CONTROL_BYTES) protocol('oversized channel control message'); + const now = performance.now(); + if (now - this.receivedControlPeriod >= CONTROL_INTERVAL_MS) { + this.receivedControlPeriod = now; this.receivedControlCount = 0; + } + if (++this.receivedControlCount > CONTROL_BURST * 4) protocol('channel control rate exceeded'); + } + + schedule(): void { + if (this.scheduled || this.stopped || this.currentControl) return; + this.scheduled = true; + queueMicrotask(() => { void this.drive(); }); + } + + private async drive(): Promise { + if (this.currentControl) { this.scheduled = false; return; } + this.driving = true; + try { + let controlBurst = 0; + while (!this.stopped) { + const channels = [...this.channels.values()]; + let ready: WindowChannel | undefined; + for (let i = 0; i < channels.length; i++) { + const index = (this.cursor + i) % channels.length; + const c = channels[index]; + if (c.connected && c.sends.length && (!c.pauseNewMessages || c.sends[0].offset > 0) + && (!c.limits || c.sent - c.remoteConsumed < c.limits.hostReceive.windowBytes)) { + ready = c; this.cursor = (index + 1) % channels.length; break; + } + } + const now = performance.now(); + if (now - this.controlPeriod >= CONTROL_INTERVAL_MS) { + this.controlPeriod = now; this.controlCount = 0; + } + const controlIndex = this.controls[0]?.exempt && this.controlCount >= CONTROL_BURST + ? this.controls.findIndex(control => !control.exempt) : this.controls.length ? 0 : -1; + if (controlIndex >= 0 && (!ready || controlBurst < 16)) { + const control = this.controls.splice(controlIndex, 1)[0]; + this.currentControl = control; + if (control.exempt) this.controlCount++; + try { + if ([...this.channels.values()].some(channel => !!channel.limits)) { + await this.transport.waitWritable?.({ signal: this.sendAbort.signal }); + } + if (this.stopped) { control.reject(new ClientClosedError()); return; } + if (control.canceled) continue; + if (control.creditOwner?.isClosed) { control.resolve(); continue; } + control.started = true; + await this.transport.send(control.message); control.resolve(); + } + catch (error) { + control.reject(error as Error); + if (!(error instanceof BackpressureError)) throw error; + } + finally { if (this.currentControl === control) this.currentControl = undefined; } + controlBurst++; continue; + } + if (!ready) { + if (!this.controls.length) break; + await new Promise(resolve => setTimeout(resolve, + Math.max(1, CONTROL_INTERVAL_MS - (now - this.controlPeriod)))); + continue; + } + controlBurst = 0; + const send = ready.sends[0]; + if (ready.limits) await this.transport.waitWritable?.({ signal: this.sendAbort.signal }); + if (this.stopped || ready.engine !== this) return; + if (ready.sends[0] !== send) continue; + if (ready.pauseNewMessages && send.offset === 0) continue; + if (ready.limits && ready.sent - ready.remoteConsumed >= ready.limits.hostReceive.windowBytes) continue; + send.text ??= send.serialized; + const previousOffset = send.offset; + const epoch = ready.epoch; + const limit = ready.limits?.hostReceive; + let text = send.text, final = true; + if (limit && (send.offset > 0 || utf16Bytes(text) > limit.maximumFrameBytes)) { + const remaining = limit.windowBytes - (ready.sent - ready.remoteConsumed); + const frame = fragment(text, send.offset, ready.uri, limit.maximumFrameBytes, remaining); + text = frame.text; send.offset = frame.offset; final = frame.final; + } + const size = utf16Bytes(text); + if (!Number.isSafeInteger(ready.sent + size)) protocol('channel send counter overflow'); + if (limit) ready.sent += size; + ready.changed(); + try { await this.transport.send(text); } + catch (error) { + if (!(error instanceof BackpressureError)) throw error; + if (this.stopped || ready.engine !== this) return; + if (limit && epoch === ready.epoch) ready.sent -= size; + send.offset = previousOffset; + if (previousOffset === 0) { + if (ready.sends[0] === send) ready.sends.shift(); + if (send.notification.method === 'tcp/eof') ready.tcpSentEof = false; + for (const cleanup of send.cleanup) cleanup(); + for (const completion of send.completions) completion.reject(error); + } else { + await new Promise(resolve => setTimeout(resolve, 5)); + } + ready.changed(); + continue; + } + if (this.stopped || ready.engine !== this) return; + if (final) { + ready.sends.shift(); + for (const cleanup of send.cleanup) cleanup(); + for (const completion of send.completions) completion.resolve(); + if (send.offset > 0) this.emit(ready, 'fragmentedMessageSent', utf16Bytes(send.text)); + } + ready.changed(); + } + } catch (cause) { + const error = cause instanceof TransportError ? cause + : new TransportError('io', `transport send failed: ${String(cause)}`, { cause }); + this.fail(error); + } finally { + this.driving = false; + this.scheduled = false; + if (!this.stopped && this.controls.length) this.schedule(); + } + } + + stop(error: Error): void { + if (this.stopped) return; + this.stopped = true; + this.sendAbort.abort(error); + this.currentControl?.reject(error); + for (const control of this.controls) control.reject(error); + this.controls = []; + for (const channel of this.channels.values()) { + channel.disconnect(error, this.retain && !channel.uri.startsWith('ahp-tcp:')); + } + this.events.close(); + this.observations.close(); + } +} + +function fragment(text: string, start: number, channel: URI, maximum: number, remaining: number): + { text: string; offset: number; final: boolean } { + const boundary = (end: number) => end < text.length && end > start + && text.charCodeAt(end - 1) >= 0xd800 && text.charCodeAt(end - 1) <= 0xdbff ? end - 1 : end; + const make = (end: number) => JSON.stringify({ jsonrpc: '2.0', method: 'channel/frame', + params: { channel, data: text.slice(start, end), ...(end === text.length ? { final: true } : {}) } }); + const minEnd = start + (text.codePointAt(start)! > 0xffff ? 2 : 1); + const minimum = utf16Bytes(make(minEnd)); + if (minimum > maximum) protocol('fragment cannot carry one character within frame limit'); + const target = Math.min(maximum, Math.max(remaining, minimum)); + let low = minEnd, high = text.length, best = minEnd; + while (low <= high) { + const mid = Math.floor((low + high) / 2), end = boundary(mid); + if (end <= start) { low = mid + 1; continue; } + if (utf16Bytes(make(end)) <= target) { best = end; low = mid + 1; } + else high = mid - 1; + } + + return { text: make(best), offset: best, final: best === text.length }; +} diff --git a/clients/typescript/src/client/hosts/host-client-handle.ts b/clients/typescript/src/client/hosts/host-client-handle.ts index ad6087062..853c29c4c 100644 --- a/clients/typescript/src/client/hosts/host-client-handle.ts +++ b/clients/typescript/src/client/hosts/host-client-handle.ts @@ -11,10 +11,11 @@ * @module client/hosts/host-client-handle */ -import type { CommandMap } from '../../types/common/messages.js'; +import type { CommandMap, ClientNotificationMap } from '../../types/common/messages.js'; import type { StateAction } from '../../types/common/actions.js'; import type { URI } from '../../types/common/state.js'; import type { AhpClient, DispatchHandle } from '../client.js'; +import type { FlowControlStatus, FlowControlEvent } from '../flow-control.js'; import { HostReconnectedError, HostShutDownError, @@ -105,4 +106,32 @@ export class HostClientHandle { this.checkAlive(); return this.client; } + + async dispatchAsync(channel: URI, action: StateAction, + options: { clientSeq?: number; signal?: AbortSignal } = {}): Promise { + this.checkAlive(); + const result = await this.client.dispatchAsync(channel, action, options); + this.checkAlive(); + return result; + } + async notifyAsync(method: M, + params: ClientNotificationMap[M]['params'], options: { signal?: AbortSignal } = {}): Promise { + this.checkAlive(); + await this.client.notifyAsync(method, params, options); + this.checkAlive(); + } + async waitWritable(channel: URI, options: { signal?: AbortSignal } = {}): Promise { + this.checkAlive(); + await this.client.waitWritable(channel, options); + this.checkAlive(); + } + + flowControl(channel: URI): FlowControlStatus { + this.checkAlive(); + return this.client.flowControl(channel); + } + flowControlEvents(): AsyncIterableIterator { + this.checkAlive(); + return this.client.flowControlEvents(); + } } diff --git a/clients/typescript/src/client/hosts/index.ts b/clients/typescript/src/client/hosts/index.ts index 3511d5931..1d89f5f91 100644 --- a/clients/typescript/src/client/hosts/index.ts +++ b/clients/typescript/src/client/hosts/index.ts @@ -33,6 +33,7 @@ export { export type { HostConfig, HostEvent, + HostFlowControlEvent, HostHandle, HostId, HostState, diff --git a/clients/typescript/src/client/hosts/multi.ts b/clients/typescript/src/client/hosts/multi.ts index 86141276e..0dccc1ba6 100644 --- a/clients/typescript/src/client/hosts/multi.ts +++ b/clients/typescript/src/client/hosts/multi.ts @@ -7,7 +7,10 @@ import type { URI } from '../../types/common/state.js'; import type { StateAction } from '../../types/common/actions.js'; import type { SubscribeResult } from '../../types/common/commands.js'; -import type { DispatchHandle } from '../client.js'; +import type { DispatchHandle, SubscribeOptions, Subscription } from '../client.js'; +import type { DeliveryEvent } from '../events.js'; +import type { FlowControlEvent, FlowControlStatus, WindowOptions } from '../flow-control.js'; +import type { ClientNotificationMap } from '../../types/common/messages.js'; import { AsyncBroadcastQueue } from '../async-queue.js'; import { HostClientHandle } from './host-client-handle.js'; import { @@ -59,6 +62,7 @@ export class MultiHostClient { private readonly pendingHostIds = new Set(); private readonly fanOut: AsyncBroadcastQueue; private readonly hostEventsQueue: AsyncBroadcastQueue; + private readonly flowEventsQueue: AsyncBroadcastQueue; private readonly clientIdStore: ClientIdStore; /** * Aborted by {@link shutdown}. Passed into {@link ClientIdStore.load} @@ -85,6 +89,7 @@ export class MultiHostClient { const buffer = options.eventBuffer ?? DEFAULT_EVENT_BUFFER; this.fanOut = new AsyncBroadcastQueue(buffer); this.hostEventsQueue = new AsyncBroadcastQueue(buffer); + this.flowEventsQueue = new AsyncBroadcastQueue(buffer); } /** @@ -155,7 +160,7 @@ export class MultiHostClient { if (this.hosts.has(id)) { throw new DuplicateHostError(id); } - const runtime = new HostRuntime(resolved, clientId, this.fanOut, this.hostEventsQueue); + const runtime = new HostRuntime(resolved, clientId, this.fanOut, this.hostEventsQueue, this.flowEventsQueue); this.hosts.set(id, runtime); runtime.start(); return snapshotHandle(runtime.shared); @@ -258,10 +263,41 @@ export class MultiHostClient { } /** Convenience: subscribe to `uri` on `hostId`. */ - async subscribe(hostId: HostId, uri: URI): Promise { + subscribe(hostId: HostId, uri: URI, options: SubscribeOptions & { windowed: true | WindowOptions }): + Promise<{ result: SubscribeResult; subscription: Subscription }>; + subscribe(hostId: HostId, uri: URI, options?: SubscribeOptions): Promise; + async subscribe(hostId: HostId, uri: URI, options: SubscribeOptions = {}): + Promise }> { const runtime = this.hosts.get(hostId); if (!runtime) throw new UnknownHostError(hostId); - return runtime.subscribe(uri); + const subscribed = await runtime.subscribe(uri, options); + return options.windowed ? subscribed : subscribed.result; + } + + flowControl(hostId: HostId, uri: URI): FlowControlStatus { + const runtime = this.hosts.get(hostId); + if (!runtime) throw new UnknownHostError(hostId); + return runtime.flowControl(uri); + } + flowControlEvents(): AsyncIterableIterator { + return this.flowEventsQueue.reader(); + } + waitWritable(hostId: HostId, uri: URI, options: { signal?: AbortSignal } = {}): Promise { + const runtime = this.hosts.get(hostId); + if (!runtime) throw new UnknownHostError(hostId); + return runtime.waitWritable(uri, options); + } + notifyAsync(hostId: HostId, method: M, + params: ClientNotificationMap[M]['params'], options: { signal?: AbortSignal } = {}): Promise { + const runtime = this.hosts.get(hostId); + if (!runtime) throw new UnknownHostError(hostId); + return runtime.notifyAsync(method, params, options); + } + dispatchAsync(hostId: HostId, uri: URI, action: StateAction, + options: { clientSeq?: number; signal?: AbortSignal } = {}): Promise { + const runtime = this.hosts.get(hostId); + if (!runtime) throw new UnknownHostError(hostId); + return runtime.dispatchAsync(uri, action, options); } /** Convenience: unsubscribe from `uri` on `hostId`. */ @@ -357,6 +393,7 @@ export class MultiHostClient { await Promise.all(runtimes.map(r => r.shutdown('shutdown'))); this.fanOut.close(); this.hostEventsQueue.close(); + this.flowEventsQueue.close(); } private assertOpen(): void { diff --git a/clients/typescript/src/client/hosts/runtime.ts b/clients/typescript/src/client/hosts/runtime.ts index 504924b56..59a59941c 100644 --- a/clients/typescript/src/client/hosts/runtime.ts +++ b/clients/typescript/src/client/hosts/runtime.ts @@ -29,9 +29,13 @@ import type { SessionSummaryChangedParams, } from '../../types/channels-root/notifications.js'; import { PROTOCOL_VERSION } from '../../types/version/registry.js'; -import { AhpClient, type DispatchHandle } from '../client.js'; -import type { ClientEvent } from '../events.js'; -import { RpcError } from '../error.js'; +import { AhpClient, type DispatchHandle, type SubscribeOptions, type Subscription } from '../client.js'; +import type { ClientDeliveryEvent, DeliveryEvent } from '../events.js'; +import type { WindowChannel, FlowControlEvent, FlowControlStatus } from '../flow-control.js'; +import type { ClientNotificationMap } from '../../types/common/messages.js'; +import { ClientClosedError } from '../error.js'; +import { isKnownSnapshot } from '../state-mirror.js'; +import { RpcError, TransportError } from '../error.js'; import type { AsyncBroadcastQueue } from '../async-queue.js'; import { rootReducer } from '../../types/channels-root/reducer.js'; import type { RootAction } from '../../types/action-origin.generated.js'; @@ -284,11 +288,15 @@ const EMPTY_ROOT_STATE: RootState = { agents: [] }; * @internal */ export class HostRuntime { + private retainedChannels = new Map(); + private readonly subscriptionOptions = new Map(); + private readonly trackedChannels = new Map void }>(); private readonly config: ResolvedHostConfig; readonly shared: HostShared; readonly handleSource: HostClientHandleSource; private readonly fanOut: AsyncBroadcastQueue; private readonly hostEvents: AsyncBroadcastQueue; + private readonly flowEvents?: AsyncBroadcastQueue; private readonly shutdownController = new AbortController(); private manualReconnectController = new AbortController(); @@ -306,10 +314,12 @@ export class HostRuntime { resolvedClientId: string, fanOut: AsyncBroadcastQueue, hostEvents: AsyncBroadcastQueue, + flowEvents?: AsyncBroadcastQueue, ) { this.config = config; this.fanOut = fanOut; this.hostEvents = hostEvents; + this.flowEvents = flowEvents; this.shared = makeInitialShared(config, resolvedClientId); this.handleSource = { hostId: this.shared.id, @@ -382,18 +392,23 @@ export class HostRuntime { * subscription list first, so the next successful (re)connect will * subscribe to it automatically. */ - async subscribe(uri: URI): Promise { + async subscribe(uri: URI, options: SubscribeOptions = {}): + Promise<{ result: SubscribeResult; subscription: Subscription }> { if (this.shared.shutdownReason !== null) { throw new HostShutDownError(this.shared.id); } const client = this.shared.currentClient; if (!client) { this.trackSubscription(uri); + this.subscriptionOptions.set(uri, options); throw new HostNotConnectedError(this.shared.id); } - const { result } = await client.subscribe(uri); + const subscribed = await client.subscribe(uri, options); + if (subscribed.subscription.windowChannel?.isClosed) throw new ClientClosedError('subscription closed while subscribing'); + this.subscriptionOptions.set(uri, options); this.trackSubscription(uri); - return result; + if (subscribed.subscription.windowChannel) this.trackChannel(subscribed.subscription.windowChannel); + return subscribed; } /** @@ -407,6 +422,9 @@ export class HostRuntime { } const client = this.shared.currentClient; this.untrackSubscription(uri); + this.subscriptionOptions.delete(uri); + this.retainedChannels.get(uri)?.close(); + this.retainedChannels.delete(uri); if (client) await client.unsubscribe(uri); } @@ -430,6 +448,34 @@ export class HostRuntime { return client.dispatch(channel, action, clientSeq); } + flowControl(uri: URI): FlowControlStatus { + const client = this.shared.currentClient; + if (client) return client.flowControl(uri); + const channel = this.retainedChannels.get(uri); + if (channel) return channel.status; + throw new HostNotConnectedError(this.shared.id); + } + + async waitWritable(uri: URI, options: { signal?: AbortSignal } = {}): Promise { + const client = this.shared.currentClient; + if (!client) throw new HostNotConnectedError(this.shared.id); + return client.waitWritable(uri, options); + } + + async notifyAsync(method: M, + params: ClientNotificationMap[M]['params'], options: { signal?: AbortSignal } = {}): Promise { + const client = this.shared.currentClient; + if (!client) throw new HostNotConnectedError(this.shared.id); + return client.notifyAsync(method, params, options); + } + + async dispatchAsync(uri: URI, action: StateAction, + options: { clientSeq?: number; signal?: AbortSignal } = {}): Promise { + const client = this.shared.currentClient; + if (!client) throw new HostNotConnectedError(this.shared.id); + return client.dispatchAsync(uri, action, options); + } + // ─── Supervisor loop ─────────────────────────────────────────────────────── private async runSupervisor(): Promise { @@ -448,7 +494,7 @@ export class HostRuntime { this.reconnectAck = null; pendingAck?.resolve(); - let connection: { client: AhpClient; events: AsyncIterableIterator } | null = null; + let connection: { client: AhpClient; events: AsyncIterableIterator } | null = null; try { connection = await this.connectOnce(); } catch (err) { @@ -517,6 +563,7 @@ export class HostRuntime { // Final cleanup. The handle source mirrors shared state already. this.handleSource.currentClient = null; this.handleSource.generation = this.shared.generation; + for (const channel of this.retainedChannels.values()) channel.close(new ClientClosedError('host shut down')); } /** @@ -536,7 +583,7 @@ export class HostRuntime { * signal's listeners are detached in a `finally` so successful * connects don't leak listeners on the long-lived shutdown signal. */ - private async connectOnce(): Promise<{ client: AhpClient; events: AsyncIterableIterator }> { + private async connectOnce(): Promise<{ client: AhpClient; events: AsyncIterableIterator }> { const shutdownSignal = this.shutdownController.signal; const manualSignal = this.manualReconnectController.signal; const link = linkAbortSignals(shutdownSignal, manualSignal); @@ -553,20 +600,29 @@ export class HostRuntime { const transport = transportResult; const client = new AhpClient(transport, this.config.clientConfig); + client.restoreWindowedSubscriptions(this.retainedChannels); + client.retainWindowedSubscriptions(); client.connect(); // Attach the events stream BEFORE the handshake so any // notifications the server pushes between the handshake response // and `runConnection` are captured rather than dropped by the // broadcast queue's no-replay-for-late-readers behaviour. - const events = client.events(); + const events = client.deliveryEvents(); + const flowEvents = client.flowControlEvents(); + void (async () => { + for await (const event of flowEvents) this.flowEvents?.publish({ ...event, hostId: this.shared.id }); + })(); let success = false; try { const prior = { serverSeq: this.shared.serverSeq, - subscriptions: [...this.shared.subscriptions], + subscriptions: this.shared.subscriptions.filter(uri => !uri.startsWith('ahp-tcp:')), }; - const canReconnect = prior.serverSeq > 0 && prior.subscriptions.length > 0; + const windows = client.windowedResumeOptions(); + const initialSubscriptions = prior.subscriptions.filter(uri => + !this.subscriptionOptions.get(uri)?.windowed && !windows.items.some(i => i.channel === uri)); + const canReconnect = (prior.serverSeq > 0 || windows.items.length > 0) && prior.subscriptions.length > 0; let reconnectResult: ReconnectResult | null = null; let initSnapshots: Snapshot[] | null = null; @@ -583,6 +639,7 @@ export class HostRuntime { clientId: this.shared.clientId, lastSeenServerSeq: prior.serverSeq, subscriptions: prior.subscriptions, + ...(windows.items.length ? { windows } : {}), }), cancelSignal, ); @@ -599,7 +656,7 @@ export class HostRuntime { client.initialize({ clientId: this.shared.clientId, protocolVersions: [PROTOCOL_VERSION], - initialSubscriptions: prior.subscriptions, + initialSubscriptions, }), cancelSignal, ); @@ -616,7 +673,7 @@ export class HostRuntime { client.initialize({ clientId: this.shared.clientId, protocolVersions: [PROTOCOL_VERSION], - initialSubscriptions: prior.subscriptions, + initialSubscriptions, }), cancelSignal, ); @@ -629,6 +686,27 @@ export class HostRuntime { initCompletionTriggers = initResult.completionTriggerCharacters ?? []; } + for (const uri of prior.subscriptions) { + const options = this.subscriptionOptions.get(uri); + const previouslyWindowed = windows.items.some(i => i.channel === uri); + if ((initSnapshots !== null && (options?.windowed || previouslyWindowed)) + || (reconnectResult !== null && options?.windowed && !previouslyWindowed + && !(reconnectResult.type === 'replay' && reconnectResult.missing.includes(uri)))) { + const channel = client.windowedSubscription(uri)?.windowChannel ?? this.retainedChannels.get(uri); + try { + const subscribed = await raceWithAbort(client.subscribe(uri, options ?? { windowed: true }), cancelSignal); + if (subscribed === ABORTED) throw new Error('re-subscribe aborted'); + if (subscribed.subscription.windowChannel) this.trackChannel(subscribed.subscription.windowChannel); + } catch (error) { + if (cancelSignal.aborted || error instanceof TransportError) throw error; + channel?.close(error as Error); + this.untrackSubscription(uri); + this.subscriptionOptions.delete(uri); + this.retainedChannels.delete(uri); + } + } + } + // Refresh session summaries. Failures are non-fatal — we keep the // cache as-is and the next snapshot/notification will catch up. let summaries: ListSessionsResult | null = null; @@ -664,14 +742,17 @@ export class HostRuntime { // Snapshot variant: refresh root state from the matching // snapshot, then drop subscriptions that were in prior set // but are absent from the returned snapshot list. - const surviving = new Set(reconnectResult.snapshots.map(s => s.resource)); + const surviving = new Set([ + ...reconnectResult.snapshots.map(s => s.resource), + ...(reconnectResult.windows?.items.map(item => item.channel) ?? []), + ]); const priorSet = new Set(prior.subscriptions); snapshotPrunedSubscriptions = this.shared.subscriptions.filter( u => surviving.has(u) || !priorSet.has(u), ); for (const snap of reconnectResult.snapshots) { if (snap.fromSeq > initServerSeq) initServerSeq = snap.fromSeq; - if (snap.resource === ROOT_RESOURCE_URI) { + if (snap.resource === ROOT_RESOURCE_URI && isKnownSnapshot(snap)) { this.shared.rootState = (snap.state as RootState) ?? EMPTY_ROOT_STATE; } } @@ -688,7 +769,7 @@ export class HostRuntime { } if (initSnapshots !== null) { const rootSnap = initSnapshots.find(s => s.resource === ROOT_RESOURCE_URI); - if (rootSnap) this.shared.rootState = (rootSnap.state as RootState) ?? EMPTY_ROOT_STATE; + if (rootSnap && isKnownSnapshot(rootSnap)) this.shared.rootState = (rootSnap.state as RootState) ?? EMPTY_ROOT_STATE; if (initProtocolVersion) this.shared.protocolVersion = initProtocolVersion; this.shared.defaultDirectory = initDefaultDirectory; this.shared.automations = initAutomations; @@ -703,6 +784,15 @@ export class HostRuntime { } else if (snapshotPrunedSubscriptions !== null) { this.shared.subscriptions = snapshotPrunedSubscriptions; } + if (reconnectResult !== null && windows.items.length) { + const kept = new Set(reconnectResult.windows?.items.map(item => item.channel) ?? []); + const dropped = new Set(windows.items.filter(item => !kept.has(item.channel)).map(item => item.channel)); + this.shared.subscriptions = this.shared.subscriptions.filter(uri => !dropped.has(uri)); + for (const uri of dropped) { + this.subscriptionOptions.delete(uri); + this.retainedChannels.delete(uri); + } + } // Mirror the new generation + client into the shared handle source. this.handleSource.generation = this.shared.generation; @@ -752,7 +842,7 @@ export class HostRuntime { * shutdown signal aborts, or a manual reconnect is requested. */ private async runConnection( - events: AsyncIterableIterator, + events: AsyncIterableIterator, ): Promise<'shutdown' | 'manualReconnect' | 'disconnected'> { const shutdownSignal = this.shutdownController.signal; const manualSignal = this.manualReconnectController.signal; @@ -789,6 +879,19 @@ export class HostRuntime { this.shared.currentClient = null; this.handleSource.currentClient = null; if (prev) { + this.retainedChannels = prev.retainWindowedSubscriptions(); + for (const [uri, channel] of this.retainedChannels) { + this.trackSubscription(uri); + this.trackChannel(channel); + if (!this.subscriptionOptions.has(uri)) { + this.subscriptionOptions.set(uri, channel.limits || channel.fallbackReason + ? { windowed: { ...channel.options } } : {}); + } + } + this.shared.subscriptions = this.shared.subscriptions.filter(uri => !uri.startsWith('ahp-tcp:')); + for (const channel of this.retainedChannels.values()) { + this.shared.serverSeq = Math.max(this.shared.serverSeq, channel.lastAppliedServerSeq); + } try { await prev.shutdown(); } catch { @@ -797,7 +900,7 @@ export class HostRuntime { } } - private handleEvent(event: ClientEvent): void { + private handleEvent(event: ClientDeliveryEvent): void { // Update local mirrors before fanning out so consumers observing // the next snapshot see the result of this event. switch (event.event.type) { @@ -818,8 +921,19 @@ export class HostRuntime { case 'authRequired': // No cache update; consumers observe via the event stream. break; + case 'snapshot': + this.shared.serverSeq = Math.max(this.shared.serverSeq, event.event.params.snapshot.fromSeq); + if (event.channel === ROOT_RESOURCE_URI && isKnownSnapshot(event.event.params.snapshot)) { + this.shared.rootState = event.event.params.snapshot.state as RootState; + } + return; + case 'tcpData': + case 'tcpEof': + case 'tcpReset': + case 'notification': + return; } - this.fanOut.publish(tagClientEvent(this.shared.id, event)); + this.fanOut.publish(tagClientEvent(this.shared.id, { channel: event.channel, event: event.event })); } private applyEnvelopeLocally(envelope: ActionEnvelope): void { @@ -841,6 +955,23 @@ export class HostRuntime { } } + private trackChannel(channel: WindowChannel): void { + const existing = this.trackedChannels.get(channel.uri); + if (existing?.channel === channel) return; + if (existing) existing.channel.closed.delete(existing.close); + const close = () => { + if (this.trackedChannels.get(channel.uri)?.channel !== channel) return; + this.trackedChannels.delete(channel.uri); + this.untrackSubscription(channel.uri); + this.subscriptionOptions.delete(channel.uri); + this.retainedChannels.delete(channel.uri); + channel.closed.delete(close); + }; + this.trackedChannels.set(channel.uri, { channel, close }); + channel.closed.add(close); + if (channel.isClosed) close(); + } + private untrackSubscription(uri: URI): void { this.shared.subscriptions = this.shared.subscriptions.filter(u => u !== uri); } diff --git a/clients/typescript/src/client/hosts/state-mirror.ts b/clients/typescript/src/client/hosts/state-mirror.ts index 6e2c6408e..ea1b23325 100644 --- a/clients/typescript/src/client/hosts/state-mirror.ts +++ b/clients/typescript/src/client/hosts/state-mirror.ts @@ -54,6 +54,8 @@ import { automationReducer } from '../../types/channels-automation/reducer.js'; import { automationRunReducer } from '../../types/channels-automation-run/reducer.js'; import { ROOT_RESOURCE_URI, type HostId, type HostSubscriptionEvent } from './types.js'; import { isCanvasState } from '../canvas-state.js'; +import { isKnownSnapshot } from '../state-mirror.js'; +import type { ReceivedDelivery } from '../events.js'; const INITIAL_ROOT: RootState = { agents: [] }; const AUTOMATIONS_URI = 'ahp-automations://' as const; @@ -113,6 +115,15 @@ export class MultiHostStateMirror { private readonly automationsMap = new Map(); private readonly automationRunsMap = new Map(); + /** Consume from a reliable owner queue rather than the legacy lossy fan-in. */ + applyDelivery(hostId: HostId, delivery: ReceivedDelivery): void { + const event = delivery.event; + if (event.type === 'snapshot') { + if (isKnownSnapshot(event.params.snapshot)) this.applySnapshot(hostId, event.params.snapshot); + } else if (event.type === 'action') this.applyEnvelope(hostId, event.params); + delivery.release(); + } + /** All known root states keyed by host. */ get rootStates(): ReadonlyMap { return this.rootStatesMap; diff --git a/clients/typescript/src/client/hosts/types.ts b/clients/typescript/src/client/hosts/types.ts index e4de8e753..ab42e31d3 100644 --- a/clients/typescript/src/client/hosts/types.ts +++ b/clients/typescript/src/client/hosts/types.ts @@ -15,6 +15,11 @@ import type { ClientIdStore } from './client-id-store.js'; import type { HostTransportFactory } from './factory.js'; import { defaultReconnectPolicy, type ReconnectPolicy } from './policy.js'; import type { AhpClientConfig } from '../client.js'; +import type { FlowControlEvent } from '../flow-control.js'; + +export interface HostFlowControlEvent extends FlowControlEvent { + readonly hostId: HostId; +} /** * Opaque, consumer-supplied identifier for a host registered with diff --git a/clients/typescript/src/client/index.ts b/clients/typescript/src/client/index.ts index e264977bc..5783d4695 100644 --- a/clients/typescript/src/client/index.ts +++ b/clients/typescript/src/client/index.ts @@ -13,6 +13,9 @@ export type { ServerRequestHandler, SubscribeOptions, } from './client.js'; +export type { WindowOptions, WindowingConfig, FlowControlStatus, FlowControlEvent, + ChannelCoalescer, CoalescerContext } from './flow-control.js'; +export type { DeliveryEvent, ReceivedDelivery, ClientDeliveryEvent } from './events.js'; export type { ClientEvent, ClosedReason, ConnectionState, SubscriptionEvent } from './events.js'; export { AhpClientError, @@ -20,6 +23,7 @@ export { RpcError, RpcTimeoutError, TransportError, + BackpressureError, } from './error.js'; export type { TransportErrorKind } from './error.js'; export { InMemoryTransport } from './transport.js'; diff --git a/clients/typescript/src/client/managed-subscriptions.ts b/clients/typescript/src/client/managed-subscriptions.ts index 437f8d777..36f51e442 100644 --- a/clients/typescript/src/client/managed-subscriptions.ts +++ b/clients/typescript/src/client/managed-subscriptions.ts @@ -9,7 +9,8 @@ import type { Snapshot, URI } from '../types/common/state.js'; import { AsyncBroadcastQueue } from './async-queue.js'; import { type AhpClient, type SubscribeOptions, type Subscription } from './client.js'; import { ClientClosedError } from './error.js'; -import type { SubscriptionEvent } from './events.js'; +import type { SubscriptionEvent, DeliveryEvent, ReceivedDelivery } from './events.js'; +import type { FlowControlStatus } from './flow-control.js'; /** Lifecycle state of a shared managed subscription. */ export type ManagedSubscriptionStatus = 'pending' | 'active' | 'failed' | 'closed'; @@ -59,6 +60,11 @@ export interface ManagedSubscription< readonly result: ManagedSubscribeResult | undefined; /** Subscribe failure, when {@link status} is `failed`. */ readonly error: Error | undefined; + /** One reliable owner stream, shared by leases; unlike lease.events it is not fan-out. */ + deliveries(): AsyncIterableIterator; + readonly flowControl: FlowControlStatus; + flowControlChanges(): AsyncIterableIterator; + waitWritable(options?: { signal?: AbortSignal }): Promise; } class ManagedSubscriptionState implements ManagedSubscription { @@ -67,6 +73,7 @@ class ManagedSubscriptionState implements ManagedSubscription { private statusValue: ManagedSubscriptionStatus = 'pending'; private resultValue: SubscribeResult | undefined; private errorValue: Error | undefined; + private source?: Subscription; private readonly resolveReady: (result: SubscribeResult) => void; private readonly rejectReady: (error: Error) => void; @@ -97,13 +104,32 @@ class ManagedSubscriptionState implements ManagedSubscription { return this.errorValue; } - activate(result: SubscribeResult): void { + activate(result: SubscribeResult, source: Subscription): void { if (this.statusValue !== 'pending') return; this.resultValue = result; + this.source = source; this.statusValue = 'active'; this.resolveReady(result); } + async *deliveries(): AsyncIterableIterator { + await this.ready; + yield* this.source!.deliveries(); + } + get flowControl(): FlowControlStatus { + return this.source?.flowControl ?? { + mode: 'ordinary', outstandingSendBytes: 0, bufferedSendBytes: 0, bufferedReceiveBytes: 0, writable: false, + }; + } + async *flowControlChanges(): AsyncIterableIterator { + await this.ready; + yield* this.source!.flowControlChanges(); + } + async waitWritable(options: { signal?: AbortSignal } = {}): Promise { + await this.ready; + await this.source!.waitWritable(options); + } + fail(error: Error): void { if (this.statusValue !== 'pending') return; this.errorValue = error; @@ -111,12 +137,13 @@ class ManagedSubscriptionState implements ManagedSubscription { this.rejectReady(error); } - close(message = 'managed subscription closed'): void { + close(message = 'managed subscription closed', error?: Error): void { if (this.statusValue === 'failed' || this.statusValue === 'closed') return; if (this.statusValue === 'pending') { this.rejectReady(new ClientClosedError(message)); } this.statusValue = 'closed'; + this.errorValue = error; } } @@ -133,6 +160,8 @@ export interface ManagedSubscriptionLease< > extends Disposable { readonly subscription: ManagedSubscription; readonly events: AsyncIterableIterator; + /** Consume the shared owner queue; choose one lease to own this stream. */ + deliveries(): AsyncIterableIterator; } interface Entry { @@ -141,7 +170,9 @@ interface Entry { readonly subscription: ManagedSubscriptionState; readonly events: AsyncBroadcastQueue; readonly holders: Map; - source?: Subscription; + source?: Subscription; + observer?: Subscription; + removeListeners?: () => void; } /** @@ -153,7 +184,7 @@ interface Entry { * lease tears down both local fan-out and the server subscription. */ export class ManagedSubscriptionManager { - private readonly client: AhpClient; + private client: AhpClient; private readonly eventBuffer: number; private readonly entries = new Map(); private nextHolderId = 1; @@ -257,6 +288,7 @@ export class ManagedSubscriptionManager { return { subscription: entry.subscription as ManagedSubscription, events: events as AsyncIterableIterator, + deliveries: () => entry.subscription.deliveries(), [Symbol.dispose]: () => { if (released) return; released = true; @@ -272,14 +304,35 @@ export class ManagedSubscriptionManager { private async start(entry: Entry, options: SubscribeOptions): Promise { try { + if (options.windowed) entry.observer = this.client.attachSubscription(entry.uri); const { result, subscription } = await this.client.subscribe(entry.uri, options); if (this.entries.get(entry.uri) !== entry || entry.holders.size === 0) { await subscription.close(); return; } entry.source = subscription; - entry.subscription.activate(result); - void this.pump(entry, subscription); + entry.subscription.activate(result, subscription); + if (subscription.windowChannel) { + const channel = subscription.windowChannel; + const rebound = (client: AhpClient) => { + this.client = client; + entry.observer = client.attachSubscription(entry.uri); + void this.pump(entry, entry.observer); + }; + const closed = (error?: Error) => { + if (this.entries.get(entry.uri) !== entry) return; + this.entries.delete(entry.uri); + entry.subscription.close('managed subscription closed', error); + entry.events.close(); + void entry.observer?.close(); + entry.removeListeners?.(); + }; + channel.rebound.add(rebound); channel.closed.add(closed); + entry.removeListeners = () => { + channel.rebound.delete(rebound); channel.closed.delete(closed); + }; + } + void this.pump(entry, entry.observer ?? subscription); } catch (cause) { if (this.entries.get(entry.uri) !== entry) return; this.entries.delete(entry.uri); @@ -297,9 +350,12 @@ export class ManagedSubscriptionManager { } } finally { if (this.entries.get(entry.uri) === entry) { + const channel = entry.source?.windowChannel; + if (channel && !channel.isClosed && (!channel.connected || source !== entry.observer)) return; this.entries.delete(entry.uri); entry.subscription.close('managed subscription event stream closed'); entry.events.close(); + entry.removeListeners?.(); } } } @@ -309,13 +365,20 @@ export class ManagedSubscriptionManager { entry.events.close(); await this.client.unsubscribe(entry.uri); await entry.source?.close(); + await entry.observer?.close(); } } function subscriptionOptionsKey(options: SubscribeOptions): string { + const window = typeof options.windowed === 'object' ? options.windowed : {}; return JSON.stringify({ maxLatencyMs: options.delivery?.maxLatencyMs ?? null, turns: options.view?.turns ?? null, + windowed: !!options.windowed, + windowBytes: window.windowBytes ?? null, + maximumFrameBytes: window.maximumFrameBytes ?? null, + receive: window.receive ?? 'auto', + coalescing: window.coalescing ?? true, }); } diff --git a/clients/typescript/src/client/state-mirror.ts b/clients/typescript/src/client/state-mirror.ts index 5fcb6e118..caa45d6fd 100644 --- a/clients/typescript/src/client/state-mirror.ts +++ b/clients/typescript/src/client/state-mirror.ts @@ -37,6 +37,7 @@ import { terminalReducer } from '../types/channels-terminal/reducer.js'; import { automationReducer } from '../types/channels-automation/reducer.js'; import { automationRunReducer } from '../types/channels-automation-run/reducer.js'; import { isCanvasState } from './canvas-state.js'; +import type { ReceivedDelivery } from './events.js'; const ROOT_URI = 'ahp-root://' as const; const AUTOMATIONS_URI = 'ahp-automations://' as const; @@ -44,6 +45,20 @@ const AUTOMATIONS_URI = 'ahp-automations://' as const; const INITIAL_ROOT: RootState = { agents: [] }; const INITIAL_AUTOMATION_CATALOG: AutomationState = { entries: [] }; +/** Guards the new ingestion surface without changing legacy applySnapshot behavior. @internal */ +export function isKnownSnapshot(snapshot: Snapshot): boolean { + const state = snapshot.state as unknown as Record; + if (!state || typeof state !== 'object') return false; + if (snapshot.resource === ROOT_URI) return Array.isArray(state.agents); + if (snapshot.resource.startsWith('ahp-session:')) return typeof state.provider === 'string' && Array.isArray(state.chats); + if (snapshot.resource.startsWith('ahp-terminal:')) return Array.isArray(state.content) && typeof state.title === 'string'; + if (snapshot.resource.startsWith('ahp-canvas:')) return isCanvasState(snapshot.state); + if (snapshot.resource === AUTOMATIONS_URI) return Array.isArray(state.entries); + if (snapshot.resource.startsWith('ahp-changeset:')) return Array.isArray(state.files) && typeof state.status === 'string'; + if (snapshot.resource.startsWith('ahp-automation-run:')) return Array.isArray(state.sessions) && typeof state.automation === 'string'; + return false; +} + /** Reducer-driven state container synchronised with server events. */ export class AhpStateMirror { private rootState: RootState = INITIAL_ROOT; @@ -55,6 +70,15 @@ export class AhpStateMirror { private readonly automationsMap = new Map(); private readonly automationRunsMap = new Map(); + /** Apply a typed delivery before releasing its credit. Unknown snapshot state is ignored. */ + applyDelivery(delivery: ReceivedDelivery): void { + const event = delivery.event; + if (event.type === 'snapshot') { + if (isKnownSnapshot(event.params.snapshot)) this.applySnapshot(event.params.snapshot); + } else if (event.type === 'action') this.apply(event.params); + delivery.release(); + } + /** Current root state. */ get root(): RootState { return this.rootState; diff --git a/clients/typescript/src/client/transport.ts b/clients/typescript/src/client/transport.ts index 1f943e677..de3a078ce 100644 --- a/clients/typescript/src/client/transport.ts +++ b/clients/typescript/src/client/transport.ts @@ -38,7 +38,7 @@ export type JsonRpcMessage = * through JSON. */ export type TransportFrame = - | { readonly kind: 'parsed'; readonly message: ProtocolMessage } + | { readonly kind: 'parsed'; readonly message: ProtocolMessage; readonly rawText?: string } | { readonly kind: 'text'; readonly text: string } | { readonly kind: 'binary'; readonly data: Uint8Array }; @@ -48,6 +48,12 @@ export type TransportFrame = * a clean close. */ export interface AhpTransport { + /** Opt in only when inbound raw text is retained and outbound strings are sent verbatim. */ + readonly supportsWindowedDelivery?: boolean; + /** Receive message limit in UTF-16 bytes; unknown defaults to 1048576. */ + readonly maximumFrameBytes?: number; + /** Optional physical send-buffer backpressure. Implementations should honor cancellation. */ + waitWritable?(options?: { signal?: AbortSignal }): Promise; /** * Send a single message. The argument may be either a {@link JsonRpcMessage} * (which the transport serialises) or a pre-serialised JSON-RPC string @@ -94,6 +100,7 @@ export function decodeMessage(text: string): ProtocolMessage { // ─── In-memory transport (for tests) ───────────────────────────────────────── class InMemoryHalf implements AhpTransport { + readonly supportsWindowedDelivery = true; private inbox: Array = []; private waiter: ((frame: TransportFrame | null) => void) | null = null; private closed = false; diff --git a/clients/typescript/src/ws/transport.ts b/clients/typescript/src/ws/transport.ts index 538c13816..e997fb061 100644 --- a/clients/typescript/src/ws/transport.ts +++ b/clients/typescript/src/ws/transport.ts @@ -18,6 +18,8 @@ interface PendingRead { /** Options for {@link WebSocketTransport.connect}. */ export interface WebSocketTransportOptions { + /** AHP receive message limit in UTF-16 bytes. Global WebSocket exposes no limit; default 1048576. */ + maximumFrameBytes?: number; /** * WebSocket subprotocols to negotiate. Browsers do not allow setting * arbitrary HTTP headers; subprotocols and query-string parameters are @@ -47,6 +49,8 @@ export interface WebSocketCloseInfo { * query-string parameters for auth. */ export class WebSocketTransport implements AhpTransport { + readonly supportsWindowedDelivery = true; + readonly maximumFrameBytes: number; private readonly socket: WebSocket; private readonly inbox: Array = []; private waiters: PendingRead[] = []; @@ -63,6 +67,11 @@ export class WebSocketTransport implements AhpTransport { */ static connect(url: string | URL, options: WebSocketTransportOptions = {}): Promise { return new Promise((resolve, reject) => { + if (options.maximumFrameBytes !== undefined + && (!Number.isSafeInteger(options.maximumFrameBytes) || options.maximumFrameBytes < 1024)) { + reject(new TypeError('maximumFrameBytes must be a safe integer of at least 1024')); + return; + } const SocketCtor = (globalThis as { WebSocket?: typeof WebSocket }).WebSocket; if (!SocketCtor) { reject(new TransportError('io', 'global WebSocket is not available; install or polyfill it')); @@ -85,7 +94,7 @@ export class WebSocketTransport implements AhpTransport { }; const onOpen = () => { cleanup(); - resolve(new WebSocketTransport(socket)); + resolve(new WebSocketTransport(socket, options)); }; const onErrorBeforeOpen = () => { cleanup(); @@ -108,14 +117,18 @@ export class WebSocketTransport implements AhpTransport { * * The socket MUST already be in the `OPEN` state. */ - static fromSocket(socket: WebSocket): WebSocketTransport { + static fromSocket(socket: WebSocket, options: WebSocketTransportOptions = {}): WebSocketTransport { if (socket.readyState !== socket.OPEN) { throw new TransportError('io', 'socket is not OPEN'); } - return new WebSocketTransport(socket); + return new WebSocketTransport(socket, options); } - private constructor(socket: WebSocket) { + private constructor(socket: WebSocket, options: WebSocketTransportOptions) { + this.maximumFrameBytes = options.maximumFrameBytes ?? 1_048_576; + if (!Number.isSafeInteger(this.maximumFrameBytes) || this.maximumFrameBytes < 1024) { + throw new TypeError('maximumFrameBytes must be a safe integer of at least 1024'); + } this.socket = socket; socket.binaryType = 'arraybuffer'; @@ -168,6 +181,29 @@ export class WebSocketTransport implements AhpTransport { return this.socket.bufferedAmount; } + /** Bound queued physical sends when the client's windowed scheduler uses this transport. */ + async waitWritable(options: { signal?: AbortSignal } = {}): Promise { + const check = () => { + if (options.signal?.aborted) throw options.signal.reason ?? new DOMException('Aborted', 'AbortError'); + if (this.error) throw this.error; + if (this.closed || this.socket.readyState !== this.socket.OPEN) throw new TransportError('closed', 'transport closed'); + }; + check(); + if (this.socket.bufferedAmount < 1_048_576) return; + while (this.socket.bufferedAmount > 262_144) { + await new Promise((resolve, reject) => { + const cleanup = () => options.signal?.removeEventListener('abort', abort); + const timer = setTimeout(() => { cleanup(); resolve(); }, 5); + const abort = () => { + clearTimeout(timer); cleanup(); + reject(options.signal?.reason ?? new DOMException('Aborted', 'AbortError')); + }; + options.signal?.addEventListener('abort', abort, { once: true }); + }); + check(); + } + } + /** * Information about the most recent close, set once the socket emits * `close`. `null` while the connection is open. diff --git a/clients/typescript/test/windowed-delivery-corpus.test.ts b/clients/typescript/test/windowed-delivery-corpus.test.ts new file mode 100644 index 000000000..02d8bb978 --- /dev/null +++ b/clients/typescript/test/windowed-delivery-corpus.test.ts @@ -0,0 +1,212 @@ +import test from 'node:test'; +import assert from 'node:assert/strict'; +import { readFileSync, readdirSync } from 'node:fs'; +import { AhpClient, InMemoryTransport, type ReceivedDelivery, type Subscription, + type DeliveryEvent, type WindowOptions } from '../src/client/index.js'; +import { utf16Bytes } from '../src/client/flow-control.js'; +import { countsAgainstWindow } from '../src/types/window-exempt.generated.js'; +import type { ClientNotificationMap } from '../src/types/common/messages.js'; + +const corpus = new URL('../../../types/test-cases/windowed-delivery/', import.meta.url); +type ObjectValue = Record; +const delay = (ms: number) => new Promise(resolve => setTimeout(resolve, ms)); + +function matches(actual: any, expected: any): boolean { + if (expected === null || typeof expected !== 'object') return actual === expected; + if (Array.isArray(expected)) { + try { assert.deepEqual(actual, expected); return true; } catch { return false; } + } + return actual !== null && typeof actual === 'object' + && Object.entries(expected).every(([key, value]) => matches(actual[key], value)); +} + +async function within(promise: Promise): Promise { + let timer: ReturnType; + try { + return await Promise.race([promise, new Promise((_, reject) => { + timer = setTimeout(() => reject(new Error('corpus expectation timed out after 1 second')), 1000); + })]); + } finally { clearTimeout(timer!); } +} + +class Driver { + readonly transport = InMemoryTransport.pair(); + readonly client: AhpClient; + readonly sent: ObjectValue[] = []; + readonly requests: ObjectValue[] = []; + readonly sentBytes = new Map(); + readonly subscriptions = new Map>>(); + readonly readers = new Map>(); + readonly deliveries = new Map(); + readonly pendingNext = new Map>>(); + readonly sends: Promise[] = []; + readonly recording: Promise; + + constructor(config: ObjectValue) { + this.client = new AhpClient(this.transport[0], { windowing: config }); + this.client.connect(); + this.recording = this.record(); + } + + private async record(): Promise { + while (true) { + const frame = await this.transport[1].recv(); + if (!frame) return; + assert.equal(frame.kind, 'text'); + if (frame.kind !== 'text') throw new Error('not text'); + const message = JSON.parse(frame.text); + Object.defineProperty(message, 'wireBytes', { value: utf16Bytes(frame.text) }); + this.sent.push(message); + if ('id' in message) this.requests.push(message); + else if (countsAgainstWindow(message.method)) { + const uri = message.params.channel; + this.sentBytes.set(uri, (this.sentBytes.get(uri) ?? 0) + utf16Bytes(frame.text)); + } + } + } + + async wait(predicate: () => boolean): Promise { + const deadline = Date.now() + 1000; + while (!predicate()) { + if (Date.now() >= deadline) throw new Error('corpus expectation timed out after 1 second'); + await delay(1); + } + } + + async reader(channel: string): Promise> { + if (!this.readers.has(channel)) { + const subscription = await within(this.subscriptions.get(channel)!); + this.readers.set(channel, subscription.deliveries()); + } + return this.readers.get(channel)!; + } + + async step(step: ObjectValue): Promise { + const channel = step.channel as string; + switch (step.op) { + case 'subscribe': { + const windowed = step.window !== undefined + ? { ...step.window, receive: step.receive ?? 'auto', coalescing: step.coalescing ?? true } as WindowOptions + : undefined; + const promise = this.client.subscribe(channel, windowed ? { windowed } : {}) + .then(({ subscription }) => subscription as Subscription); + void promise.catch(() => undefined); + this.subscriptions.set(channel, promise); + break; + } + case 'respond': { + await this.wait(() => this.requests.some(r => r.method === step.to)); + const index = this.requests.findIndex(r => r.method === step.to); + const request = this.requests.splice(index, 1)[0]; + this.transport[1].send(JSON.stringify({ jsonrpc: '2.0', id: request.id, result: step.result })); + break; + } + case 'hostSends': this.transport[1].send(step.text); break; + case 'unsubscribe': + await this.client.unsubscribe(channel); + this.readers.delete(channel); + this.pendingNext.delete(channel); + this.subscriptions.delete(channel); + break; + case 'hostCreditsAll': + this.transport[1].send(JSON.stringify({ jsonrpc: '2.0', method: 'channel/credit', + params: { channel, consumedBytes: this.sentBytes.get(channel) ?? 0 } })); + break; + case 'expectSent': { + let candidate: ObjectValue; + let fragments = ''; + while (true) { + await this.wait(() => this.sent.length > 0); + candidate = this.sent.shift()!; + if (step.reassemble && candidate.method === 'channel/frame') { + assert.equal(candidate.params.channel, step.match.params.channel); + const limit = this.client.flowControl(candidate.params.channel).limits!.hostReceive.maximumFrameBytes; + assert.ok(candidate.wireBytes <= limit); + fragments += candidate.params.data; + if (!candidate.params.final) continue; + candidate = JSON.parse(fragments); + } + if (matches(candidate, step.match)) break; + if (candidate.method === 'channel/credit') continue; + assert.fail(`unexpected sent message: ${JSON.stringify(candidate)}; expected ${JSON.stringify(step.match)}`); + } + for (const path of step.absent ?? []) { + let value: any = candidate!; + for (const component of path.split('.')) value = value?.[component]; + assert.equal(value, undefined, `${path} must be absent`); + } + break; + } + case 'expectNotSent': { + const until = Date.now() + step.withinMs; + do { + assert.ok(!this.sent.some(s => matches(s, step.match)), `unexpected message ${JSON.stringify(step.match)}`); + await delay(1); + } while (Date.now() < until); + break; + } + case 'next': { + const reader = await this.reader(channel); + const next = this.pendingNext.get(channel) ?? reader.next(); + this.pendingNext.delete(channel); + const result = await within(next); + assert.ok(!result.done); + assert.ok(matches(result.value.notification, step.expect), + `delivery ${JSON.stringify(result.value.notification)} did not match ${JSON.stringify(step.expect)}`); + if (step.expectBytes !== undefined) assert.equal(result.value.bytes, step.expectBytes); + const history = this.deliveries.get(channel) ?? []; + history.push(result.value); this.deliveries.set(channel, history); + break; + } + case 'advance': { + const promise = (await this.reader(channel)).next(); + void promise.catch(() => undefined); + this.pendingNext.set(channel, promise); + break; + } + case 'release': this.deliveries.get(channel)![step.index].release(); break; + case 'send': { + const promise = this.client.notifyAsync(step.method as keyof ClientNotificationMap, step.params); + void promise.catch(() => undefined); + this.sends.push(promise); + break; + } + case 'expectStatus': + await this.wait(() => matches(this.client.flowControl(channel), step.match)); + break; + case 'expectClosed': + await this.wait(() => this.client.connectionState.status === 'closed'); + assert.equal((this.client.connectionState as any).reason.error.kind, 'protocol'); + break; + default: assert.fail(`unknown corpus operation ${step.op}`); + } + } + + async close(): Promise { + await this.client.shutdown(); + await this.transport[1].close(); + await this.recording; + await Promise.allSettled(this.sends); + } +} + +const sizes = JSON.parse(readFileSync(new URL('sizes.json', corpus), 'utf8')); +for (const size of sizes.cases) { + test(`window corpus size: ${size.name}`, () => { + assert.equal(utf16Bytes(size.text), size.utf16Bytes); + assert.equal(new TextEncoder().encode(size.text).length, size.utf8Bytes); + }); +} + +for (const file of readdirSync(new URL('scenarios/', corpus)).filter(name => name.endsWith('.json')).sort()) { + test(`window corpus: ${file}`, async () => { + const scenario = JSON.parse(readFileSync(new URL(`scenarios/${file}`, corpus), 'utf8')); + const driver = new Driver(scenario.windowing); + try { + for (let i = 0; i < scenario.steps.length; i++) { + try { await driver.step(scenario.steps[i]); } + catch (cause) { throw new Error(`${file}, step ${i}: ${scenario.steps[i].op}`, { cause }); } + } + } finally { await driver.close(); } + }); +} diff --git a/clients/typescript/test/windowing-hosts-lifecycle.test.ts b/clients/typescript/test/windowing-hosts-lifecycle.test.ts new file mode 100644 index 000000000..c82646467 --- /dev/null +++ b/clients/typescript/test/windowing-hosts-lifecycle.test.ts @@ -0,0 +1,110 @@ +import test, { type TestContext } from 'node:test'; +import assert from 'node:assert/strict'; +import { InMemoryTransport, type AhpTransport } from '../src/client/index.js'; +import { MultiHostClient, immediateForeverPolicy } from '../src/client/hosts/index.js'; + +const BAD = 'ahp-terminal:/bad'; +const GOOD = 'ahp-terminal:/good'; +const flowControl = { clientReceive: { windowBytes: 4096, maximumFrameBytes: 1024 }, + hostReceive: { windowBytes: 4096, maximumFrameBytes: 1024 } }; +const sleep = (ms = 1) => new Promise(resolve => setTimeout(resolve, ms)); +async function wait(predicate: () => boolean) { + for (let i = 0; i < 1000; i++) { if (predicate()) return; await sleep(); } + assert.fail('host lifecycle expectation timed out'); +} + +async function setup(t: TestContext, behavior: 'error' | 'timeout' | 'hang' | 'normal') { + const multi = new MultiHostClient(); + const servers: AhpTransport[] = []; + const sent: Array<{ connection: number; message: any }> = []; + t.after(async () => { for (const server of servers) await server.close(); await multi.shutdown(); }); + await multi.addHost({ + id: 'lifecycle', label: 'Lifecycle', reconnectPolicy: immediateForeverPolicy(), + clientConfig: { requestTimeoutMs: behavior === 'timeout' ? 40 : 0 }, + transportFactory: async () => { + const [client, server] = InMemoryTransport.pair(); + const connection = servers.length; servers.push(server); + void (async () => { + while (true) { + const frame = await server.recv(); if (!frame) return; + if (frame.kind !== 'text') continue; + const message = JSON.parse(frame.text); + sent.push({ connection, message }); + if (!('id' in message)) continue; + if (message.method === 'reconnect' && behavior !== 'normal') { + await server.send({ jsonrpc: '2.0', id: message.id, error: { code: -32001, message: 'initialize again' } }); + continue; + } + if (connection > 0 && message.method === 'subscribe' && message.params.channel === BAD) { + if (behavior === 'error') { + await server.send({ jsonrpc: '2.0', id: message.id, error: { code: -32003, message: 'channel denied' } }); + continue; + } + if (behavior === 'timeout' || behavior === 'hang') continue; + } + const result = message.method === 'initialize' + ? { protocolVersion: '1.1.0', serverSeq: 1, snapshots: [] } + : message.method === 'subscribe' ? { flowControl } + : message.method === 'reconnect' ? { type: 'replay', actions: [], missing: [] } + : message.method === 'listSessions' ? { items: [] } : {}; + await server.send({ jsonrpc: '2.0', id: message.id, result }); + if (message.method === 'subscribe') await server.send({ + jsonrpc: '2.0', method: 'channel/snapshot', params: { channel: message.params.channel, + snapshot: { resource: message.params.channel, state: {}, fromSeq: connection + 2 } }, + }); + } + })(); + return client; + }, + }); + await wait(() => multi.host('lifecycle')?.state.status === 'connected'); + return { multi, servers, sent }; +} + +for (const behavior of ['error', 'timeout'] as const) { + test(`${behavior} on one re-subscribe drops only that URI and the host connects`, async t => { + const { multi, servers, sent } = await setup(t, behavior); + const bad = (await multi.subscribe('lifecycle', BAD, { windowed: true })).subscription; + const good = (await multi.subscribe('lifecycle', GOOD, { windowed: true })).subscription; + await bad.next(); await good.next(); + await servers[0].close(); + await wait(() => multi.host('lifecycle')!.generation >= 2 && multi.host('lifecycle')!.state.status === 'connected'); + assert.ok(!multi.host('lifecycle')!.subscriptions.includes(BAD)); + assert.ok(multi.host('lifecycle')!.subscriptions.includes(GOOD)); + await assert.rejects(bad.next(), behavior === 'error' ? /denied/ : /timed out/); + assert.equal((await good.next()).value.type, 'snapshot'); + await sleep(10); + assert.equal(servers.length, 2); + assert.equal(sent.filter(s => s.connection === 1 && s.message.method === 'subscribe' + && s.message.params.channel === BAD).length, 1); + }); +} + +test('shutdown interrupts a re-subscribe with no request timeout', async t => { + const { multi, servers, sent } = await setup(t, 'hang'); + await multi.subscribe('lifecycle', BAD, { windowed: true }); + await servers[0].close(); + await wait(() => sent.some(s => s.connection === 1 && s.message.method === 'subscribe')); + let timer: ReturnType; + try { + await Promise.race([multi.shutdown(), new Promise((_, reject) => { + timer = setTimeout(() => reject(new Error('shutdown hung behind re-subscribe')), 500); + })]); + } finally { clearTimeout(timer!); } +}); + +test('closing the managed owner removes its URI and future reconnect never re-subscribes it', async t => { + const { multi, servers, sent } = await setup(t, 'normal'); + const owner = (await multi.subscribe('lifecycle', BAD, { windowed: true })).subscription; + await owner.close(); + assert.ok(!multi.host('lifecycle')!.subscriptions.includes(BAD)); + await servers[0].close(); + await wait(() => multi.host('lifecycle')!.generation >= 2); + const requests = sent.filter(s => s.connection === 1).map(s => s.message); + assert.ok(requests.some(m => m.method === 'reconnect')); + for (const request of requests) { + assert.ok(!request.params?.subscriptions?.includes(BAD)); + assert.ok(!request.params?.windows?.items.some((item: any) => item.channel === BAD)); + assert.ok(!(request.method === 'subscribe' && request.params.channel === BAD)); + } +}); diff --git a/clients/typescript/test/windowing-hosts.test.ts b/clients/typescript/test/windowing-hosts.test.ts new file mode 100644 index 000000000..12ee722fa --- /dev/null +++ b/clients/typescript/test/windowing-hosts.test.ts @@ -0,0 +1,254 @@ +import test from 'node:test'; +import assert from 'node:assert/strict'; +import { InMemoryTransport, ManagedSubscriptionManager, type AhpTransport, type FlowControlEvent, type ReceivedDelivery } from '../src/client/index.js'; +import { MultiHostClient, MultiHostStateMirror, immediateForeverPolicy, HostReconnectedError } from '../src/client/hosts/index.js'; + +const CHANNEL = 'ahp-terminal:/persistent'; +const sleep = (ms = 1) => new Promise(resolve => setTimeout(resolve, ms)); +async function wait(predicate: () => boolean): Promise { + for (let i = 0; i < 1000; i++) { if (predicate()) return; await sleep(); } + assert.fail('host expectation timed out'); +} +function snapshot(channel: string, fromSeq: number) { + return { jsonrpc: '2.0', method: 'channel/snapshot', params: { channel, snapshot: { + resource: channel, + state: channel === 'ahp-root://' ? { agents: [], activeSessions: fromSeq } : { title: 'Shell', content: [], lifecycle: { status: 'running' } }, + fromSeq, + } } }; +} +function action(channel: string, serverSeq: number, data: string) { + return { jsonrpc: '2.0', method: 'action', params: { + channel, serverSeq, action: { type: 'terminal/data', data }, + } }; +} +const accepted = { + clientReceive: { windowBytes: 16384, maximumFrameBytes: 1024 }, + hostReceive: { windowBytes: 1, maximumFrameBytes: 1024 }, +}; + +for (const recovery of ['snapshot', 'replay', 'live', 'dropped'] as const) { + test(`multi-host reconnect ${recovery}: owner queue, sequence, epochs and windows survive`, async t => { + const multi = new MultiHostClient(); + t.after(() => multi.shutdown()); + const telemetry = multi.flowControlEvents(); + const servers: AhpTransport[] = []; + const received: { message: any; connection: number }[] = []; + const newTexts: string[] = []; + const loops: Promise[] = []; + await multi.addHost({ + id: 'host', label: 'Host', reconnectPolicy: immediateForeverPolicy(), + clientConfig: { windowing: { maximumFrameBytes: 1024 } }, + transportFactory: async () => { + const [client, server] = InMemoryTransport.pair(); + const connection = servers.length; + servers.push(server); + loops.push((async () => { + while (true) { + const frame = await server.recv(); + if (!frame) return; + if (frame.kind !== 'text') continue; + const message = JSON.parse(frame.text); + received.push({ message, connection }); + if (!('id' in message)) continue; + if (message.method === 'initialize') { + await server.send(JSON.stringify({ jsonrpc: '2.0', id: message.id, result: { + protocolVersion: '1.1.0', serverSeq: 1, + snapshots: [snapshot('ahp-root://', 1).params.snapshot], tcpConnections: {}, + } })); + } else if (message.method === 'subscribe') { + await server.send(JSON.stringify({ jsonrpc: '2.0', id: message.id, result: { flowControl: accepted } })); + await server.send(JSON.stringify(snapshot(message.params.channel, 7))); + } else if (message.method === 'listSessions') { + await server.send(JSON.stringify({ jsonrpc: '2.0', id: message.id, result: { items: [] } })); + } else if (message.method === 'reconnect') { + const items = recovery === 'dropped' ? [] : message.params.windows.items.map((item: any) => ({ + channel: item.channel, flowControl: accepted, recovery, + })); + await server.send(JSON.stringify({ jsonrpc: '2.0', id: message.id, result: { + type: 'snapshot', snapshots: [snapshot('ahp-root://', 20).params.snapshot], windows: { items }, + } })); + if (recovery !== 'dropped') { + const text = JSON.stringify(recovery === 'snapshot' + ? snapshot(CHANNEL, 21) : action(CHANNEL, 9, recovery)); + newTexts.push(text); + await server.send(text); + } + } + } + })()); + return client; + }, + }); + await wait(() => multi.host('host')?.state.status === 'connected'); + const oldClient = multi.client('host')!; + const { subscription } = await multi.subscribe('host', CHANNEL, { windowed: { receive: 'manual' } }); + const reader = subscription.deliveries(); + const initial = await reader.next(); assert.ok(!initial.done); + await servers[0].send(JSON.stringify(action(CHANNEL, 8, 'retained'))); + await wait(() => subscription.windowChannel!.lastAppliedServerSeq === 8); + // The owner may have an unsent write when the transport disappears. It is never replayed. + await multi.notifyAsync('host', 'dispatchAction', { channel: CHANNEL, clientSeq: 1, + action: { type: 'terminal/input', data: 'sent once' } }); + const pending = multi.notifyAsync('host', 'dispatchAction', { channel: CHANNEL, clientSeq: 2, + action: { type: 'terminal/input', data: 'not sent' } }); + void pending.catch(() => undefined); + await servers[0].close(); + await assert.rejects(pending); + await wait(() => multi.host('host')!.generation >= 2 && multi.host('host')!.state.status === 'connected'); + assert.throws(() => oldClient.dispatch(CHANNEL, { type: 'terminal/input', data: 'stale' }), HostReconnectedError); + const reconnect = received.find(r => r.connection === 1 && r.message.method === 'reconnect')!.message; + assert.deepEqual(reconnect.params.windows.items, [{ + channel: CHANNEL, flowControl: { receive: { windowBytes: 262144, maximumFrameBytes: 1024 } }, + lastAppliedServerSeq: 8, + }]); + assert.ok(!received.some(r => r.connection === 1 && r.message.method === 'dispatchAction')); + initial.value.release(); + await sleep(10); + assert.ok(!received.some(r => r.connection === 1 && r.message.method === 'channel/credit')); + const retained = await reader.next(); assert.ok(!retained.done && retained.value.event.type === 'action'); + assert.equal((retained.value.event.params.action as any).data, 'retained'); + retained.value.release(); + if (recovery === 'dropped') { + await assert.rejects(reader.next(), /dropped on reconnect/); + assert.ok(!multi.host('host')!.subscriptions.includes(CHANNEL)); + } else { + assert.equal(multi.client('host')!.rawClient().windowedSubscription(CHANNEL), subscription); + assert.ok(multi.host('host')!.subscriptions.includes(CHANNEL)); + const resumed = await reader.next(); assert.ok(!resumed.done); + assert.equal(resumed.value.event.type, recovery === 'snapshot' ? 'snapshot' : 'action'); + resumed.value.release(); + await wait(() => received.some(r => r.connection === 1 && r.message.method === 'channel/credit')); + const credits = received.filter(r => r.connection === 1 && r.message.method === 'channel/credit'); + assert.equal(credits.at(-1)!.message.params.consumedBytes, newTexts[0].length * 2); + } + assert.equal(multi.host('host')!.activeSessions, 20); + const observed: Array = []; + while (!observed.some(e => e.type === (recovery === 'dropped' ? 'dropped' : 'resumed'))) { + const next = await telemetry.next(); assert.ok(!next.done); observed.push(next.value); + } + assert.ok(observed.every(e => e.hostId === 'host')); + await multi.shutdown(); + await Promise.all(loops); + }); +} + +test('TCP owner closes on connection loss and is never initialized or reconnected', async t => { + const multi = new MultiHostClient(); + t.after(() => multi.shutdown()); + const servers: AhpTransport[] = []; + const messages: any[] = []; + await multi.addHost({ + id: 'tcp', label: 'TCP', reconnectPolicy: immediateForeverPolicy(), + transportFactory: async () => { + const [client, server] = InMemoryTransport.pair(); + servers.push(server); + void (async () => { + while (true) { + const frame = await server.recv(); if (!frame) return; + if (frame.kind !== 'text') continue; + const message = JSON.parse(frame.text); messages.push(message); + if (!('id' in message)) continue; + let result: any = {}; + if (message.method === 'initialize') result = { protocolVersion: '1.1.0', serverSeq: 1, snapshots: [], tcpConnections: {} }; + else if (message.method === 'listSessions') result = { items: [] }; + else if (message.method === 'subscribe') result = { flowControl: accepted }; + else if (message.method === 'reconnect') result = { type: 'replay', actions: [], missing: [] }; + await server.send(JSON.stringify({ jsonrpc: '2.0', id: message.id, result })); + if (message.method === 'subscribe') await server.send(JSON.stringify(snapshot(message.params.channel, 5))); + } + })(); + return client; + }, + }); + await wait(() => multi.host('tcp')?.state.status === 'connected'); + const { subscription } = await multi.subscribe('tcp', 'ahp-tcp:/closed', { windowed: true }); + await subscription.next(); + const next = subscription.next(); void next.catch(() => undefined); + await servers[0].close(); + await assert.rejects(next); + await wait(() => multi.host('tcp')!.generation >= 2); + assert.ok(!multi.host('tcp')!.subscriptions.includes('ahp-tcp:/closed')); + for (const message of messages.filter(m => ['initialize', 'reconnect'].includes(m.method))) { + assert.ok(!(message.params.subscriptions ?? message.params.initialSubscriptions ?? []).includes('ahp-tcp:/closed')); + assert.ok(!message.params.windows?.items.some((i: any) => i.channel.startsWith('ahp-tcp:'))); + } +}); + +test('host-aware mirror applies reliable receipts without URI collisions', () => { + const mirror = new MultiHostStateMirror(); + let releases = 0; + for (const host of ['one', 'two']) { + const n = snapshot(CHANNEL, 1); + const delivery: ReceivedDelivery = { notification: n, bytes: 100, + event: { type: 'snapshot', params: n.params as any }, release: () => { releases++; } }; + mirror.applyDelivery(host, delivery); + } + assert.equal(mirror.terminals.size, 2); + assert.equal(releases, 2); +}); + +for (const dropped of [false, true]) { + test(`managed subscriptions acquired through the host's raw client ${dropped ? 'report drops' : 'rebind leases and observers'}`, async t => { + const multi = new MultiHostClient(); + t.after(() => multi.shutdown()); + const servers: AhpTransport[] = []; + const requests: any[] = []; + await multi.addHost({ + id: 'managed', label: 'Managed', reconnectPolicy: immediateForeverPolicy(), + transportFactory: async () => { + const [client, server] = InMemoryTransport.pair(); + servers.push(server); + void (async () => { + while (true) { + const frame = await server.recv(); if (!frame) return; + if (frame.kind !== 'text') continue; + const request = JSON.parse(frame.text); + if (!('id' in request)) continue; + requests.push(request); + let result: any = {}; + if (request.method === 'initialize') result = { protocolVersion: '1.1.0', serverSeq: 1, snapshots: [] }; + else if (request.method === 'subscribe') result = { flowControl: accepted }; + else if (request.method === 'listSessions') result = { items: [] }; + else if (request.method === 'reconnect') result = { + type: 'replay', actions: [], missing: [], windows: { items: dropped ? [] : [{ + channel: CHANNEL, flowControl: accepted, recovery: 'replay', + }] }, + }; + await server.send(JSON.stringify({ jsonrpc: '2.0', id: request.id, result })); + if (request.method === 'subscribe') await server.send(JSON.stringify(snapshot(request.params.channel, 7))); + if (request.method === 'reconnect' && !dropped) await server.send(JSON.stringify(action(CHANNEL, 8, 'new'))); + } + })(); + return client; + }, + }); + await wait(() => multi.host('managed')?.state.status === 'connected'); + const manager = new ManagedSubscriptionManager(multi.client('managed')!.rawClient()); + t.after(() => manager.close()); + const lease = manager.acquire(CHANNEL, 'owner', { windowed: { receive: 'manual' } }); + await lease.subscription.ready; + const owner = lease.deliveries(); + const initial = await owner.next(); assert.ok(!initial.done); + initial.value.release(); + await servers[0].close(); + await wait(() => multi.host('managed')!.generation >= 2); + assert.ok(requests.find(r => r.method === 'reconnect').params.windows.items.some((i: any) => i.channel === CHANNEL)); + if (dropped) { + await assert.rejects(owner.next(), /dropped/); + assert.equal(lease.subscription.status, 'closed'); + assert.ok(lease.subscription.error); + assert.equal(manager.activeSubscriptions().length, 0); + } else { + const next = await owner.next(); assert.ok(!next.done && next.value.event.type === 'action'); + assert.equal((next.value.event.params.action as any).data, 'new'); + next.value.release(); + assert.equal(lease.subscription.status, 'active'); + assert.equal(manager.activeSubscriptions().length, 1); + assert.equal((await lease.events.next()).value.type, 'action'); + const second = manager.acquire(CHANNEL, 'second', { windowed: { receive: 'manual' } }); + assert.equal(second.subscription, lease.subscription); + second[Symbol.dispose](); + } + lease[Symbol.dispose](); + }); +} diff --git a/clients/typescript/test/windowing-lifecycle.test.ts b/clients/typescript/test/windowing-lifecycle.test.ts new file mode 100644 index 000000000..45858a7d5 --- /dev/null +++ b/clients/typescript/test/windowing-lifecycle.test.ts @@ -0,0 +1,230 @@ +import test, { type TestContext } from 'node:test'; +import assert from 'node:assert/strict'; +import { AhpClient, BackpressureError, ClientClosedError, RpcError, RpcTimeoutError, + type AhpTransport, type JsonRpcMessage, type TransportFrame } from '../src/client/index.js'; + +const URI = 'ahp-terminal:/lifecycle'; +const limits = { clientReceive: { windowBytes: 4096, maximumFrameBytes: 1024 }, + hostReceive: { windowBytes: 1, maximumFrameBytes: 1024 } }; +const sleep = (ms = 1) => new Promise(resolve => setTimeout(resolve, ms)); +async function wait(predicate: () => boolean) { + for (let i = 0; i < 1000; i++) { if (predicate()) return; await sleep(); } + assert.fail('windowing lifecycle expectation timed out'); +} +class Transport implements AhpTransport { + readonly supportsWindowedDelivery = true; + readonly sent: any[] = []; + frames: TransportFrame[] = []; + waiter?: (value: TransportFrame | null) => void; + closed = false; + send(message: JsonRpcMessage | string): void { + this.sent.push(typeof message === 'string' ? JSON.parse(message) : message); + } + recv(): Promise { + if (this.frames.length) return Promise.resolve(this.frames.shift()!); + if (this.closed) return Promise.resolve(null); + return new Promise(resolve => { this.waiter = resolve; }); + } + push(message: object) { + const frame: TransportFrame = { kind: 'text', text: JSON.stringify(message) }; + if (this.waiter) { const waiter = this.waiter; this.waiter = undefined; waiter(frame); } + else this.frames.push(frame); + } + close() { this.closed = true; this.waiter?.(null); this.waiter = undefined; } + accept(id: number, channel = URI) { + this.push({ jsonrpc: '2.0', id, result: { flowControl: limits } }); + this.snapshot(channel); + } + snapshot(channel = URI) { + this.push({ jsonrpc: '2.0', method: 'channel/snapshot', params: { + channel, snapshot: { resource: channel, state: {}, fromSeq: 1 }, + } }); + } +} +function setup(t: TestContext, timeout = 0, cap = 1_048_576) { + const transport = new Transport(); + const client = new AhpClient(transport, { requestTimeoutMs: timeout, windowing: { maxBufferedSendBytes: cap } }); + client.connect(); + t.after(() => client.shutdown()); + return { transport, client }; +} +async function open(client: AhpClient, transport: Transport, uri = URI) { + const subscription = client.subscribe(uri, { windowed: { receive: 'manual' } }); + await wait(() => transport.sent.some(m => m.method === 'subscribe' && m.params.channel === uri)); + transport.accept(transport.sent.findLast(m => m.method === 'subscribe' && m.params.channel === uri).id, uri); + return (await subscription).subscription; +} + +test('close clears pending credits and old receipts cannot credit a new URI lifetime', async t => { + const { client, transport } = setup(t); + const first = await open(client, transport); + const firstReceipt = await first.deliveries().next(); assert.ok(!firstReceipt.done); + firstReceipt.value.release(); // Schedules the internal 5ms timer. + await client.unsubscribe(URI); + const next = await open(client, transport); + const held = await next.deliveries().next(); assert.ok(!held.done); + await sleep(15); + assert.ok(!transport.sent.some(m => m.method === 'channel/credit')); + await next.close(); + const third = await open(client, transport); + held.value.release(); + await sleep(10); + assert.ok(!transport.sent.some(m => m.method === 'channel/credit')); + assert.equal(third.flowControl.outstandingSendBytes, 0); +}); + +test('a large head is accepted; only followers are capped; controls bypass the cap', async t => { + const { client, transport } = setup(t, 0, 16); + const uri = 'ahp-tcp:/lifecycle'; + const sub = await open(client, transport, uri); + const receipt = await sub.deliveries().next(); assert.ok(!receipt.done); + await sub.waitWritable(); + const large = client.notifyAsync('tcp/data', { channel: uri, data: btoa('x'.repeat(4096)) }); + void large.catch(() => undefined); + await wait(() => transport.sent.some(m => m.method === 'channel/frame')); + assert.throws(() => client.notify('tcp/data', { channel: uri, data: 'YQ==' }), BackpressureError); + assert.equal(client.connectionState.status, 'connected'); + receipt.value.release(); + await wait(() => transport.sent.some(m => m.method === 'channel/credit')); + const ping = client.ping(); + await wait(() => transport.sent.some(m => m.method === 'ping')); + transport.push({ jsonrpc: '2.0', id: transport.sent.find(m => m.method === 'ping').id, result: null }); + await ping; + client.notify('tcp/reset', { channel: uri }); + await assert.rejects(large, ClientClosedError); + await wait(() => transport.sent.some(m => m.method === 'tcp/reset')); + await client.unsubscribe(uri); + await wait(() => transport.sent.some(m => m.method === 'unsubscribe')); + assert.equal(client.connectionState.status, 'connected'); +}); + +for (const cause of ['cancel', 'timeout', 'unsubscribe'] as const) { + test(`${cause} during subscribe tolerates late acceptance and data`, async t => { + const { client, transport } = setup(t, cause === 'timeout' ? 20 : 0); + const abort = new AbortController(); + const request = client.subscribe(URI, { windowed: true, signal: abort.signal }); + void request.catch(() => undefined); + await wait(() => transport.sent.some(m => m.method === 'subscribe')); + const id = transport.sent.find(m => m.method === 'subscribe').id; + if (cause === 'cancel') abort.abort(); + if (cause === 'unsubscribe') await client.unsubscribe(URI); + await assert.rejects(request, cause === 'timeout' ? RpcTimeoutError : cause === 'unsubscribe' ? ClientClosedError : /abort/i); + transport.accept(id); + await sleep(10); + assert.equal(client.connectionState.status, 'connected'); + assert.equal(client.windowedSubscription(URI), undefined); + assert.ok(transport.sent.some(m => m.method === 'unsubscribe')); + client.dispatch(URI, { type: 'terminal/input', data: 'not held' }); + assert.ok(transport.sent.some(m => m.method === 'dispatchAction')); + }); +} + +test('a timed out reconnect tolerates its late window result and subsequent snapshot', async t => { + const { client, transport } = setup(t, 20); + const reconnect = client.reconnect({ clientId: 'lifecycle', lastSeenServerSeq: 0, subscriptions: [URI], + windows: { items: [{ channel: URI, flowControl: { receive: { windowBytes: 4096, maximumFrameBytes: 1024 } } }] } }); + await assert.rejects(reconnect, RpcTimeoutError); + const id = transport.sent.find(m => m.method === 'reconnect').id; + transport.push({ jsonrpc: '2.0', id, result: { type: 'replay', actions: [], missing: [], windows: { + items: [{ channel: URI, flowControl: limits, recovery: 'snapshot' }], + } } }); + transport.snapshot(); + await sleep(10); + assert.equal(client.connectionState.status, 'connected'); + assert.equal(client.windowedSubscription(URI), undefined); +}); + +test('RPC subscribe errors remove only never-accepted state', async t => { + const { client, transport } = setup(t); + const pending = client.subscribe(URI, { windowed: true }); + await wait(() => transport.sent.some(m => m.method === 'subscribe')); + transport.push({ jsonrpc: '2.0', id: transport.sent[0].id, error: { code: -32001, message: 'denied' } }); + await assert.rejects(pending, RpcError); + assert.equal(client.windowedSubscription(URI), undefined); + client.dispatch(URI, { type: 'terminal/input', data: 'ordinary' }); + assert.ok(transport.sent.some(m => m.method === 'dispatchAction')); + assert.equal(client.connectionState.status, 'connected'); +}); + +test('ordinary dispatch, notify and unsubscribe reach the wire before immediate shutdown', async t => { + const { client, transport } = setup(t); + const events = client.events(); + const states = client.stateChanges(); + const subscription = client.attachSubscription(URI); + client.dispatch(URI, { type: 'terminal/input', data: 'dispatch' }); + client.notify('dispatchAction', { channel: URI, clientSeq: 2, action: { type: 'terminal/input', data: 'notify' } }); + const unsubscribe = client.unsubscribe(URI); + await client.shutdown(); + await unsubscribe; + assert.deepEqual(transport.sent.map(m => m.method), ['dispatchAction', 'dispatchAction', 'unsubscribe']); + assert.equal((await events.next()).done, true); + assert.equal((await subscription.next()).done, true); + for await (const _state of states) { /* Drain closed state history. */ } + assert.equal((await states.next()).done, true); +}); + +test('closing channels prunes scheduler/owner maps and bounds retired-channel bookkeeping', async t => { + const { client, transport } = setup(t); + for (let i = 0; i < 50; i++) { + const uri = `ahp-terminal:/cycle-${i}`; + const sub = await open(client, transport, uri); + await sub.close(); + assert.equal(client.windowedSubscription(uri), undefined); + } + for (let i = 0; i < 4200; i++) client.notify('unsubscribe', { channel: `ahp-terminal:/retired-${i}` }); + assert.equal((client as any).windowEngine.channels.size, 0); + assert.equal((client as any).windowOwners.size, 0); + assert.ok((client as any).unsubscribedWindowed.size <= 4096); +}); + +test('a queued credit waiting for physical writability is discarded on unsubscribe', async t => { + const { client, transport } = setup(t); + const sub = await open(client, transport); + const receipt = await sub.deliveries().next(); assert.ok(!receipt.done); + let writable!: () => void; + const gate = new Promise(resolve => { writable = resolve; }); + (transport as AhpTransport).waitWritable = () => gate; + receipt.value.release(); + await wait(() => !!(client as any).windowEngine.currentControl?.creditOwner); + await client.unsubscribe(URI); + const replacement = client.subscribe(URI, { windowed: true }); + writable(); + await wait(() => transport.sent.filter(m => m.method === 'subscribe').length === 2); + transport.accept(transport.sent.findLast(m => m.method === 'subscribe').id); + await replacement; + assert.ok(!transport.sent.some(m => m.method === 'channel/credit')); +}); + +test('internal RPC queue overflow rejects the caller without disconnecting', async t => { + const { client, transport } = setup(t); + await open(client, transport); + let writable!: () => void; + (transport as AhpTransport).waitWritable = () => new Promise(resolve => { writable = resolve; }); + const promises: Promise[] = []; + const first = client.ping(); void first.catch(() => undefined); promises.push(first); + await wait(() => !!writable); + for (let i = 0; i < 4096; i++) { + const ping = client.ping(); void ping.catch(() => undefined); promises.push(ping); + } + await assert.rejects(client.ping(), BackpressureError); + assert.equal(client.connectionState.status, 'connected'); + await client.shutdown(); + writable(); + await Promise.allSettled(promises); +}); + +test('transport-local backpressure rejects a send without closing the connection', async t => { + const { client, transport } = setup(t); + await open(client, transport); + const send = transport.send.bind(transport); + transport.send = message => { + const notification = typeof message === 'string' ? JSON.parse(message) : message; + if (notification.method === 'dispatchAction') throw new BackpressureError(URI, 16); + send(message); + }; + await assert.rejects(client.notifyAsync('dispatchAction', { + channel: URI, clientSeq: 1, action: { type: 'terminal/input', data: 'local pressure' }, + }), BackpressureError); + assert.equal(client.connectionState.status, 'connected'); + assert.equal(client.flowControl(URI).outstandingSendBytes, 0); +}); diff --git a/clients/typescript/test/windowing.test.ts b/clients/typescript/test/windowing.test.ts new file mode 100644 index 000000000..7e049a0bf --- /dev/null +++ b/clients/typescript/test/windowing.test.ts @@ -0,0 +1,414 @@ +import test, { type TestContext } from 'node:test'; +import assert from 'node:assert/strict'; +import { AhpClient, AhpStateMirror, BackpressureError, ClientClosedError, + ManagedSubscriptionManager, type AhpTransport, type JsonRpcMessage, type TransportFrame, + type WindowOptions, type AhpClientConfig, type ReceivedDelivery } from '../src/client/index.js'; +import type { ChannelFlowControl } from '../src/types/common/state.js'; +import { utf16Bytes } from '../src/client/flow-control.js'; + +const URI = 'ahp-terminal:/window'; +const sleep = (ms = 0) => new Promise(resolve => setTimeout(resolve, ms)); +async function eventually(predicate: () => boolean): Promise { + for (let i = 0; i < 1000; i++) { if (predicate()) return; await sleep(1); } + assert.fail('condition did not become true'); +} +const flow = (windowBytes = 4096, maximumFrameBytes = 1024): ChannelFlowControl => ({ + clientReceive: { windowBytes, maximumFrameBytes }, hostReceive: { windowBytes, maximumFrameBytes }, +}); +function snapshot(channel = URI, fromSeq = 5) { + return { jsonrpc: '2.0', method: 'channel/snapshot', + params: { channel, snapshot: { resource: channel, state: {}, fromSeq } } }; +} +function action(data: string, serverSeq = 6, channel = URI) { + return { jsonrpc: '2.0', method: 'action', params: { channel, + action: { type: 'terminal/data', data }, serverSeq } }; +} + +class Peer implements AhpTransport { + supportsWindowedDelivery = true; + maximumFrameBytes = 1_048_576; + closed = false; + frames: TransportFrame[] = []; + sent: { text: string; message: any }[] = []; + waiter?: (value: TransportFrame | null) => void; + accepted: ChannelFlowControl | undefined = flow(); + resultSnapshot = false; + onSend?: (message: any) => void; + send(message: JsonRpcMessage | string): void { + const text = typeof message === 'string' ? message : JSON.stringify(message); + const parsed = JSON.parse(text); + this.sent.push({ text, message: parsed }); + if (parsed.method === 'subscribe') { + this.push({ jsonrpc: '2.0', id: parsed.id, result: this.accepted + ? { flowControl: this.accepted } : this.resultSnapshot ? { snapshot: snapshot(parsed.params.channel).params.snapshot } : {} }); + if (this.accepted) this.push(snapshot(parsed.params.channel)); + } else if (parsed.method === 'initialize') { + this.push({ jsonrpc: '2.0', id: parsed.id, + result: { protocolVersion: '0.9.0', snapshots: [], serverSeq: 0 } }); + } + this.onSend?.(parsed); + } + recv(): Promise { + if (this.frames.length) return Promise.resolve(this.frames.shift()!); + if (this.closed) return Promise.resolve(null); + return new Promise(resolve => { this.waiter = resolve; }); + } + push(message: object | string): void { + this.pushFrame({ kind: 'text', text: typeof message === 'string' ? message : JSON.stringify(message) }); + } + pushFrame(frame: TransportFrame): void { + if (this.waiter) { const waiter = this.waiter; this.waiter = undefined; waiter(frame); } + else this.frames.push(frame); + } + close(): void { this.closed = true; this.waiter?.(null); this.waiter = undefined; } + credit(channel = URI): void { + const consumedBytes = this.sent.filter(s => !('id' in s.message) && s.message.params.channel === channel + && !['channel/credit', 'unsubscribe', 'tcp/reset'].includes(s.message.method)) + .reduce((total, s) => total + utf16Bytes(s.text), 0); + this.push({ jsonrpc: '2.0', method: 'channel/credit', params: { channel, consumedBytes } }); + } +} + +function setup(t: TestContext, config: AhpClientConfig = {}) { + const peer = new Peer(); + const client = new AhpClient(peer, config); + client.connect(); + t.after(() => client.shutdown()); + return { peer, client }; +} +async function subscribe(client: AhpClient, options: WindowOptions = {}, channel = URI) { + return (await client.subscribe(channel, { windowed: options })).subscription; +} +async function read(reader: AsyncIterableIterator): Promise { + const result = await reader.next(); assert.ok(!result.done); return result.value; +} + +test('offers on older negotiated versions without altering version offers', async t => { + const { peer, client } = setup(t); + await client.initialize({ clientId: 'test', protocolVersions: ['0.9.0'] }); + const sub = await subscribe(client); + assert.equal(sub.flowControl.mode, 'windowed'); + assert.deepEqual(peer.sent[0].message.params.protocolVersions, ['0.9.0']); + assert.ok(peer.sent[1].message.params.flowControl); +}); + +test('parsed-only transports fall back; TCP cannot fall back', async t => { + const { peer, client } = setup(t); + peer.supportsWindowedDelivery = false; peer.accepted = undefined; peer.resultSnapshot = true; + const sub = await subscribe(client); + assert.equal(sub.flowControl.fallbackReason, 'transport'); + assert.equal(peer.sent[0].message.params.flowControl, undefined); + assert.equal((await sub.next()).value.type, 'snapshot'); + await assert.rejects(subscribe(client, {}, 'ahp-tcp:/unsupported'), /TCP requires/); +}); + +test('parsed frames retain exact raw text; missing raw text closes an accepted connection', async t => { + const { peer, client } = setup(t); + const sub = await subscribe(client, { receive: 'manual' }); + const reader = sub.deliveries(); (await read(reader)).release(); + const text = ` ${JSON.stringify(action('é😀'))} \n`; + peer.pushFrame({ kind: 'parsed', message: JSON.parse(text), rawText: text }); + const value = await read(reader); + assert.equal(value.bytes, utf16Bytes(text)); value.release(); + peer.pushFrame({ kind: 'parsed', message: JSON.parse(JSON.stringify(action('missing'))) }); + await eventually(() => client.connectionState.status === 'closed'); + assert.ok(peer.closed); +}); + +test('binary text is measured after UTF-8 decoding, not by transport length', async t => { + const { peer, client } = setup(t); + const sub = await subscribe(client, { receive: 'manual' }); + const reader = sub.deliveries(); (await read(reader)).release(); + const text = JSON.stringify(action('漢字😀')); + peer.pushFrame({ kind: 'binary', data: new TextEncoder().encode(text) }); + assert.equal((await read(reader)).bytes, text.length * 2); +}); + +test('global events are observers and a tiny legacy buffer never drops owner deliveries', async t => { + const { peer, client } = setup(t, { subscriptionBuffer: 1 }); + const observer = client.events(); + const sub = await subscribe(client, { coalescing: false }); + for (let i = 0; i < 20; i++) peer.push(action(String(i), 6 + i)); + await eventually(() => sub.flowControl.bufferedReceiveBytes > 1000); + assert.equal(peer.sent.filter(s => s.message.method === 'channel/credit').length, 0); + await observer.next(); + assert.equal(peer.sent.filter(s => s.message.method === 'channel/credit').length, 0); + assert.equal((await sub.next()).value.type, 'snapshot'); + for (let i = 0; i < 20; i++) { + const result = await sub.next(); assert.ok(!result.done && result.value.type === 'action'); + assert.equal((result.value.params.action as any).data, String(i)); + } + await sub.return(); + assert.equal((await sub.next()).done, true); +}); + +test('ordinary fallback pins the initial snapshot while keeping ordinary lossy buffering', async t => { + const { peer, client } = setup(t, { subscriptionBuffer: 1 }); + peer.accepted = undefined; peer.resultSnapshot = true; + const sub = await subscribe(client); + for (let i = 0; i < 10; i++) peer.push(action(String(i))); + await sleep(10); + assert.equal((await sub.next()).value.type, 'snapshot'); + assert.equal(sub.flowControl.mode, 'ordinary'); + assert.equal(sub.flowControl.fallbackReason, 'declined'); +}); + +test('custom coalescers run first, preserve raw debit, and are disabled per subscription', async t => { + const { peer, client } = setup(t, { windowing: { coalescers: [{ + method: 'action', + merge(a, b, context) { + assert.equal(context.direction, 'receive'); + return { ...b, params: { ...(b.params as any), + action: { type: 'terminal/data', data: (a.params as any).action.data + '|' + (b.params as any).action.data } } }; + }, + }] } }); + const sub = await subscribe(client, { receive: 'manual' }); + const reader = sub.deliveries(); (await read(reader)).release(); + const a = JSON.stringify(action('a')), b = JSON.stringify(action('b', 7)); + peer.push(a); peer.push(b); + await sleep(5); + const value = await read(reader); + assert.equal((value.notification.params as any).action.data, 'a|b'); + assert.equal(value.bytes, utf16Bytes(a) + utf16Bytes(b)); + value.release(); value.release(); + const updated = await subscribe(client, { receive: 'manual', coalescing: false }); + assert.equal(updated, sub); + (await read(reader)).release(); + peer.push(action('c', 8)); peer.push(action('d', 9)); + assert.equal((await read(reader)).event.type, 'action'); + assert.equal((await read(reader)).event.type, 'action'); +}); + +test('a custom coalescer cannot change method, channel or exempt classification', async t => { + const { peer, client } = setup(t, { windowing: { coalescers: [{ + method: 'action', merge: () => ({ jsonrpc: '2.0', method: 'tcp/reset', params: { channel: URI } }), + }] } }); + const sub = await subscribe(client); + await sub.next(); + peer.push(action('a')); peer.push(action('b')); + await eventually(() => client.connectionState.status === 'closed'); + assert.equal((client.connectionState as any).reason.error.kind, 'protocol'); +}); + +test('bounded synchronous sends throw, async sends wait, status observers never unblock sends', async t => { + const { peer, client } = setup(t, { windowing: { maxBufferedSendBytes: 900 } }); + peer.accepted = flow(1); + const sub = await subscribe(client); + const changes = sub.flowControlChanges(), events = client.flowControlEvents(); + await client.notifyAsync('dispatchAction', { channel: URI, clientSeq: 1, action: { type: 'terminal/input', data: 'first' } }); + const params = { channel: URI, clientSeq: 2, action: { type: 'terminal/input' as const, data: 'x'.repeat(200) } }; + const pending = client.notifyAsync('dispatchAction', params); void pending.catch(() => undefined); + const third = client.notifyAsync('dispatchAction', { ...params, clientSeq: 3 }); void third.catch(() => undefined); + assert.throws(() => client.notify('dispatchAction', { ...params, clientSeq: 4 }), BackpressureError); + assert.ok(!sub.flowControl.writable); + await changes.next(); await events.next(); + const abort = new AbortController(); + const waiting = sub.waitWritable({ signal: abort.signal }); + abort.abort(); + await assert.rejects(waiting, /abort/i); + peer.credit(); + await pending; + peer.credit(); + await third; + assert.equal(sub.flowControl.bufferedSendBytes, 0); + peer.credit(); + await sub.waitWritable(); + assert.ok(sub.flowControl.writable); +}); + +test('queued sends snapshot caller params and independently canceled messages are not coalesced', async t => { + const { peer, client } = setup(t); + peer.accepted = flow(1); + await subscribe(client); + await client.notifyAsync('dispatchAction', { channel: URI, clientSeq: 1, + action: { type: 'terminal/input', data: 'block' } }); + const params = { channel: URI, clientSeq: 2, action: { type: 'terminal/input' as const, data: 'original' } }; + const sent = client.notifyAsync('dispatchAction', params); + params.action.data = 'mutated'; + const abort = new AbortController(); + const canceled = client.notifyAsync('dispatchAction', { ...params, clientSeq: 3 }, { signal: abort.signal }); + abort.abort(); await assert.rejects(canceled, /abort/i); + peer.credit(); await sent; + assert.equal(peer.sent.at(-1)!.message.params.action.data, 'original'); +}); + +test('the fair scheduler interleaves channels, with surrogate-safe fragments and no premature cancellation', async t => { + const { peer, client } = setup(t); + peer.accepted = flow(100_000); + await subscribe(client); await subscribe(client, {}, 'ahp-terminal:/second'); + const abort = new AbortController(); + peer.onSend = message => { if (message.method === 'channel/frame') abort.abort(); }; + const large = client.notifyAsync('dispatchAction', { channel: URI, clientSeq: 1, + action: { type: 'terminal/input', data: '😀"\\'.repeat(1000) } }, { signal: abort.signal }); + const small = client.notifyAsync('dispatchAction', { channel: 'ahp-terminal:/second', clientSeq: 2, + action: { type: 'terminal/input', data: 'small' } }); + await Promise.all([large, small]); + const data = peer.sent.filter(s => !('id' in s.message)); + assert.equal(data[0].message.method, 'channel/frame'); + assert.equal(data[1].message.params.channel, 'ahp-terminal:/second'); + const frames = data.filter(s => s.message.method === 'channel/frame'); + for (const frame of frames) { + assert.ok(utf16Bytes(frame.text) <= 1024); + assert.ok(!/^[\uDC00-\uDFFF]|[\uD800-\uDBFF]$/.test(frame.message.params.data)); + } + assert.equal(frames.at(-1)!.message.params.final, true); + assert.equal(JSON.parse(frames.map(f => f.message.params.data).join('')).params.action.data, '😀"\\'.repeat(1000)); +}); + +test('re-subscribe reduced frame limits apply to the remainder of an in-progress message', async t => { + const { peer, client } = setup(t); + peer.accepted = flow(1, 2048); + await subscribe(client); + const pending = client.notifyAsync('dispatchAction', { channel: URI, clientSeq: 1, + action: { type: 'terminal/input', data: 'x'.repeat(3000) } }); + await eventually(() => peer.sent.some(s => s.message.method === 'channel/frame')); + peer.accepted = flow(100_000, 1024); + await subscribe(client); + peer.credit(); + await pending; + const frames = peer.sent.filter(s => s.message.method === 'channel/frame'); + const joined = JSON.parse(frames.map(s => s.message.params.data).join('')); + assert.equal(joined.params.action.data, 'x'.repeat(3000)); + assert.ok(frames.slice(1).every(s => utf16Bytes(s.text) <= 1024)); +}); + +test('TCP EOF is counted and ordered; reset bypasses a full window and cancels queued data', async t => { + const { peer, client } = setup(t); + peer.accepted = flow(1); + const channel = 'ahp-tcp:/test'; + await subscribe(client, { receive: 'manual' }, channel); + await client.notifyAsync('tcp/data', { channel, data: 'YQ==' }); + const eof = client.notifyAsync('tcp/eof', { channel }); void eof.catch(() => undefined); + assert.ok(!peer.sent.some(s => s.message.method === 'tcp/eof')); + client.notify('tcp/reset', { channel }); + await assert.rejects(eof, ClientClosedError); + await eventually(() => peer.sent.some(s => s.message.method === 'tcp/reset')); + assert.throws(() => client.notify('tcp/data', { channel, data: 'Yg==' }), ClientClosedError); +}); + +test('managed lease fan-out is observer-only; the shared reliable owner controls credit', async t => { + const { peer, client } = setup(t); + const manager = new ManagedSubscriptionManager(client); + const lease = manager.acquire(URI, 'owner', { windowed: { receive: 'manual' } }); + const observer = manager.acquire(URI, 'observer', { windowed: { receive: 'manual' } }); + await lease.subscription.ready; + peer.push(action('payload')); + await observer.events.next(); + await sleep(10); + assert.ok(!peer.sent.some(s => s.message.method === 'channel/credit')); + const reader = lease.deliveries(); + (await read(reader)).release(); + (await read(reader)).release(); + await eventually(() => peer.sent.some(s => s.message.method === 'channel/credit')); + assert.throws(() => manager.acquire(URI, 'conflict', { windowed: { receive: 'auto' } }), /differ/); + lease[Symbol.dispose](); observer[Symbol.dispose](); + await manager.close(); +}); + +test('mirrors consume snapshots/actions before release and safely ignore unfamiliar state', async t => { + const { peer, client } = setup(t); + const sub = await subscribe(client, { receive: 'manual' }); + const reader = sub.deliveries(), mirror = new AhpStateMirror(); + const unknown = await read(reader); + mirror.applyDelivery(unknown); + peer.push({ ...snapshot(), params: { channel: URI, + snapshot: { resource: URI, state: { title: 'Known', content: [], lifecycle: { status: 'running' } }, fromSeq: 6 } } }); + mirror.applyDelivery(await read(reader)); + peer.push(action('hello', 7)); mirror.applyDelivery(await read(reader)); + assert.equal((mirror.getTerminal(URI)!.content[0] as any).value, 'hello'); + await eventually(() => sub.flowControl.bufferedReceiveBytes === 0); +}); + +test('late windowed frames after unsubscribe are ignored, not treated as new subscription violations', async t => { + const { peer, client } = setup(t); + await subscribe(client); + await client.unsubscribe(URI); + peer.push({ jsonrpc: '2.0', method: 'channel/frame', params: { channel: URI, data: 'old', final: true } }); + await sleep(10); + assert.equal(client.connectionState.status, 'connected'); +}); + +test('closing a windowed owner unsubscribes and re-subscribe starts fresh accounting', async t => { + const { peer, client } = setup(t); + const first = await subscribe(client); + await first.next(); + await client.notifyAsync('dispatchAction', { channel: URI, clientSeq: 1, + action: { type: 'terminal/input', data: 'first' } }); + assert.ok(first.flowControl.outstandingSendBytes > 0); + await first.close(); + peer.push(action('between consumers', 8)); + const second = await subscribe(client); + assert.notEqual(first, second); + assert.equal((await first.next()).done, true); + assert.equal((await second.next()).value.type, 'snapshot'); + assert.equal(second.flowControl.outstandingSendBytes, 0); + assert.ok(peer.sent.some(sent => sent.message.method === 'unsubscribe')); + peer.push({ jsonrpc: '2.0', method: 'channel/credit', params: { channel: URI, consumedBytes: 0 } }); + await eventually(() => second.flowControl.outstandingSendBytes === 0); +}); + +test('an explicit ordinary re-subscribe drains queued unsent messages without resetting their order', async t => { + const { peer, client } = setup(t); + peer.accepted = flow(1); + await subscribe(client); + await client.notifyAsync('dispatchAction', { channel: URI, clientSeq: 1, + action: { type: 'terminal/input', data: 'first' } }); + const second = client.notifyAsync('dispatchAction', { channel: URI, clientSeq: 2, + action: { type: 'terminal/input', data: 'second' } }); + const third = client.notifyAsync('dispatchAction', { channel: URI, clientSeq: 3, + action: { type: 'terminal/input', data: 'third' } }); + peer.accepted = undefined; peer.resultSnapshot = true; + await client.subscribe(URI, { windowed: false }); + await Promise.all([second, third]); + assert.deepEqual(peer.sent.filter(s => s.message.method === 'dispatchAction').map(s => s.message.params.clientSeq), [1, 2, 3]); +}); + +test('TCP windowing, initial-subscription and reconnect restrictions also cover raw requests', async t => { + const { client } = setup(t); + await assert.rejects(client.request('subscribe', { channel: 'ahp-tcp:/new' }), /windowed/); + await assert.rejects(client.reconnect({ clientId: 'test', lastSeenServerSeq: 0, subscriptions: ['ahp-tcp:/new'] }), /reconnect/); + await assert.rejects(client.initialize({ clientId: 'test', protocolVersions: ['1.1.0'], + initialSubscriptions: ['ahp-tcp:/new'] }), /initial subscription/); +}); + +test('a delayed old transport write cannot remove a new connection queued send', async t => { + const { peer, client } = setup(t); + peer.accepted = flow(1); + const owner = await subscribe(client); + const normalSend = peer.send.bind(peer); + let finishOldWrite!: () => void; + peer.send = message => { + normalSend(message); + const value = typeof message === 'string' ? JSON.parse(message) : message; + if (value.method === 'dispatchAction') return new Promise(resolve => { finishOldWrite = resolve; }); + }; + const oldWrite = client.notifyAsync('dispatchAction', { channel: URI, clientSeq: 1, + action: { type: 'terminal/input', data: 'old' } }); + void oldWrite.catch(() => undefined); + await eventually(() => !!finishOldWrite); + const retained = client.retainWindowedSubscriptions(); + await client.shutdown(); + await assert.rejects(oldWrite, ClientClosedError); + const next = setup(t); + next.peer.accepted = flow(1); + next.client.restoreWindowedSubscriptions(retained); + next.peer.onSend = request => { + if (request.method === 'reconnect') next.peer.push({ jsonrpc: '2.0', id: request.id, result: { + type: 'replay', actions: [], missing: [], windows: { items: [{ + channel: URI, flowControl: flow(1), recovery: 'replay', + }] }, + } }); + }; + await next.client.reconnect({ clientId: 'test', lastSeenServerSeq: 5, subscriptions: [URI], + windows: next.client.windowedResumeOptions() }); + await next.client.notifyAsync('dispatchAction', { channel: URI, clientSeq: 2, + action: { type: 'terminal/input', data: 'new' } }); + const queued = next.client.notifyAsync('dispatchAction', { channel: URI, clientSeq: 3, + action: { type: 'terminal/input', data: 'queued' } }); + finishOldWrite(); + await sleep(10); + assert.ok(owner.flowControl.bufferedSendBytes > 0); + next.peer.credit(); + await queued; + assert.deepEqual(next.peer.sent.filter(s => s.message.method === 'dispatchAction').map(s => s.message.params.clientSeq), [2, 3]); +}); diff --git a/clients/typescript/test/ws-transport.test.ts b/clients/typescript/test/ws-transport.test.ts index d5b93538a..c36e6294a 100644 --- a/clients/typescript/test/ws-transport.test.ts +++ b/clients/typescript/test/ws-transport.test.ts @@ -100,3 +100,37 @@ test('lastClose reflects the wasClean flag from the close event', async () => { assert.equal(info.wasClean, false); assert.equal(info.reason, 'gone'); }); + +test('WebSocket windowing advertises the 1 MiB fallback and a configurable override', async () => { + const { WebSocketTransport } = await import('../src/ws/index.js'); + const { transport } = await makeTransport(); + assert.equal(transport.supportsWindowedDelivery, true); + assert.equal(transport.maximumFrameBytes, 1_048_576); + const overridden = WebSocketTransport.fromSocket(new MockWebSocket() as unknown as WebSocket, { + maximumFrameBytes: 16384, + }); + assert.equal(overridden.maximumFrameBytes, 16384); + assert.throws(() => WebSocketTransport.fromSocket(new MockWebSocket() as unknown as WebSocket, { + maximumFrameBytes: 512, + }), TypeError); + await assert.rejects(WebSocketTransport.connect('ws://unused.invalid', { maximumFrameBytes: 0 }), TypeError); +}); + +test('WebSocket physical writability waits for drain and honors cancellation and close', async () => { + const { ws, transport } = await makeTransport(); + ws.bufferedAmount = 2_097_152; + let writable = false; + const pending = transport.waitWritable().then(() => { writable = true; }); + await new Promise(resolve => setTimeout(resolve, 10)); + assert.equal(writable, false); + ws.bufferedAmount = 262_144; + await pending; + ws.bufferedAmount = 2_097_152; + const controller = new AbortController(); + const canceled = transport.waitWritable({ signal: controller.signal }); + controller.abort(); + await assert.rejects(canceled, /abort/i); + const closed = transport.waitWritable(); + ws.close(); + await assert.rejects(closed, TransportError); +}); diff --git a/docs/.changes/20261014-dotnet-windowed-delivery.json b/docs/.changes/20261014-dotnet-windowed-delivery.json new file mode 100644 index 000000000..24691730a --- /dev/null +++ b/docs/.changes/20261014-dotnet-windowed-delivery.json @@ -0,0 +1,7 @@ +{ + "type": "added", + "message": "Windowed delivery for subscriptions: configure with `ClientConfig.Windowing` and opt in with `SubscribeWindowedAsync`. Credit follows the consumer, sends wait for the host window (`WaitWritableAsync`), queued `terminal/data`, `chat/delta`, and `tcp/data` are coalesced, and flow control is reported through the existing `Meter`. Off by default.", + "targets": [ + "dotnet" + ] +} diff --git a/docs/.changes/20261014-go-windowed-delivery.json b/docs/.changes/20261014-go-windowed-delivery.json new file mode 100644 index 000000000..1457ac2ce --- /dev/null +++ b/docs/.changes/20261014-go-windowed-delivery.json @@ -0,0 +1,7 @@ +{ + "type": "added", + "message": "Windowed delivery for subscriptions: configure with `Config.Windowing` and opt in with `SubscribeWithOptions`. Credit follows the consumer (`Next` and `Delivery.Release`), sends wait for the host window (`WaitWritable`), and queued `terminal/data`, `chat/delta`, and `tcp/data` are coalesced. Off by default.", + "targets": [ + "go" + ] +} diff --git a/docs/.changes/20261014-rust-windowed-delivery.json b/docs/.changes/20261014-rust-windowed-delivery.json new file mode 100644 index 000000000..c77c839d8 --- /dev/null +++ b/docs/.changes/20261014-rust-windowed-delivery.json @@ -0,0 +1,7 @@ +{ + "type": "added", + "message": "Windowed delivery for subscriptions: configure with `Client::builder().windowing(...)` and opt in with `subscribe_windowed`. Dropping or releasing a `Delivery` returns credit, sends wait for the host window (`wait_writable`), and queued `terminal/data`, `chat/delta`, and `tcp/data` are coalesced. Off by default.", + "targets": [ + "rust" + ] +} diff --git a/docs/.changes/20261014-swift-windowed-delivery.json b/docs/.changes/20261014-swift-windowed-delivery.json new file mode 100644 index 000000000..9ec829353 --- /dev/null +++ b/docs/.changes/20261014-swift-windowed-delivery.json @@ -0,0 +1,7 @@ +{ + "type": "added", + "message": "Windowed delivery for subscriptions: opt in with `subscribe(_:options: .windowed())`. Credit follows the consumer (`AHPDelivery.release()` in manual mode), sends wait for the host window (`waitUntilWritable()`), and queued `terminal/data`, `chat/delta`, and `tcp/data` are coalesced. Off by default.", + "targets": [ + "swift" + ] +} diff --git a/docs/.changes/20261014-terminal-snapshot-decoding.json b/docs/.changes/20261014-terminal-snapshot-decoding.json new file mode 100644 index 000000000..adeba953a --- /dev/null +++ b/docs/.changes/20261014-terminal-snapshot-decoding.json @@ -0,0 +1,9 @@ +{ + "type": "fixed", + "message": "Terminal snapshots no longer decode as session state. Both have a `lifecycle` member; a snapshot is now treated as session state only if it also has `provider`.", + "targets": [ + "go", + "kotlin", + "dotnet" + ] +} diff --git a/docs/.changes/20261014-typescript-windowed-delivery.json b/docs/.changes/20261014-typescript-windowed-delivery.json new file mode 100644 index 000000000..82eb324a7 --- /dev/null +++ b/docs/.changes/20261014-typescript-windowed-delivery.json @@ -0,0 +1,7 @@ +{ + "type": "added", + "message": "Windowed delivery for subscriptions: opt in with `subscribe(uri, { windowed: true })`. Credit follows the consumer (`deliveries()` and `release()` in manual mode), sends wait for the host window (`notifyAsync`, `dispatchAsync`, `waitWritable()`), and queued `terminal/data`, `chat/delta`, and `tcp/data` are coalesced. Off by default.", + "targets": [ + "typescript" + ] +} diff --git a/docs/specification/windowed-delivery.md b/docs/specification/windowed-delivery.md index ddf703b86..8f22a8dae 100644 --- a/docs/specification/windowed-delivery.md +++ b/docs/specification/windowed-delivery.md @@ -669,11 +669,18 @@ file in parts. ## Implementation Notes -Generated client models contain the message types for windowed delivery. They -do not implement negotiation, counting, scheduling, or retained queues. A peer -MUST implement the rules on this page before it offers or accepts windowed -delivery. After acceptance, a peer MUST NOT treat data messages as uncounted -and MUST NOT change the subscription to ordinary delivery. +Generated client models contain the message types for windowed delivery. The +TypeScript, Rust, Go, Swift, and .NET client libraries also implement the client +side of windowed delivery: negotiation, counting, credit, fragments, fair +sending, and reconnection. Each client turns it on only for subscriptions that +ask for it. A peer MUST implement the rules on this page before it offers or +accepts windowed delivery. After acceptance, a peer MUST NOT treat data messages +as uncounted, and MUST NOT change the subscription to ordinary delivery on its +own; only a new `subscribe` can change it. + +The shared corpus in `types/test-cases/windowed-delivery/` describes the +expected client behavior as scripted conversations with a host. Each client +library runs it. These message types are optional additions. Peers that do not offer windowed delivery continue to use ordinary subscriptions. In Rust, struct literals for diff --git a/scripts/generate-csharp.ts b/scripts/generate-csharp.ts index 8d9717cdb..0fc2822ca 100644 --- a/scripts/generate-csharp.ts +++ b/scripts/generate-csharp.ts @@ -1503,11 +1503,10 @@ internal sealed class SnapshotStateConverter : JsonConverter { result.Chat = root.Deserialize(AhpJsonTypeInfo.Get(options)); } - else if (root.TryGetProperty("lifecycle", out _)) + else if (root.TryGetProperty("lifecycle", out _) && root.TryGetProperty("provider", out _)) { - // SessionState is discriminated on its required \`lifecycle\` field. - // (It no longer carries \`summary\`; that field was removed when the - // session state was flattened.) + // SessionState is discriminated on its required \`lifecycle\` and + // \`provider\` fields; TerminalState also has \`lifecycle\`. result.Session = root.Deserialize(AhpJsonTypeInfo.Get(options)); } else if (root.TryGetProperty("instanceId", out _) && diff --git a/scripts/generate-go.ts b/scripts/generate-go.ts index caaee7a4f..a5ea12504 100644 --- a/scripts/generate-go.ts +++ b/scripts/generate-go.ts @@ -1390,7 +1390,7 @@ func (s *SnapshotState) UnmarshalJSON(data []byte) error { \t\t\treturn err \t\t} \t\ts.Automations = &v -\tcase containsAll(probe, "lifecycle"): +\tcase containsAll(probe, "lifecycle", "provider"): \t\tvar v SessionState \t\tif err := json.Unmarshal(data, &v); err != nil { \t\t\treturn err diff --git a/scripts/generate-kotlin.ts b/scripts/generate-kotlin.ts index 0d30f0765..556937e43 100644 --- a/scripts/generate-kotlin.ts +++ b/scripts/generate-kotlin.ts @@ -867,7 +867,8 @@ internal object SnapshotStateSerializer : KSerializer { // Try the most distinctive shape first. AutomationRunState has required // \`automation\`, \`origin\`, and \`sessions\`; AutomationState has // required \`entries\`; SessionState has required - // \`lifecycle\`; ChatState has required \`turns\`; ChangesetState has + // \`lifecycle\` + \`provider\` (TerminalState also has \`lifecycle\`); + // ChatState has required \`turns\`; ChangesetState has // required \`status\` + \`files\`; ResourceWatchState has required // \`root\` + \`recursive\`; AnnotationsState has required \`annotations\` // (checked after session, whose optional annotations summary reuses the @@ -880,7 +881,7 @@ internal object SnapshotStateSerializer : KSerializer { SnapshotState.AutomationRun(input.json.decodeFromJsonElement(AutomationRunState.serializer(), element)) obj.containsKey("entries") -> SnapshotState.Automations(input.json.decodeFromJsonElement(AutomationState.serializer(), element)) - obj.containsKey("lifecycle") -> SnapshotState.Session(input.json.decodeFromJsonElement(SessionState.serializer(), element)) + obj.containsKey("lifecycle") && obj.containsKey("provider") -> SnapshotState.Session(input.json.decodeFromJsonElement(SessionState.serializer(), element)) obj.containsKey("turns") -> SnapshotState.Chat(input.json.decodeFromJsonElement(ChatState.serializer(), element)) obj.containsKey("instanceId") && obj.containsKey("extensionId") && obj.containsKey("canvasId") -> SnapshotState.Canvas(input.json.decodeFromJsonElement(CanvasState.serializer(), element)) diff --git a/types/test-cases/round-trips/076-terminal-snapshot.json b/types/test-cases/round-trips/076-terminal-snapshot.json new file mode 100644 index 000000000..0d090ffe2 --- /dev/null +++ b/types/test-cases/round-trips/076-terminal-snapshot.json @@ -0,0 +1,38 @@ +{ + "name": "terminal-snapshot", + "group": "A", + "description": "A terminal snapshot decodes as terminal state, not session state, even though both have a lifecycle member.", + "type": "Snapshot", + "input": { + "resource": "ahp-terminal:/t1", + "state": { + "title": "Shell", + "content": [], + "lifecycle": { + "status": "running" + }, + "claim": { + "kind": "client", + "clientId": "client-1" + } + }, + "fromSeq": 5 + }, + "acceptableOutputs": [ + { + "resource": "ahp-terminal:/t1", + "state": { + "title": "Shell", + "content": [], + "lifecycle": { + "status": "running" + }, + "claim": { + "kind": "client", + "clientId": "client-1" + } + }, + "fromSeq": 5 + } + ] +} diff --git a/types/test-cases/windowed-delivery/README.md b/types/test-cases/windowed-delivery/README.md new file mode 100644 index 000000000..9e7d1a96e --- /dev/null +++ b/types/test-cases/windowed-delivery/README.md @@ -0,0 +1,78 @@ +# Windowed delivery conformance corpus + +Language-neutral tests for the client side of +[windowed delivery](../../../docs/specification/windowed-delivery.md). Every +client library that implements windowed delivery runs these files through a +small driver against a scripted fake host. Setup differs per language; the +buffering, credit, coalescing, and sending behavior must not. + +- `sizes.json` lists raw JSON texts and their size in UTF-16 bytes. A client + measures each text exactly as received, without parsing and serializing it + again. +- `scenarios/*.json` are scripted conversations between one client and a fake + host. + +`types/windowed-delivery-corpus.test.ts` checks that the corpus is well formed. + +## Driver setup + +The driver connects a client to an in-memory transport that it controls: + +- Text that the fake host sends reaches the client byte for byte, and the + transport reports raw text sizes, so the client can offer windowed delivery. +- The driver records every message the client sends, in order. +- The connection is ready before the first step. If the client requires an + `initialize` handshake, the driver completes it first; the handshake is not + part of the scenario. +- `windowing` at the top level of a scenario is the connection's windowing + configuration (`windowBytes`, `maximumFrameBytes`, and optionally + `maxBufferedSendBytes`). + +## Steps + +Steps run in order. Steps that inject host messages do not wait for the client +to process them; this lets scenarios test messages that arrive back to back. +Every step that expects something waits up to **1 second** before it fails. + +| `op` | Members | Driver action | +|---|---|---| +| `subscribe` | `channel`, `window?`, `receive?`, `coalescing?` | Start a subscription and do not wait for the result. With `window` (an object, possibly empty), opt into windowed delivery with those per-subscription overrides. Without it, subscribe without opting in. `receive` is `auto` (default) or `manual`. `coalescing` defaults to `true`. If the channel is already subscribed, change the existing subscription instead (re-subscribe). | +| `respond` | `to`, `result` | Respond to the oldest unanswered request with method `to`, using its `id`. | +| `hostSends` | `text` | Deliver `text` to the client exactly as written. | +| `hostCreditsAll` | `channel` | Send `channel/credit` with `consumedBytes` equal to the total size of every data message and fragment that the client has sent on `channel` so far, as measured by the driver from the sent text. | +| `expectSent` | `match`, `absent?`, `reassemble?` | Take the client's sent messages in order until one matches `match`. Unmatched `channel/credit` messages are skipped, because credit batching differs between clients. Any other unmatched message fails the step. `absent` lists dotted paths that must not exist in the matched message. With `reassemble: true`, `channel/frame` messages for the channel are collected until `final: true`; each must be at most the accepted host `maximumFrameBytes`, and the joined text is parsed and matched. | +| `expectNotSent` | `match`, `withinMs` | Fail if a message matching `match` is sent, or was already sent and not taken by `expectSent`, within `withinMs` milliseconds. | +| `next` | `channel`, `expect`, `expectBytes?` | Take the next delivery from the subscription. In `auto` mode, this releases the previous delivery. The delivery, written as a JSON-RPC notification `{ "method", "params" }`, must match `expect`. `expectBytes` is the delivery's charged size: the sum of the original messages or fragments it came from. | +| `advance` | `channel` | `auto` mode only: release the current delivery without taking another, as a consumer does when it starts waiting for the next delivery. | +| `release` | `channel`, `index` | `manual` mode only: release the delivery with this zero-based index, counting every delivery taken on the channel during the scenario, across subscriptions. | +| `send` | `channel`, `method`, `params` | Send a notification on the subscription through the client's flow-controlled send path. Do not wait for it to be written. | +| `unsubscribe` | `channel` | Unsubscribe from the channel, as a consumer that is done with it would. Deliveries already taken stay held until released. | +| `expectStatus` | `channel`, `match` | The subscription's flow-control status must match `match`. Waits until it does. | +| `expectClosed` | | The client closes the connection because of a protocol error. | + +### Matching + +`match` and `expect` are partial structural matches. Objects match when every +listed member matches; unlisted members are ignored. Arrays and scalar values +must be equal. + +### Status fields + +`expectStatus` uses these names, written in camelCase, whatever the language's +native names are: + +| Field | Meaning | +|---|---| +| `mode` | `windowed` or `ordinary` | +| `fallbackReason` | `declined` or `transport` when `mode` is `ordinary` after an offer | +| `limits` | The accepted `ChannelFlowControl` (`clientReceive`, `hostReceive`) | +| `bufferedReceiveBytes` | Data received on the subscription and not yet released | +| `outstandingSendBytes` | Data sent and not yet credited by the host | +| `bufferedSendBytes` | Data queued and not yet sent | +| `writable` | `true` when the host window and the send queue both have room | + +### Sizes + +Sizes in receive steps are exact, because the host text is literal. The size of +what the client sends depends on its serializer, so send scenarios check +structure only and use windows with enough margin for any serializer. diff --git a/types/test-cases/windowed-delivery/scenarios/001-accept-snapshot-credit.json b/types/test-cases/windowed-delivery/scenarios/001-accept-snapshot-credit.json new file mode 100644 index 000000000..4bb666ae5 --- /dev/null +++ b/types/test-cases/windowed-delivery/scenarios/001-accept-snapshot-credit.json @@ -0,0 +1,117 @@ +{ + "description": "The client offers its configured limits, the host lowers them, and the snapshot is the first delivery. In auto mode, credit for the snapshot is held until the consumer advances.", + "windowing": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + }, + "steps": [ + { + "op": "subscribe", + "channel": "ahp-terminal:/t1", + "window": {}, + "receive": "auto" + }, + { + "op": "expectSent", + "match": { + "method": "subscribe", + "params": { + "channel": "ahp-terminal:/t1", + "flowControl": { + "receive": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + } + } + } + } + }, + { + "op": "respond", + "to": "subscribe", + "result": { + "flowControl": { + "clientReceive": { + "windowBytes": 512, + "maximumFrameBytes": 1024 + }, + "hostReceive": { + "windowBytes": 512, + "maximumFrameBytes": 1024 + } + } + } + }, + { + "op": "hostSends", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"channel/snapshot\",\"params\":{\"channel\":\"ahp-terminal:/t1\",\"snapshot\":{\"resource\":\"ahp-terminal:/t1\",\"state\":{\"title\":\"Shell\",\"content\":[],\"lifecycle\":{\"status\":\"running\"},\"claim\":{\"kind\":\"client\",\"clientId\":\"client-1\"}},\"fromSeq\":5}}}" + }, + { + "op": "expectStatus", + "channel": "ahp-terminal:/t1", + "match": { + "mode": "windowed", + "limits": { + "clientReceive": { + "windowBytes": 512, + "maximumFrameBytes": 1024 + }, + "hostReceive": { + "windowBytes": 512, + "maximumFrameBytes": 1024 + } + } + } + }, + { + "op": "next", + "channel": "ahp-terminal:/t1", + "expect": { + "method": "channel/snapshot", + "params": { + "channel": "ahp-terminal:/t1", + "snapshot": { + "resource": "ahp-terminal:/t1", + "state": { + "title": "Shell", + "content": [], + "lifecycle": { + "status": "running" + }, + "claim": { + "kind": "client", + "clientId": "client-1" + } + }, + "fromSeq": 5 + } + } + }, + "expectBytes": 520 + }, + { + "op": "expectNotSent", + "match": { + "method": "channel/credit", + "params": { + "channel": "ahp-terminal:/t1" + } + }, + "withinMs": 100 + }, + { + "op": "advance", + "channel": "ahp-terminal:/t1" + }, + { + "op": "expectSent", + "match": { + "method": "channel/credit", + "params": { + "channel": "ahp-terminal:/t1", + "consumedBytes": 520 + } + } + } + ] +} diff --git a/types/test-cases/windowed-delivery/scenarios/002-declined-falls-back.json b/types/test-cases/windowed-delivery/scenarios/002-declined-falls-back.json new file mode 100644 index 000000000..14a14a00e --- /dev/null +++ b/types/test-cases/windowed-delivery/scenarios/002-declined-falls-back.json @@ -0,0 +1,105 @@ +{ + "description": "A result without flowControl declines windowed delivery. The subscription uses ordinary delivery and never sends credit. The result snapshot is still the first delivery.", + "windowing": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + }, + "steps": [ + { + "op": "subscribe", + "channel": "ahp-terminal:/t1", + "window": {} + }, + { + "op": "expectSent", + "match": { + "method": "subscribe", + "params": { + "channel": "ahp-terminal:/t1", + "flowControl": { + "receive": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + } + } + } + } + }, + { + "op": "respond", + "to": "subscribe", + "result": { + "snapshot": { + "resource": "ahp-terminal:/t1", + "state": { + "title": "Shell", + "content": [], + "lifecycle": { + "status": "running" + }, + "claim": { + "kind": "client", + "clientId": "client-1" + } + }, + "fromSeq": 5 + } + } + }, + { + "op": "expectStatus", + "channel": "ahp-terminal:/t1", + "match": { + "mode": "ordinary", + "fallbackReason": "declined" + } + }, + { + "op": "next", + "channel": "ahp-terminal:/t1", + "expect": { + "method": "channel/snapshot", + "params": { + "channel": "ahp-terminal:/t1", + "snapshot": { + "resource": "ahp-terminal:/t1", + "fromSeq": 5 + } + } + } + }, + { + "op": "hostSends", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"action\",\"params\":{\"channel\":\"ahp-terminal:/t1\",\"action\":{\"type\":\"terminal/data\",\"data\":\"Hello\\r\\n\"},\"serverSeq\":6}}" + }, + { + "op": "next", + "channel": "ahp-terminal:/t1", + "expect": { + "method": "action", + "params": { + "channel": "ahp-terminal:/t1", + "action": { + "type": "terminal/data", + "data": "Hello\r\n" + }, + "serverSeq": 6 + } + } + }, + { + "op": "advance", + "channel": "ahp-terminal:/t1" + }, + { + "op": "expectNotSent", + "match": { + "method": "channel/credit", + "params": { + "channel": "ahp-terminal:/t1" + } + }, + "withinMs": 150 + } + ] +} diff --git a/types/test-cases/windowed-delivery/scenarios/003-not-opted-in.json b/types/test-cases/windowed-delivery/scenarios/003-not-opted-in.json new file mode 100644 index 000000000..bedb5cb59 --- /dev/null +++ b/types/test-cases/windowed-delivery/scenarios/003-not-opted-in.json @@ -0,0 +1,25 @@ +{ + "description": "Without a window option, subscribe carries no flowControl offer.", + "windowing": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + }, + "steps": [ + { + "op": "subscribe", + "channel": "ahp-terminal:/t1" + }, + { + "op": "expectSent", + "match": { + "method": "subscribe", + "params": { + "channel": "ahp-terminal:/t1" + } + }, + "absent": [ + "params.flowControl" + ] + } + ] +} diff --git a/types/test-cases/windowed-delivery/scenarios/004-per-subscription-overrides.json b/types/test-cases/windowed-delivery/scenarios/004-per-subscription-overrides.json new file mode 100644 index 000000000..c120613b3 --- /dev/null +++ b/types/test-cases/windowed-delivery/scenarios/004-per-subscription-overrides.json @@ -0,0 +1,31 @@ +{ + "description": "Per-subscription window options override the connection defaults in the offer.", + "windowing": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + }, + "steps": [ + { + "op": "subscribe", + "channel": "ahp-terminal:/t1", + "window": { + "windowBytes": 4096 + } + }, + { + "op": "expectSent", + "match": { + "method": "subscribe", + "params": { + "channel": "ahp-terminal:/t1", + "flowControl": { + "receive": { + "windowBytes": 4096, + "maximumFrameBytes": 16384 + } + } + } + } + } + ] +} diff --git a/types/test-cases/windowed-delivery/scenarios/005-result-and-data-back-to-back.json b/types/test-cases/windowed-delivery/scenarios/005-result-and-data-back-to-back.json new file mode 100644 index 000000000..8217e5122 --- /dev/null +++ b/types/test-cases/windowed-delivery/scenarios/005-result-and-data-back-to-back.json @@ -0,0 +1,109 @@ +{ + "description": "The host sends the result, the snapshot and an action without waiting. The client must count both data messages.", + "windowing": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + }, + "steps": [ + { + "op": "subscribe", + "channel": "ahp-terminal:/t1", + "window": {} + }, + { + "op": "expectSent", + "match": { + "method": "subscribe", + "params": { + "channel": "ahp-terminal:/t1", + "flowControl": { + "receive": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + } + } + } + } + }, + { + "op": "respond", + "to": "subscribe", + "result": { + "flowControl": { + "clientReceive": { + "windowBytes": 4096, + "maximumFrameBytes": 1024 + }, + "hostReceive": { + "windowBytes": 512, + "maximumFrameBytes": 1024 + } + } + } + }, + { + "op": "hostSends", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"channel/snapshot\",\"params\":{\"channel\":\"ahp-terminal:/t1\",\"snapshot\":{\"resource\":\"ahp-terminal:/t1\",\"state\":{\"title\":\"Shell\",\"content\":[],\"lifecycle\":{\"status\":\"running\"},\"claim\":{\"kind\":\"client\",\"clientId\":\"client-1\"}},\"fromSeq\":5}}}" + }, + { + "op": "hostSends", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"action\",\"params\":{\"channel\":\"ahp-terminal:/t1\",\"action\":{\"type\":\"terminal/data\",\"data\":\"Hello\\r\\n\"},\"serverSeq\":6}}" + }, + { + "op": "next", + "channel": "ahp-terminal:/t1", + "expect": { + "method": "channel/snapshot", + "params": { + "channel": "ahp-terminal:/t1", + "snapshot": { + "resource": "ahp-terminal:/t1", + "state": { + "title": "Shell", + "content": [], + "lifecycle": { + "status": "running" + }, + "claim": { + "kind": "client", + "clientId": "client-1" + } + }, + "fromSeq": 5 + } + } + }, + "expectBytes": 520 + }, + { + "op": "next", + "channel": "ahp-terminal:/t1", + "expect": { + "method": "action", + "params": { + "channel": "ahp-terminal:/t1", + "action": { + "type": "terminal/data", + "data": "Hello\r\n" + }, + "serverSeq": 6 + } + }, + "expectBytes": 284 + }, + { + "op": "advance", + "channel": "ahp-terminal:/t1" + }, + { + "op": "expectSent", + "match": { + "method": "channel/credit", + "params": { + "channel": "ahp-terminal:/t1", + "consumedBytes": 804 + } + } + } + ] +} diff --git a/types/test-cases/windowed-delivery/scenarios/006-manual-release.json b/types/test-cases/windowed-delivery/scenarios/006-manual-release.json new file mode 100644 index 000000000..9e6df2fa1 --- /dev/null +++ b/types/test-cases/windowed-delivery/scenarios/006-manual-release.json @@ -0,0 +1,150 @@ +{ + "description": "In manual mode, credit follows release() in any order, never the act of receiving.", + "windowing": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + }, + "steps": [ + { + "op": "subscribe", + "channel": "ahp-terminal:/t1", + "window": {}, + "receive": "manual" + }, + { + "op": "expectSent", + "match": { + "method": "subscribe", + "params": { + "channel": "ahp-terminal:/t1", + "flowControl": { + "receive": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + } + } + } + } + }, + { + "op": "respond", + "to": "subscribe", + "result": { + "flowControl": { + "clientReceive": { + "windowBytes": 4096, + "maximumFrameBytes": 1024 + }, + "hostReceive": { + "windowBytes": 512, + "maximumFrameBytes": 1024 + } + } + } + }, + { + "op": "hostSends", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"channel/snapshot\",\"params\":{\"channel\":\"ahp-terminal:/t1\",\"snapshot\":{\"resource\":\"ahp-terminal:/t1\",\"state\":{\"title\":\"Shell\",\"content\":[],\"lifecycle\":{\"status\":\"running\"},\"claim\":{\"kind\":\"client\",\"clientId\":\"client-1\"}},\"fromSeq\":5}}}" + }, + { + "op": "hostSends", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"action\",\"params\":{\"channel\":\"ahp-terminal:/t1\",\"action\":{\"type\":\"terminal/data\",\"data\":\"Hello\\r\\n\"},\"serverSeq\":6}}" + }, + { + "op": "next", + "channel": "ahp-terminal:/t1", + "expect": { + "method": "channel/snapshot", + "params": { + "channel": "ahp-terminal:/t1", + "snapshot": { + "resource": "ahp-terminal:/t1", + "state": { + "title": "Shell", + "content": [], + "lifecycle": { + "status": "running" + }, + "claim": { + "kind": "client", + "clientId": "client-1" + } + }, + "fromSeq": 5 + } + } + }, + "expectBytes": 520 + }, + { + "op": "next", + "channel": "ahp-terminal:/t1", + "expect": { + "method": "action", + "params": { + "channel": "ahp-terminal:/t1", + "action": { + "type": "terminal/data", + "data": "Hello\r\n" + }, + "serverSeq": 6 + } + }, + "expectBytes": 284 + }, + { + "op": "expectNotSent", + "match": { + "method": "channel/credit", + "params": { + "channel": "ahp-terminal:/t1" + } + }, + "withinMs": 100 + }, + { + "op": "expectStatus", + "channel": "ahp-terminal:/t1", + "match": { + "bufferedReceiveBytes": 804 + } + }, + { + "op": "release", + "channel": "ahp-terminal:/t1", + "index": 1 + }, + { + "op": "expectSent", + "match": { + "method": "channel/credit", + "params": { + "channel": "ahp-terminal:/t1", + "consumedBytes": 284 + } + } + }, + { + "op": "release", + "channel": "ahp-terminal:/t1", + "index": 0 + }, + { + "op": "expectSent", + "match": { + "method": "channel/credit", + "params": { + "channel": "ahp-terminal:/t1", + "consumedBytes": 804 + } + } + }, + { + "op": "expectStatus", + "channel": "ahp-terminal:/t1", + "match": { + "bufferedReceiveBytes": 0 + } + } + ] +} diff --git a/types/test-cases/windowed-delivery/scenarios/007-spec-fragments.json b/types/test-cases/windowed-delivery/scenarios/007-spec-fragments.json new file mode 100644 index 000000000..722d55b87 --- /dev/null +++ b/types/test-cases/windowed-delivery/scenarios/007-spec-fragments.json @@ -0,0 +1,138 @@ +{ + "description": "The spec example: the 284-byte action arrives as fragments of 360 and 374 bytes. The non-final fragment is credited once stored; the joined message is not charged again.", + "windowing": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + }, + "steps": [ + { + "op": "subscribe", + "channel": "ahp-terminal:/t1", + "window": {}, + "receive": "auto" + }, + { + "op": "expectSent", + "match": { + "method": "subscribe", + "params": { + "channel": "ahp-terminal:/t1", + "flowControl": { + "receive": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + } + } + } + } + }, + { + "op": "respond", + "to": "subscribe", + "result": { + "flowControl": { + "clientReceive": { + "windowBytes": 512, + "maximumFrameBytes": 1024 + }, + "hostReceive": { + "windowBytes": 512, + "maximumFrameBytes": 1024 + } + } + } + }, + { + "op": "hostSends", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"channel/snapshot\",\"params\":{\"channel\":\"ahp-terminal:/t1\",\"snapshot\":{\"resource\":\"ahp-terminal:/t1\",\"state\":{\"title\":\"Shell\",\"content\":[],\"lifecycle\":{\"status\":\"running\"},\"claim\":{\"kind\":\"client\",\"clientId\":\"client-1\"}},\"fromSeq\":5}}}" + }, + { + "op": "next", + "channel": "ahp-terminal:/t1", + "expect": { + "method": "channel/snapshot", + "params": { + "channel": "ahp-terminal:/t1", + "snapshot": { + "resource": "ahp-terminal:/t1", + "state": { + "title": "Shell", + "content": [], + "lifecycle": { + "status": "running" + }, + "claim": { + "kind": "client", + "clientId": "client-1" + } + }, + "fromSeq": 5 + } + } + }, + "expectBytes": 520 + }, + { + "op": "advance", + "channel": "ahp-terminal:/t1" + }, + { + "op": "expectSent", + "match": { + "method": "channel/credit", + "params": { + "channel": "ahp-terminal:/t1", + "consumedBytes": 520 + } + } + }, + { + "op": "hostSends", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"channel/frame\",\"params\":{\"channel\":\"ahp-terminal:/t1\",\"data\":\"{\\\"jsonrpc\\\":\\\"2.0\\\",\\\"method\\\":\\\"action\\\",\\\"params\\\":{\\\"channel\\\":\\\"ahp-terminal:/t1\\\",\"}}" + }, + { + "op": "expectSent", + "match": { + "method": "channel/credit", + "params": { + "channel": "ahp-terminal:/t1", + "consumedBytes": 880 + } + } + }, + { + "op": "hostSends", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"channel/frame\",\"params\":{\"channel\":\"ahp-terminal:/t1\",\"data\":\"\\\"action\\\":{\\\"type\\\":\\\"terminal/data\\\",\\\"data\\\":\\\"Hello\\\\r\\\\n\\\"},\\\"serverSeq\\\":6}}\",\"final\":true}}" + }, + { + "op": "next", + "channel": "ahp-terminal:/t1", + "expect": { + "method": "action", + "params": { + "channel": "ahp-terminal:/t1", + "action": { + "type": "terminal/data", + "data": "Hello\r\n" + }, + "serverSeq": 6 + } + }, + "expectBytes": 734 + }, + { + "op": "advance", + "channel": "ahp-terminal:/t1" + }, + { + "op": "expectSent", + "match": { + "method": "channel/credit", + "params": { + "channel": "ahp-terminal:/t1", + "consumedBytes": 1254 + } + } + } + ] +} diff --git a/types/test-cases/windowed-delivery/scenarios/008-fragmented-larger-than-window.json b/types/test-cases/windowed-delivery/scenarios/008-fragmented-larger-than-window.json new file mode 100644 index 000000000..785a7e5c6 --- /dev/null +++ b/types/test-cases/windowed-delivery/scenarios/008-fragmented-larger-than-window.json @@ -0,0 +1,152 @@ +{ + "description": "A message larger than the 512-byte window arrives in several fragments containing non-ASCII text and a surrogate pair. Each non-final fragment is credited when stored, so the host can make progress.", + "windowing": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + }, + "steps": [ + { + "op": "subscribe", + "channel": "ahp-terminal:/t1", + "window": {}, + "receive": "auto" + }, + { + "op": "expectSent", + "match": { + "method": "subscribe", + "params": { + "channel": "ahp-terminal:/t1", + "flowControl": { + "receive": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + } + } + } + } + }, + { + "op": "respond", + "to": "subscribe", + "result": { + "flowControl": { + "clientReceive": { + "windowBytes": 512, + "maximumFrameBytes": 1024 + }, + "hostReceive": { + "windowBytes": 512, + "maximumFrameBytes": 1024 + } + } + } + }, + { + "op": "hostSends", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"channel/snapshot\",\"params\":{\"channel\":\"ahp-terminal:/t1\",\"snapshot\":{\"resource\":\"ahp-terminal:/t1\",\"state\":{\"title\":\"Shell\",\"content\":[],\"lifecycle\":{\"status\":\"running\"},\"claim\":{\"kind\":\"client\",\"clientId\":\"client-1\"}},\"fromSeq\":5}}}" + }, + { + "op": "next", + "channel": "ahp-terminal:/t1", + "expect": { + "method": "channel/snapshot", + "params": { + "channel": "ahp-terminal:/t1", + "snapshot": { + "resource": "ahp-terminal:/t1", + "state": { + "title": "Shell", + "content": [], + "lifecycle": { + "status": "running" + }, + "claim": { + "kind": "client", + "clientId": "client-1" + } + }, + "fromSeq": 5 + } + } + }, + "expectBytes": 520 + }, + { + "op": "advance", + "channel": "ahp-terminal:/t1" + }, + { + "op": "expectSent", + "match": { + "method": "channel/credit", + "params": { + "channel": "ahp-terminal:/t1", + "consumedBytes": 520 + } + } + }, + { + "op": "hostSends", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"channel/frame\",\"params\":{\"channel\":\"ahp-terminal:/t1\",\"data\":\"{\\\"jsonrpc\\\":\\\"2.0\\\",\\\"method\\\":\\\"action\\\",\\\"params\\\":{\\\"channel\\\":\\\"ahp-terminal:/t1\\\",\\\"action\\\":{\\\"type\\\":\\\"terminal/data\\\",\\\"data\\\":\\\"Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grü\"}}" + }, + { + "op": "expectSent", + "match": { + "method": "channel/credit", + "params": { + "channel": "ahp-terminal:/t1", + "consumedBytes": 1544 + } + } + }, + { + "op": "hostSends", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"channel/frame\",\"params\":{\"channel\":\"ahp-terminal:/t1\",\"data\":\"ße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 \"}}" + }, + { + "op": "expectSent", + "match": { + "method": "channel/credit", + "params": { + "channel": "ahp-terminal:/t1", + "consumedBytes": 2568 + } + } + }, + { + "op": "hostSends", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"channel/frame\",\"params\":{\"channel\":\"ahp-terminal:/t1\",\"data\":\"Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 end\\\\r\\\\n\\\"},\\\"serverSeq\\\":6}}\",\"final\":true}}" + }, + { + "op": "next", + "channel": "ahp-terminal:/t1", + "expect": { + "method": "action", + "params": { + "channel": "ahp-terminal:/t1", + "action": { + "type": "terminal/data", + "data": "Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 Grüße 😀 end\r\n" + }, + "serverSeq": 6 + } + }, + "expectBytes": 2534 + }, + { + "op": "advance", + "channel": "ahp-terminal:/t1" + }, + { + "op": "expectSent", + "match": { + "method": "channel/credit", + "params": { + "channel": "ahp-terminal:/t1", + "consumedBytes": 3054 + } + } + } + ] +} diff --git a/types/test-cases/windowed-delivery/scenarios/009-unknown-counts-exempt-does-not.json b/types/test-cases/windowed-delivery/scenarios/009-unknown-counts-exempt-does-not.json new file mode 100644 index 000000000..0a4d08f50 --- /dev/null +++ b/types/test-cases/windowed-delivery/scenarios/009-unknown-counts-exempt-does-not.json @@ -0,0 +1,124 @@ +{ + "description": "An unknown notification for the channel counts against the window and is delivered. A host channel/credit is a control message and does not count.", + "windowing": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + }, + "steps": [ + { + "op": "subscribe", + "channel": "ahp-terminal:/t1", + "window": {}, + "receive": "auto" + }, + { + "op": "expectSent", + "match": { + "method": "subscribe", + "params": { + "channel": "ahp-terminal:/t1", + "flowControl": { + "receive": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + } + } + } + } + }, + { + "op": "respond", + "to": "subscribe", + "result": { + "flowControl": { + "clientReceive": { + "windowBytes": 512, + "maximumFrameBytes": 1024 + }, + "hostReceive": { + "windowBytes": 512, + "maximumFrameBytes": 1024 + } + } + } + }, + { + "op": "hostSends", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"channel/snapshot\",\"params\":{\"channel\":\"ahp-terminal:/t1\",\"snapshot\":{\"resource\":\"ahp-terminal:/t1\",\"state\":{\"title\":\"Shell\",\"content\":[],\"lifecycle\":{\"status\":\"running\"},\"claim\":{\"kind\":\"client\",\"clientId\":\"client-1\"}},\"fromSeq\":5}}}" + }, + { + "op": "next", + "channel": "ahp-terminal:/t1", + "expect": { + "method": "channel/snapshot", + "params": { + "channel": "ahp-terminal:/t1", + "snapshot": { + "resource": "ahp-terminal:/t1", + "state": { + "title": "Shell", + "content": [], + "lifecycle": { + "status": "running" + }, + "claim": { + "kind": "client", + "clientId": "client-1" + } + }, + "fromSeq": 5 + } + } + }, + "expectBytes": 520 + }, + { + "op": "advance", + "channel": "ahp-terminal:/t1" + }, + { + "op": "expectSent", + "match": { + "method": "channel/credit", + "params": { + "channel": "ahp-terminal:/t1", + "consumedBytes": 520 + } + } + }, + { + "op": "hostSends", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"channel/credit\",\"params\":{\"channel\":\"ahp-terminal:/t1\",\"consumedBytes\":0}}" + }, + { + "op": "hostSends", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"ext/ping\",\"params\":{\"channel\":\"ahp-terminal:/t1\",\"n\":1}}" + }, + { + "op": "next", + "channel": "ahp-terminal:/t1", + "expect": { + "method": "ext/ping", + "params": { + "channel": "ahp-terminal:/t1", + "n": 1 + } + }, + "expectBytes": 166 + }, + { + "op": "advance", + "channel": "ahp-terminal:/t1" + }, + { + "op": "expectSent", + "match": { + "method": "channel/credit", + "params": { + "channel": "ahp-terminal:/t1", + "consumedBytes": 686 + } + } + } + ] +} diff --git a/types/test-cases/windowed-delivery/scenarios/010-empty-and-unknown-snapshot-state.json b/types/test-cases/windowed-delivery/scenarios/010-empty-and-unknown-snapshot-state.json new file mode 100644 index 000000000..da8000a55 --- /dev/null +++ b/types/test-cases/windowed-delivery/scenarios/010-empty-and-unknown-snapshot-state.json @@ -0,0 +1,83 @@ +{ + "description": "Snapshots with an empty or unrecognized state are delivered and credited, not rejected.", + "windowing": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + }, + "steps": [ + { + "op": "subscribe", + "channel": "ahp-terminal:/t1", + "window": {} + }, + { + "op": "expectSent", + "match": { + "method": "subscribe", + "params": { + "channel": "ahp-terminal:/t1", + "flowControl": { + "receive": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + } + } + } + } + }, + { + "op": "respond", + "to": "subscribe", + "result": { + "flowControl": { + "clientReceive": { + "windowBytes": 4096, + "maximumFrameBytes": 1024 + }, + "hostReceive": { + "windowBytes": 512, + "maximumFrameBytes": 1024 + } + } + } + }, + { + "op": "hostSends", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"channel/snapshot\",\"params\":{\"channel\":\"ahp-terminal:/t1\",\"snapshot\":{\"resource\":\"ahp-terminal:/t1\",\"state\":{},\"fromSeq\":9}}}" + }, + { + "op": "hostSends", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"channel/snapshot\",\"params\":{\"channel\":\"ahp-terminal:/t1\",\"snapshot\":{\"resource\":\"ahp-terminal:/t1\",\"state\":{\"somethingNew\":{\"a\":1}},\"fromSeq\":9}}}" + }, + { + "op": "next", + "channel": "ahp-terminal:/t1", + "expect": { + "method": "channel/snapshot" + }, + "expectBytes": 302 + }, + { + "op": "next", + "channel": "ahp-terminal:/t1", + "expect": { + "method": "channel/snapshot" + }, + "expectBytes": 346 + }, + { + "op": "advance", + "channel": "ahp-terminal:/t1" + }, + { + "op": "expectSent", + "match": { + "method": "channel/credit", + "params": { + "channel": "ahp-terminal:/t1", + "consumedBytes": 648 + } + } + } + ] +} diff --git a/types/test-cases/windowed-delivery/scenarios/011-coalesce-terminal-data.json b/types/test-cases/windowed-delivery/scenarios/011-coalesce-terminal-data.json new file mode 100644 index 000000000..5ccfcb9a9 --- /dev/null +++ b/types/test-cases/windowed-delivery/scenarios/011-coalesce-terminal-data.json @@ -0,0 +1,124 @@ +{ + "description": "Adjacent queued terminal/data actions are concatenated into one delivery that carries the last serverSeq and the sum of the original sizes.", + "windowing": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + }, + "steps": [ + { + "op": "subscribe", + "channel": "ahp-terminal:/t1", + "window": {} + }, + { + "op": "expectSent", + "match": { + "method": "subscribe", + "params": { + "channel": "ahp-terminal:/t1", + "flowControl": { + "receive": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + } + } + } + } + }, + { + "op": "respond", + "to": "subscribe", + "result": { + "flowControl": { + "clientReceive": { + "windowBytes": 65536, + "maximumFrameBytes": 16384 + }, + "hostReceive": { + "windowBytes": 65536, + "maximumFrameBytes": 16384 + } + } + } + }, + { + "op": "hostSends", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"channel/snapshot\",\"params\":{\"channel\":\"ahp-terminal:/t1\",\"snapshot\":{\"resource\":\"ahp-terminal:/t1\",\"state\":{\"title\":\"Shell\",\"content\":[],\"lifecycle\":{\"status\":\"running\"},\"claim\":{\"kind\":\"client\",\"clientId\":\"client-1\"}},\"fromSeq\":5}}}" + }, + { + "op": "next", + "channel": "ahp-terminal:/t1", + "expect": { + "method": "channel/snapshot", + "params": { + "channel": "ahp-terminal:/t1", + "snapshot": { + "resource": "ahp-terminal:/t1", + "state": { + "title": "Shell", + "content": [], + "lifecycle": { + "status": "running" + }, + "claim": { + "kind": "client", + "clientId": "client-1" + } + }, + "fromSeq": 5 + } + } + }, + "expectBytes": 520 + }, + { + "op": "hostSends", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"action\",\"params\":{\"channel\":\"ahp-terminal:/t1\",\"action\":{\"type\":\"terminal/data\",\"data\":\"ab\"},\"serverSeq\":6}}" + }, + { + "op": "hostSends", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"action\",\"params\":{\"channel\":\"ahp-terminal:/t1\",\"action\":{\"type\":\"terminal/data\",\"data\":\"cd\"},\"serverSeq\":7}}" + }, + { + "op": "hostSends", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"action\",\"params\":{\"channel\":\"ahp-terminal:/t1\",\"action\":{\"type\":\"terminal/data\",\"data\":\"ef\\r\\n\"},\"serverSeq\":8}}" + }, + { + "op": "expectStatus", + "channel": "ahp-terminal:/t1", + "match": { + "bufferedReceiveBytes": 1338 + } + }, + { + "op": "next", + "channel": "ahp-terminal:/t1", + "expect": { + "method": "action", + "params": { + "channel": "ahp-terminal:/t1", + "action": { + "type": "terminal/data", + "data": "abcdef\r\n" + }, + "serverSeq": 8 + } + }, + "expectBytes": 818 + }, + { + "op": "advance", + "channel": "ahp-terminal:/t1" + }, + { + "op": "expectSent", + "match": { + "method": "channel/credit", + "params": { + "channel": "ahp-terminal:/t1", + "consumedBytes": 1338 + } + } + } + ] +} diff --git a/types/test-cases/windowed-delivery/scenarios/012-coalescing-disabled.json b/types/test-cases/windowed-delivery/scenarios/012-coalescing-disabled.json new file mode 100644 index 000000000..0e85d8100 --- /dev/null +++ b/types/test-cases/windowed-delivery/scenarios/012-coalescing-disabled.json @@ -0,0 +1,123 @@ +{ + "description": "With coalescing turned off, each action is delivered separately.", + "windowing": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + }, + "steps": [ + { + "op": "subscribe", + "channel": "ahp-terminal:/t1", + "window": {}, + "coalescing": false + }, + { + "op": "expectSent", + "match": { + "method": "subscribe", + "params": { + "channel": "ahp-terminal:/t1", + "flowControl": { + "receive": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + } + } + } + } + }, + { + "op": "respond", + "to": "subscribe", + "result": { + "flowControl": { + "clientReceive": { + "windowBytes": 65536, + "maximumFrameBytes": 16384 + }, + "hostReceive": { + "windowBytes": 65536, + "maximumFrameBytes": 16384 + } + } + } + }, + { + "op": "hostSends", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"channel/snapshot\",\"params\":{\"channel\":\"ahp-terminal:/t1\",\"snapshot\":{\"resource\":\"ahp-terminal:/t1\",\"state\":{\"title\":\"Shell\",\"content\":[],\"lifecycle\":{\"status\":\"running\"},\"claim\":{\"kind\":\"client\",\"clientId\":\"client-1\"}},\"fromSeq\":5}}}" + }, + { + "op": "next", + "channel": "ahp-terminal:/t1", + "expect": { + "method": "channel/snapshot", + "params": { + "channel": "ahp-terminal:/t1", + "snapshot": { + "resource": "ahp-terminal:/t1", + "state": { + "title": "Shell", + "content": [], + "lifecycle": { + "status": "running" + }, + "claim": { + "kind": "client", + "clientId": "client-1" + } + }, + "fromSeq": 5 + } + } + }, + "expectBytes": 520 + }, + { + "op": "hostSends", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"action\",\"params\":{\"channel\":\"ahp-terminal:/t1\",\"action\":{\"type\":\"terminal/data\",\"data\":\"ab\"},\"serverSeq\":6}}" + }, + { + "op": "hostSends", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"action\",\"params\":{\"channel\":\"ahp-terminal:/t1\",\"action\":{\"type\":\"terminal/data\",\"data\":\"cd\"},\"serverSeq\":7}}" + }, + { + "op": "expectStatus", + "channel": "ahp-terminal:/t1", + "match": { + "bufferedReceiveBytes": 1060 + } + }, + { + "op": "next", + "channel": "ahp-terminal:/t1", + "expect": { + "method": "action", + "params": { + "channel": "ahp-terminal:/t1", + "action": { + "type": "terminal/data", + "data": "ab" + }, + "serverSeq": 6 + } + }, + "expectBytes": 270 + }, + { + "op": "next", + "channel": "ahp-terminal:/t1", + "expect": { + "method": "action", + "params": { + "channel": "ahp-terminal:/t1", + "action": { + "type": "terminal/data", + "data": "cd" + }, + "serverSeq": 7 + } + }, + "expectBytes": 270 + } + ] +} diff --git a/types/test-cases/windowed-delivery/scenarios/013-no-coalesce-across-origin.json b/types/test-cases/windowed-delivery/scenarios/013-no-coalesce-across-origin.json new file mode 100644 index 000000000..90c2eb996 --- /dev/null +++ b/types/test-cases/windowed-delivery/scenarios/013-no-coalesce-across-origin.json @@ -0,0 +1,146 @@ +{ + "description": "An action with an origin is never merged, so its neighbors are not adjacent and stay separate.", + "windowing": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + }, + "steps": [ + { + "op": "subscribe", + "channel": "ahp-terminal:/t1", + "window": {} + }, + { + "op": "expectSent", + "match": { + "method": "subscribe", + "params": { + "channel": "ahp-terminal:/t1", + "flowControl": { + "receive": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + } + } + } + } + }, + { + "op": "respond", + "to": "subscribe", + "result": { + "flowControl": { + "clientReceive": { + "windowBytes": 65536, + "maximumFrameBytes": 16384 + }, + "hostReceive": { + "windowBytes": 65536, + "maximumFrameBytes": 16384 + } + } + } + }, + { + "op": "hostSends", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"channel/snapshot\",\"params\":{\"channel\":\"ahp-terminal:/t1\",\"snapshot\":{\"resource\":\"ahp-terminal:/t1\",\"state\":{\"title\":\"Shell\",\"content\":[],\"lifecycle\":{\"status\":\"running\"},\"claim\":{\"kind\":\"client\",\"clientId\":\"client-1\"}},\"fromSeq\":5}}}" + }, + { + "op": "next", + "channel": "ahp-terminal:/t1", + "expect": { + "method": "channel/snapshot", + "params": { + "channel": "ahp-terminal:/t1", + "snapshot": { + "resource": "ahp-terminal:/t1", + "state": { + "title": "Shell", + "content": [], + "lifecycle": { + "status": "running" + }, + "claim": { + "kind": "client", + "clientId": "client-1" + } + }, + "fromSeq": 5 + } + } + }, + "expectBytes": 520 + }, + { + "op": "hostSends", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"action\",\"params\":{\"channel\":\"ahp-terminal:/t1\",\"action\":{\"type\":\"terminal/data\",\"data\":\"ab\"},\"serverSeq\":6}}" + }, + { + "op": "hostSends", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"action\",\"params\":{\"channel\":\"ahp-terminal:/t1\",\"action\":{\"type\":\"terminal/data\",\"data\":\"XY\"},\"serverSeq\":7,\"origin\":{\"clientId\":\"client-1\",\"clientSeq\":3}}}" + }, + { + "op": "hostSends", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"action\",\"params\":{\"channel\":\"ahp-terminal:/t1\",\"action\":{\"type\":\"terminal/data\",\"data\":\"ef\\r\\n\"},\"serverSeq\":8}}" + }, + { + "op": "expectStatus", + "channel": "ahp-terminal:/t1", + "match": { + "bufferedReceiveBytes": 1432 + } + }, + { + "op": "next", + "channel": "ahp-terminal:/t1", + "expect": { + "method": "action", + "params": { + "channel": "ahp-terminal:/t1", + "action": { + "type": "terminal/data", + "data": "ab" + }, + "serverSeq": 6 + } + }, + "expectBytes": 270 + }, + { + "op": "next", + "channel": "ahp-terminal:/t1", + "expect": { + "method": "action", + "params": { + "channel": "ahp-terminal:/t1", + "action": { + "type": "terminal/data", + "data": "XY" + }, + "serverSeq": 7, + "origin": { + "clientId": "client-1", + "clientSeq": 3 + } + } + }, + "expectBytes": 364 + }, + { + "op": "next", + "channel": "ahp-terminal:/t1", + "expect": { + "method": "action", + "params": { + "channel": "ahp-terminal:/t1", + "action": { + "type": "terminal/data", + "data": "ef\r\n" + }, + "serverSeq": 8 + } + }, + "expectBytes": 278 + } + ] +} diff --git a/types/test-cases/windowed-delivery/scenarios/014-coalesce-chat-delta.json b/types/test-cases/windowed-delivery/scenarios/014-coalesce-chat-delta.json new file mode 100644 index 000000000..6cd7564cc --- /dev/null +++ b/types/test-cases/windowed-delivery/scenarios/014-coalesce-chat-delta.json @@ -0,0 +1,110 @@ +{ + "description": "Adjacent chat/delta actions for the same turn and part are concatenated; a different part is not merged.", + "windowing": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + }, + "steps": [ + { + "op": "subscribe", + "channel": "ahp-chat:/c1", + "window": {} + }, + { + "op": "expectSent", + "match": { + "method": "subscribe", + "params": { + "channel": "ahp-chat:/c1", + "flowControl": { + "receive": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + } + } + } + } + }, + { + "op": "respond", + "to": "subscribe", + "result": { + "flowControl": { + "clientReceive": { + "windowBytes": 65536, + "maximumFrameBytes": 16384 + }, + "hostReceive": { + "windowBytes": 65536, + "maximumFrameBytes": 16384 + } + } + } + }, + { + "op": "hostSends", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"channel/snapshot\",\"params\":{\"channel\":\"ahp-chat:/c1\",\"snapshot\":{\"resource\":\"ahp-chat:/c1\",\"state\":{\"resource\":\"ahp-chat:/c1\",\"title\":\"s\",\"status\":1,\"modifiedAt\":\"2026-01-01T00:00:00Z\",\"turns\":[]},\"fromSeq\":1}}}" + }, + { + "op": "next", + "channel": "ahp-chat:/c1", + "expect": { + "method": "channel/snapshot" + }, + "expectBytes": 476 + }, + { + "op": "hostSends", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"action\",\"params\":{\"channel\":\"ahp-chat:/c1\",\"action\":{\"type\":\"chat/delta\",\"turnId\":\"turn-1\",\"partId\":\"p1\",\"content\":\"Hel\"},\"serverSeq\":2}}" + }, + { + "op": "hostSends", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"action\",\"params\":{\"channel\":\"ahp-chat:/c1\",\"action\":{\"type\":\"chat/delta\",\"turnId\":\"turn-1\",\"partId\":\"p1\",\"content\":\"lo\"},\"serverSeq\":3}}" + }, + { + "op": "hostSends", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"action\",\"params\":{\"channel\":\"ahp-chat:/c1\",\"action\":{\"type\":\"chat/delta\",\"turnId\":\"turn-1\",\"partId\":\"p2\",\"content\":\"!\"},\"serverSeq\":4}}" + }, + { + "op": "expectStatus", + "channel": "ahp-chat:/c1", + "match": { + "bufferedReceiveBytes": 1454 + } + }, + { + "op": "next", + "channel": "ahp-chat:/c1", + "expect": { + "method": "action", + "params": { + "action": { + "type": "chat/delta", + "partId": "p1", + "content": "Hello" + }, + "serverSeq": 3 + } + }, + "expectBytes": 654 + }, + { + "op": "next", + "channel": "ahp-chat:/c1", + "expect": { + "method": "action", + "params": { + "channel": "ahp-chat:/c1", + "action": { + "type": "chat/delta", + "turnId": "turn-1", + "partId": "p2", + "content": "!" + }, + "serverSeq": 4 + } + }, + "expectBytes": 324 + } + ] +} diff --git a/types/test-cases/windowed-delivery/scenarios/015-send-blocks-on-window.json b/types/test-cases/windowed-delivery/scenarios/015-send-blocks-on-window.json new file mode 100644 index 000000000..8b0926c84 --- /dev/null +++ b/types/test-cases/windowed-delivery/scenarios/015-send-blocks-on-window.json @@ -0,0 +1,151 @@ +{ + "description": "Data sent by the consumer waits when the host window is full and resumes on credit. Each message is about 750 bytes against a 1000-byte host window.", + "windowing": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + }, + "steps": [ + { + "op": "subscribe", + "channel": "ahp-terminal:/t1", + "window": {} + }, + { + "op": "expectSent", + "match": { + "method": "subscribe", + "params": { + "channel": "ahp-terminal:/t1", + "flowControl": { + "receive": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + } + } + } + } + }, + { + "op": "respond", + "to": "subscribe", + "result": { + "flowControl": { + "clientReceive": { + "windowBytes": 65536, + "maximumFrameBytes": 16384 + }, + "hostReceive": { + "windowBytes": 1000, + "maximumFrameBytes": 16384 + } + } + } + }, + { + "op": "expectStatus", + "channel": "ahp-terminal:/t1", + "match": { + "writable": true + } + }, + { + "op": "send", + "channel": "ahp-terminal:/t1", + "method": "dispatchAction", + "params": { + "channel": "ahp-terminal:/t1", + "clientSeq": 1, + "action": { + "type": "terminal/input", + "data": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" + } + } + }, + { + "op": "send", + "channel": "ahp-terminal:/t1", + "method": "dispatchAction", + "params": { + "channel": "ahp-terminal:/t1", + "clientSeq": 2, + "action": { + "type": "terminal/input", + "data": "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb" + } + } + }, + { + "op": "send", + "channel": "ahp-terminal:/t1", + "method": "dispatchAction", + "params": { + "channel": "ahp-terminal:/t1", + "clientSeq": 3, + "action": { + "type": "terminal/input", + "data": "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc" + } + } + }, + { + "op": "expectSent", + "match": { + "method": "dispatchAction", + "params": { + "channel": "ahp-terminal:/t1", + "clientSeq": 1, + "action": { + "type": "terminal/input", + "data": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" + } + } + } + }, + { + "op": "expectSent", + "match": { + "method": "dispatchAction", + "params": { + "channel": "ahp-terminal:/t1", + "clientSeq": 2, + "action": { + "type": "terminal/input", + "data": "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb" + } + } + } + }, + { + "op": "expectNotSent", + "match": { + "method": "dispatchAction" + }, + "withinMs": 150 + }, + { + "op": "expectStatus", + "channel": "ahp-terminal:/t1", + "match": { + "writable": false + } + }, + { + "op": "hostCreditsAll", + "channel": "ahp-terminal:/t1" + }, + { + "op": "expectSent", + "match": { + "method": "dispatchAction", + "params": { + "channel": "ahp-terminal:/t1", + "clientSeq": 3, + "action": { + "type": "terminal/input", + "data": "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc" + } + } + } + } + ] +} diff --git a/types/test-cases/windowed-delivery/scenarios/016-outbound-fragmenting.json b/types/test-cases/windowed-delivery/scenarios/016-outbound-fragmenting.json new file mode 100644 index 000000000..87e88501c --- /dev/null +++ b/types/test-cases/windowed-delivery/scenarios/016-outbound-fragmenting.json @@ -0,0 +1,73 @@ +{ + "description": "A message larger than the host maximumFrameBytes is sent as channel/frame fragments. Each fragment fits the limit and the joined text is the original message.", + "windowing": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + }, + "steps": [ + { + "op": "subscribe", + "channel": "ahp-terminal:/t1", + "window": {} + }, + { + "op": "expectSent", + "match": { + "method": "subscribe", + "params": { + "channel": "ahp-terminal:/t1", + "flowControl": { + "receive": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + } + } + } + } + }, + { + "op": "respond", + "to": "subscribe", + "result": { + "flowControl": { + "clientReceive": { + "windowBytes": 65536, + "maximumFrameBytes": 16384 + }, + "hostReceive": { + "windowBytes": 1048576, + "maximumFrameBytes": 1024 + } + } + } + }, + { + "op": "send", + "channel": "ahp-terminal:/t1", + "method": "dispatchAction", + "params": { + "channel": "ahp-terminal:/t1", + "clientSeq": 1, + "action": { + "type": "terminal/input", + "data": "ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 " + } + } + }, + { + "op": "expectSent", + "reassemble": true, + "match": { + "method": "dispatchAction", + "params": { + "channel": "ahp-terminal:/t1", + "clientSeq": 1, + "action": { + "type": "terminal/input", + "data": "ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 ü\"\\x😀 " + } + } + } + } + ] +} diff --git a/types/test-cases/windowed-delivery/scenarios/017-tcp-data-send-coalescing.json b/types/test-cases/windowed-delivery/scenarios/017-tcp-data-send-coalescing.json new file mode 100644 index 000000000..7b2fb7fa7 --- /dev/null +++ b/types/test-cases/windowed-delivery/scenarios/017-tcp-data-send-coalescing.json @@ -0,0 +1,110 @@ +{ + "description": "While the host window is full, queued tcp/data sends are merged by concatenating the decoded bytes, not the base64 text.", + "windowing": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + }, + "steps": [ + { + "op": "subscribe", + "channel": "ahp-tcp:/t1", + "window": {} + }, + { + "op": "expectSent", + "match": { + "method": "subscribe", + "params": { + "channel": "ahp-tcp:/t1", + "flowControl": { + "receive": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + } + } + } + } + }, + { + "op": "respond", + "to": "subscribe", + "result": { + "flowControl": { + "clientReceive": { + "windowBytes": 65536, + "maximumFrameBytes": 16384 + }, + "hostReceive": { + "windowBytes": 512, + "maximumFrameBytes": 16384 + } + } + } + }, + { + "op": "send", + "channel": "ahp-tcp:/t1", + "method": "tcp/data", + "params": { + "channel": "ahp-tcp:/t1", + "data": "eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4" + } + }, + { + "op": "expectSent", + "match": { + "method": "tcp/data", + "params": { + "channel": "ahp-tcp:/t1", + "data": "eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4" + } + } + }, + { + "op": "send", + "channel": "ahp-tcp:/t1", + "method": "tcp/data", + "params": { + "channel": "ahp-tcp:/t1", + "data": "aGVs" + } + }, + { + "op": "send", + "channel": "ahp-tcp:/t1", + "method": "tcp/data", + "params": { + "channel": "ahp-tcp:/t1", + "data": "bG8=" + } + }, + { + "op": "expectNotSent", + "match": { + "method": "tcp/data" + }, + "withinMs": 150 + }, + { + "op": "hostCreditsAll", + "channel": "ahp-tcp:/t1" + }, + { + "op": "expectSent", + "match": { + "method": "tcp/data", + "params": { + "channel": "ahp-tcp:/t1", + "data": "aGVsbG8=" + } + } + }, + { + "op": "expectNotSent", + "match": { + "method": "tcp/data" + }, + "withinMs": 100 + } + ] +} diff --git a/types/test-cases/windowed-delivery/scenarios/018-resubscribe-keeps-counters.json b/types/test-cases/windowed-delivery/scenarios/018-resubscribe-keeps-counters.json new file mode 100644 index 000000000..72533861c --- /dev/null +++ b/types/test-cases/windowed-delivery/scenarios/018-resubscribe-keeps-counters.json @@ -0,0 +1,183 @@ +{ + "description": "A second subscribe on a windowed subscription changes its limits; counters continue, so credit after the new snapshot is cumulative.", + "windowing": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + }, + "steps": [ + { + "op": "subscribe", + "channel": "ahp-terminal:/t1", + "window": {}, + "receive": "auto" + }, + { + "op": "expectSent", + "match": { + "method": "subscribe", + "params": { + "channel": "ahp-terminal:/t1", + "flowControl": { + "receive": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + } + } + } + } + }, + { + "op": "respond", + "to": "subscribe", + "result": { + "flowControl": { + "clientReceive": { + "windowBytes": 512, + "maximumFrameBytes": 1024 + }, + "hostReceive": { + "windowBytes": 512, + "maximumFrameBytes": 1024 + } + } + } + }, + { + "op": "hostSends", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"channel/snapshot\",\"params\":{\"channel\":\"ahp-terminal:/t1\",\"snapshot\":{\"resource\":\"ahp-terminal:/t1\",\"state\":{\"title\":\"Shell\",\"content\":[],\"lifecycle\":{\"status\":\"running\"},\"claim\":{\"kind\":\"client\",\"clientId\":\"client-1\"}},\"fromSeq\":5}}}" + }, + { + "op": "next", + "channel": "ahp-terminal:/t1", + "expect": { + "method": "channel/snapshot", + "params": { + "channel": "ahp-terminal:/t1", + "snapshot": { + "resource": "ahp-terminal:/t1", + "state": { + "title": "Shell", + "content": [], + "lifecycle": { + "status": "running" + }, + "claim": { + "kind": "client", + "clientId": "client-1" + } + }, + "fromSeq": 5 + } + } + }, + "expectBytes": 520 + }, + { + "op": "advance", + "channel": "ahp-terminal:/t1" + }, + { + "op": "expectSent", + "match": { + "method": "channel/credit", + "params": { + "channel": "ahp-terminal:/t1", + "consumedBytes": 520 + } + } + }, + { + "op": "subscribe", + "channel": "ahp-terminal:/t1", + "window": { + "windowBytes": 2048 + } + }, + { + "op": "expectSent", + "match": { + "method": "subscribe", + "params": { + "channel": "ahp-terminal:/t1", + "flowControl": { + "receive": { + "windowBytes": 2048, + "maximumFrameBytes": 16384 + } + } + } + } + }, + { + "op": "respond", + "to": "subscribe", + "result": { + "flowControl": { + "clientReceive": { + "windowBytes": 2048, + "maximumFrameBytes": 1024 + }, + "hostReceive": { + "windowBytes": 512, + "maximumFrameBytes": 1024 + } + } + } + }, + { + "op": "expectStatus", + "channel": "ahp-terminal:/t1", + "match": { + "limits": { + "clientReceive": { + "windowBytes": 2048 + } + } + } + }, + { + "op": "hostSends", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"channel/snapshot\",\"params\":{\"channel\":\"ahp-terminal:/t1\",\"snapshot\":{\"resource\":\"ahp-terminal:/t1\",\"state\":{\"title\":\"Shell\",\"content\":[],\"lifecycle\":{\"status\":\"running\"},\"claim\":{\"kind\":\"client\",\"clientId\":\"client-1\"}},\"fromSeq\":5}}}" + }, + { + "op": "next", + "channel": "ahp-terminal:/t1", + "expect": { + "method": "channel/snapshot", + "params": { + "channel": "ahp-terminal:/t1", + "snapshot": { + "resource": "ahp-terminal:/t1", + "state": { + "title": "Shell", + "content": [], + "lifecycle": { + "status": "running" + }, + "claim": { + "kind": "client", + "clientId": "client-1" + } + }, + "fromSeq": 5 + } + } + }, + "expectBytes": 520 + }, + { + "op": "advance", + "channel": "ahp-terminal:/t1" + }, + { + "op": "expectSent", + "match": { + "method": "channel/credit", + "params": { + "channel": "ahp-terminal:/t1", + "consumedBytes": 1040 + } + } + } + ] +} diff --git a/types/test-cases/windowed-delivery/scenarios/019-violation-oversized-message.json b/types/test-cases/windowed-delivery/scenarios/019-violation-oversized-message.json new file mode 100644 index 000000000..1ee96e52e --- /dev/null +++ b/types/test-cases/windowed-delivery/scenarios/019-violation-oversized-message.json @@ -0,0 +1,57 @@ +{ + "description": "A data message larger than the accepted maximumFrameBytes closes the connection with a protocol error.", + "windowing": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + }, + "steps": [ + { + "op": "subscribe", + "channel": "ahp-terminal:/t1", + "window": {}, + "receive": "auto" + }, + { + "op": "expectSent", + "match": { + "method": "subscribe", + "params": { + "channel": "ahp-terminal:/t1", + "flowControl": { + "receive": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + } + } + } + } + }, + { + "op": "respond", + "to": "subscribe", + "result": { + "flowControl": { + "clientReceive": { + "windowBytes": 4096, + "maximumFrameBytes": 1024 + }, + "hostReceive": { + "windowBytes": 512, + "maximumFrameBytes": 1024 + } + } + } + }, + { + "op": "hostSends", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"channel/snapshot\",\"params\":{\"channel\":\"ahp-terminal:/t1\",\"snapshot\":{\"resource\":\"ahp-terminal:/t1\",\"state\":{\"title\":\"Shell\",\"content\":[],\"lifecycle\":{\"status\":\"running\"},\"claim\":{\"kind\":\"client\",\"clientId\":\"client-1\"}},\"fromSeq\":5}}}" + }, + { + "op": "hostSends", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"action\",\"params\":{\"channel\":\"ahp-terminal:/t1\",\"action\":{\"type\":\"terminal/data\",\"data\":\"zzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzz\"},\"serverSeq\":6}}" + }, + { + "op": "expectClosed" + } + ] +} diff --git a/types/test-cases/windowed-delivery/scenarios/020-violation-credit-exceeds-sent.json b/types/test-cases/windowed-delivery/scenarios/020-violation-credit-exceeds-sent.json new file mode 100644 index 000000000..992881a02 --- /dev/null +++ b/types/test-cases/windowed-delivery/scenarios/020-violation-credit-exceeds-sent.json @@ -0,0 +1,57 @@ +{ + "description": "Credit for more bytes than the client has sent closes the connection with a protocol error.", + "windowing": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + }, + "steps": [ + { + "op": "subscribe", + "channel": "ahp-terminal:/t1", + "window": {}, + "receive": "auto" + }, + { + "op": "expectSent", + "match": { + "method": "subscribe", + "params": { + "channel": "ahp-terminal:/t1", + "flowControl": { + "receive": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + } + } + } + } + }, + { + "op": "respond", + "to": "subscribe", + "result": { + "flowControl": { + "clientReceive": { + "windowBytes": 512, + "maximumFrameBytes": 1024 + }, + "hostReceive": { + "windowBytes": 512, + "maximumFrameBytes": 1024 + } + } + } + }, + { + "op": "hostSends", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"channel/snapshot\",\"params\":{\"channel\":\"ahp-terminal:/t1\",\"snapshot\":{\"resource\":\"ahp-terminal:/t1\",\"state\":{\"title\":\"Shell\",\"content\":[],\"lifecycle\":{\"status\":\"running\"},\"claim\":{\"kind\":\"client\",\"clientId\":\"client-1\"}},\"fromSeq\":5}}}" + }, + { + "op": "hostSends", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"channel/credit\",\"params\":{\"channel\":\"ahp-terminal:/t1\",\"consumedBytes\":10}}" + }, + { + "op": "expectClosed" + } + ] +} diff --git a/types/test-cases/windowed-delivery/scenarios/021-violation-frame-channel-mismatch.json b/types/test-cases/windowed-delivery/scenarios/021-violation-frame-channel-mismatch.json new file mode 100644 index 000000000..b716443ad --- /dev/null +++ b/types/test-cases/windowed-delivery/scenarios/021-violation-frame-channel-mismatch.json @@ -0,0 +1,83 @@ +{ + "description": "A joined fragmented message whose channel differs from the fragments closes the connection with a protocol error.", + "windowing": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + }, + "steps": [ + { + "op": "subscribe", + "channel": "ahp-terminal:/t1", + "window": {}, + "receive": "auto" + }, + { + "op": "expectSent", + "match": { + "method": "subscribe", + "params": { + "channel": "ahp-terminal:/t1", + "flowControl": { + "receive": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + } + } + } + } + }, + { + "op": "respond", + "to": "subscribe", + "result": { + "flowControl": { + "clientReceive": { + "windowBytes": 4096, + "maximumFrameBytes": 1024 + }, + "hostReceive": { + "windowBytes": 512, + "maximumFrameBytes": 1024 + } + } + } + }, + { + "op": "hostSends", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"channel/snapshot\",\"params\":{\"channel\":\"ahp-terminal:/t1\",\"snapshot\":{\"resource\":\"ahp-terminal:/t1\",\"state\":{\"title\":\"Shell\",\"content\":[],\"lifecycle\":{\"status\":\"running\"},\"claim\":{\"kind\":\"client\",\"clientId\":\"client-1\"}},\"fromSeq\":5}}}" + }, + { + "op": "next", + "channel": "ahp-terminal:/t1", + "expect": { + "method": "channel/snapshot", + "params": { + "channel": "ahp-terminal:/t1", + "snapshot": { + "resource": "ahp-terminal:/t1", + "state": { + "title": "Shell", + "content": [], + "lifecycle": { + "status": "running" + }, + "claim": { + "kind": "client", + "clientId": "client-1" + } + }, + "fromSeq": 5 + } + } + }, + "expectBytes": 520 + }, + { + "op": "hostSends", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"channel/frame\",\"params\":{\"channel\":\"ahp-terminal:/t1\",\"data\":\"{\\\"jsonrpc\\\":\\\"2.0\\\",\\\"method\\\":\\\"action\\\",\\\"params\\\":{\\\"channel\\\":\\\"ahp-terminal:/other\\\",\\\"action\\\":{\\\"type\\\":\\\"terminal/data\\\",\\\"data\\\":\\\"x\\\"},\\\"serverSeq\\\":6}}\",\"final\":true}}" + }, + { + "op": "expectClosed" + } + ] +} diff --git a/types/test-cases/windowed-delivery/scenarios/022-violation-raised-limits.json b/types/test-cases/windowed-delivery/scenarios/022-violation-raised-limits.json new file mode 100644 index 000000000..35ce391d7 --- /dev/null +++ b/types/test-cases/windowed-delivery/scenarios/022-violation-raised-limits.json @@ -0,0 +1,48 @@ +{ + "description": "Accepted clientReceive limits higher than the offer close the connection with a protocol error.", + "windowing": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + }, + "steps": [ + { + "op": "subscribe", + "channel": "ahp-terminal:/t1", + "window": {} + }, + { + "op": "expectSent", + "match": { + "method": "subscribe", + "params": { + "channel": "ahp-terminal:/t1", + "flowControl": { + "receive": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + } + } + } + } + }, + { + "op": "respond", + "to": "subscribe", + "result": { + "flowControl": { + "clientReceive": { + "windowBytes": 999999, + "maximumFrameBytes": 16384 + }, + "hostReceive": { + "windowBytes": 512, + "maximumFrameBytes": 1024 + } + } + } + }, + { + "op": "expectClosed" + } + ] +} diff --git a/types/test-cases/windowed-delivery/scenarios/023-unsubscribe-then-resubscribe.json b/types/test-cases/windowed-delivery/scenarios/023-unsubscribe-then-resubscribe.json new file mode 100644 index 000000000..aec2f20bb --- /dev/null +++ b/types/test-cases/windowed-delivery/scenarios/023-unsubscribe-then-resubscribe.json @@ -0,0 +1,186 @@ +{ + "description": "After unsubscribe, releasing a delivery from the old subscription sends no credit. A new subscription on the same URI starts its counters at zero.", + "windowing": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + }, + "steps": [ + { + "op": "subscribe", + "channel": "ahp-terminal:/t1", + "window": {}, + "receive": "manual" + }, + { + "op": "expectSent", + "match": { + "method": "subscribe", + "params": { + "channel": "ahp-terminal:/t1", + "flowControl": { + "receive": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + } + } + } + } + }, + { + "op": "respond", + "to": "subscribe", + "result": { + "flowControl": { + "clientReceive": { + "windowBytes": 4096, + "maximumFrameBytes": 1024 + }, + "hostReceive": { + "windowBytes": 512, + "maximumFrameBytes": 1024 + } + } + } + }, + { + "op": "hostSends", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"channel/snapshot\",\"params\":{\"channel\":\"ahp-terminal:/t1\",\"snapshot\":{\"resource\":\"ahp-terminal:/t1\",\"state\":{\"title\":\"Shell\",\"content\":[],\"lifecycle\":{\"status\":\"running\"},\"claim\":{\"kind\":\"client\",\"clientId\":\"client-1\"}},\"fromSeq\":5}}}" + }, + { + "op": "next", + "channel": "ahp-terminal:/t1", + "expect": { + "method": "channel/snapshot", + "params": { + "channel": "ahp-terminal:/t1", + "snapshot": { + "resource": "ahp-terminal:/t1", + "state": { + "title": "Shell", + "content": [], + "lifecycle": { + "status": "running" + }, + "claim": { + "kind": "client", + "clientId": "client-1" + } + }, + "fromSeq": 5 + } + } + }, + "expectBytes": 520 + }, + { + "op": "unsubscribe", + "channel": "ahp-terminal:/t1" + }, + { + "op": "expectSent", + "match": { + "method": "unsubscribe", + "params": { + "channel": "ahp-terminal:/t1" + } + } + }, + { + "op": "subscribe", + "channel": "ahp-terminal:/t1", + "window": {}, + "receive": "manual" + }, + { + "op": "expectSent", + "match": { + "method": "subscribe", + "params": { + "channel": "ahp-terminal:/t1", + "flowControl": { + "receive": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + } + } + } + } + }, + { + "op": "respond", + "to": "subscribe", + "result": { + "flowControl": { + "clientReceive": { + "windowBytes": 4096, + "maximumFrameBytes": 1024 + }, + "hostReceive": { + "windowBytes": 512, + "maximumFrameBytes": 1024 + } + } + } + }, + { + "op": "hostSends", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"channel/snapshot\",\"params\":{\"channel\":\"ahp-terminal:/t1\",\"snapshot\":{\"resource\":\"ahp-terminal:/t1\",\"state\":{\"title\":\"Shell\",\"content\":[],\"lifecycle\":{\"status\":\"running\"},\"claim\":{\"kind\":\"client\",\"clientId\":\"client-1\"}},\"fromSeq\":5}}}" + }, + { + "op": "next", + "channel": "ahp-terminal:/t1", + "expect": { + "method": "channel/snapshot", + "params": { + "channel": "ahp-terminal:/t1", + "snapshot": { + "resource": "ahp-terminal:/t1", + "state": { + "title": "Shell", + "content": [], + "lifecycle": { + "status": "running" + }, + "claim": { + "kind": "client", + "clientId": "client-1" + } + }, + "fromSeq": 5 + } + } + }, + "expectBytes": 520 + }, + { + "op": "release", + "channel": "ahp-terminal:/t1", + "index": 0 + }, + { + "op": "expectNotSent", + "match": { + "method": "channel/credit", + "params": { + "channel": "ahp-terminal:/t1" + } + }, + "withinMs": 150 + }, + { + "op": "release", + "channel": "ahp-terminal:/t1", + "index": 1 + }, + { + "op": "expectSent", + "match": { + "method": "channel/credit", + "params": { + "channel": "ahp-terminal:/t1", + "consumedBytes": 520 + } + } + } + ] +} diff --git a/types/test-cases/windowed-delivery/scenarios/024-message-larger-than-send-buffer.json b/types/test-cases/windowed-delivery/scenarios/024-message-larger-than-send-buffer.json new file mode 100644 index 000000000..d19d138e0 --- /dev/null +++ b/types/test-cases/windowed-delivery/scenarios/024-message-larger-than-send-buffer.json @@ -0,0 +1,81 @@ +{ + "description": "A single message larger than maxBufferedSendBytes is accepted when nothing else is queued, and is sent in fragments.", + "windowing": { + "windowBytes": 262144, + "maximumFrameBytes": 16384, + "maxBufferedSendBytes": 4096 + }, + "steps": [ + { + "op": "subscribe", + "channel": "ahp-terminal:/t1", + "window": {} + }, + { + "op": "expectSent", + "match": { + "method": "subscribe", + "params": { + "channel": "ahp-terminal:/t1", + "flowControl": { + "receive": { + "windowBytes": 262144, + "maximumFrameBytes": 16384 + } + } + } + } + }, + { + "op": "respond", + "to": "subscribe", + "result": { + "flowControl": { + "clientReceive": { + "windowBytes": 65536, + "maximumFrameBytes": 16384 + }, + "hostReceive": { + "windowBytes": 1048576, + "maximumFrameBytes": 1024 + } + } + } + }, + { + "op": "expectStatus", + "channel": "ahp-terminal:/t1", + "match": { + "writable": true + } + }, + { + "op": "send", + "channel": "ahp-terminal:/t1", + "method": "dispatchAction", + "params": { + "channel": "ahp-terminal:/t1", + "clientSeq": 1, + "action": { + "type": "terminal/input", + "data": "qqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqq" + } + } + }, + { + "op": "expectSent", + "reassemble": true, + "match": { + "method": "dispatchAction", + "params": { + "channel": "ahp-terminal:/t1", + "clientSeq": 1, + "action": { + "type": "terminal/input", + "data": "qqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqq" + } + } + } + } + ] +} diff --git a/types/test-cases/windowed-delivery/sizes.json b/types/test-cases/windowed-delivery/sizes.json new file mode 100644 index 000000000..008423483 --- /dev/null +++ b/types/test-cases/windowed-delivery/sizes.json @@ -0,0 +1,77 @@ +{ + "description": "Raw JSON texts and their size in UTF-16 bytes (2 bytes per UTF-16 code unit of the text exactly as received).", + "cases": [ + { + "name": "ascii", + "text": "{\"a\":\"A\"}", + "utf16Bytes": 18, + "utf8Bytes": 9 + }, + { + "name": "whitespace and member order are measured as received", + "text": "{ \"b\" : 1,\n \"a\" : [ ] }", + "utf16Bytes": 48, + "utf8Bytes": 24 + }, + { + "name": "latin-1 é is one code unit", + "text": "{\"t\":\"é\"}", + "utf16Bytes": 18, + "utf8Bytes": 10 + }, + { + "name": "escaped é counts its escape text", + "text": "{\"t\":\"\\u00e9\"}", + "utf16Bytes": 28, + "utf8Bytes": 14 + }, + { + "name": "CJK is one code unit per character", + "text": "{\"t\":\"漢字\"}", + "utf16Bytes": 20, + "utf8Bytes": 14 + }, + { + "name": "emoji is a surrogate pair", + "text": "{\"t\":\"😀\"}", + "utf16Bytes": 20, + "utf8Bytes": 12 + }, + { + "name": "escaped surrogate pair", + "text": "{\"t\":\"\\ud83d\\ude00\"}", + "utf16Bytes": 40, + "utf8Bytes": 20 + }, + { + "name": "escaped control characters", + "text": "{\"t\":\"a\\r\\n\\t\\\"\\\\b\"}", + "utf16Bytes": 40, + "utf8Bytes": 20 + }, + { + "name": "spec snapshot example", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"channel/snapshot\",\"params\":{\"channel\":\"ahp-terminal:/t1\",\"snapshot\":{\"resource\":\"ahp-terminal:/t1\",\"state\":{\"title\":\"Shell\",\"content\":[],\"lifecycle\":{\"status\":\"running\"},\"claim\":{\"kind\":\"client\",\"clientId\":\"client-1\"}},\"fromSeq\":5}}}", + "utf16Bytes": 520, + "utf8Bytes": 260 + }, + { + "name": "spec action example", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"action\",\"params\":{\"channel\":\"ahp-terminal:/t1\",\"action\":{\"type\":\"terminal/data\",\"data\":\"Hello\\r\\n\"},\"serverSeq\":6}}", + "utf16Bytes": 284, + "utf8Bytes": 142 + }, + { + "name": "spec first fragment", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"channel/frame\",\"params\":{\"channel\":\"ahp-terminal:/t1\",\"data\":\"{\\\"jsonrpc\\\":\\\"2.0\\\",\\\"method\\\":\\\"action\\\",\\\"params\\\":{\\\"channel\\\":\\\"ahp-terminal:/t1\\\",\"}}", + "utf16Bytes": 360, + "utf8Bytes": 180 + }, + { + "name": "spec last fragment", + "text": "{\"jsonrpc\":\"2.0\",\"method\":\"channel/frame\",\"params\":{\"channel\":\"ahp-terminal:/t1\",\"data\":\"\\\"action\\\":{\\\"type\\\":\\\"terminal/data\\\",\\\"data\\\":\\\"Hello\\\\r\\\\n\\\"},\\\"serverSeq\\\":6}}\",\"final\":true}}", + "utf16Bytes": 374, + "utf8Bytes": 187 + } + ] +} diff --git a/types/windowed-delivery-corpus.test.ts b/types/windowed-delivery-corpus.test.ts new file mode 100644 index 000000000..61b48f925 --- /dev/null +++ b/types/windowed-delivery-corpus.test.ts @@ -0,0 +1,69 @@ +import assert from 'node:assert/strict'; +import { readdirSync, readFileSync } from 'node:fs'; +import { describe, it } from 'node:test'; +import { countsAgainstWindow } from './window-exempt.generated.js'; + +const dir = new URL('./test-cases/windowed-delivery/', import.meta.url); +const read = (path: string) => JSON.parse(readFileSync(new URL(path, dir), 'utf8')); + +const OPS = new Set([ + 'subscribe', 'expectSent', 'expectNotSent', 'respond', 'hostSends', 'hostCreditsAll', + 'next', 'advance', 'release', 'send', 'unsubscribe', 'expectStatus', 'expectClosed', +]); + +interface Step { + op: string; + text?: string; + channel?: string; + expectBytes?: number; + match?: { method?: string; params?: { channel?: string; consumedBytes?: number } }; +} + +describe('windowed-delivery corpus', () => { + it('sizes.json counts UTF-16 code units of the raw text', () => { + const { cases } = read('sizes.json') as { cases: { name: string; text: string; utf16Bytes: number; utf8Bytes: number }[] }; + assert.ok(cases.length > 0); + for (const c of cases) { + assert.equal(c.utf16Bytes, c.text.length * 2, c.name); + assert.equal(c.utf8Bytes, Buffer.byteLength(c.text, 'utf8'), c.name); + } + }); + + const names = readdirSync(new URL('scenarios/', dir)).filter(n => n.endsWith('.json')).sort(); + + it('has scenarios', () => assert.ok(names.length >= 20)); + + for (const name of names) { + it(`scenario ${name} is well formed`, () => { + const scenario = read(`scenarios/${name}`) as { description: string; steps: Step[] }; + assert.equal(typeof scenario.description, 'string'); + const frames = new Map(); + const lastCredit = new Map(); + for (const step of scenario.steps) { + assert.ok(OPS.has(step.op), `${name}: unknown op ${step.op}`); + if (step.op === 'hostSends') { + const message = JSON.parse(step.text!); + assert.equal(message.jsonrpc, '2.0'); + assert.equal(typeof message.method, 'string'); + if (message.method === 'channel/frame') { + const channel = message.params.channel as string; + const joined = (frames.get(channel) ?? '') + message.params.data; + if (message.params.final) { + frames.delete(channel); + const inner = JSON.parse(joined); + assert.ok(countsAgainstWindow(inner.method), `${name}: fragmented control message`); + } else { + frames.set(channel, joined); + } + } + } + if (step.op === 'expectSent' && step.match?.method === 'channel/credit') { + const channel = step.match.params!.channel!; + const consumed = step.match.params!.consumedBytes!; + assert.ok(consumed > (lastCredit.get(channel) ?? 0), `${name}: credit must increase`); + lastCredit.set(channel, consumed); + } + } + }); + } +});