From 777679bcb3ad99faeac41ccbc286ef5c22deef23 Mon Sep 17 00:00:00 2001 From: Connor Peet Date: Wed, 14 Oct 2026 20:46:51 -0700 Subject: [PATCH 1/8] fix(generators): stop decoding terminal snapshots as session state Go, Kotlin, and .NET chose SessionState for any snapshot with a lifecycle member, but TerminalState has one too. Require provider as well, and add a terminal snapshot round-trip fixture. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../Generated/State.generated.cs | 7 ++-- clients/go/ahptypes/state.generated.go | 2 +- .../generated/State.generated.kt | 5 ++- .../20261014-terminal-snapshot-decoding.json | 9 +++++ scripts/generate-csharp.ts | 7 ++-- scripts/generate-go.ts | 2 +- scripts/generate-kotlin.ts | 5 ++- .../round-trips/076-terminal-snapshot.json | 38 +++++++++++++++++++ 8 files changed, 61 insertions(+), 14 deletions(-) create mode 100644 docs/.changes/20261014-terminal-snapshot-decoding.json create mode 100644 types/test-cases/round-trips/076-terminal-snapshot.json 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/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/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/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/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 + } + ] +} From 6880545d7628bad21d254ce53c2f8d020299fe11 Mon Sep 17 00:00:00 2001 From: Connor Peet Date: Wed, 14 Oct 2026 20:46:51 -0700 Subject: [PATCH 2/8] test(channels): add a shared windowed-delivery conformance corpus Scripted host conversations and UTF-16 size cases that every client library runs, so buffering, credit, coalescing, and sending behave the same way in each language. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/specification/windowed-delivery.md | 17 +- types/test-cases/windowed-delivery/README.md | 78 ++++++++ .../scenarios/001-accept-snapshot-credit.json | 117 +++++++++++ .../scenarios/002-declined-falls-back.json | 105 ++++++++++ .../scenarios/003-not-opted-in.json | 25 +++ .../004-per-subscription-overrides.json | 31 +++ .../005-result-and-data-back-to-back.json | 109 ++++++++++ .../scenarios/006-manual-release.json | 150 ++++++++++++++ .../scenarios/007-spec-fragments.json | 138 +++++++++++++ .../008-fragmented-larger-than-window.json | 152 ++++++++++++++ .../009-unknown-counts-exempt-does-not.json | 124 ++++++++++++ .../010-empty-and-unknown-snapshot-state.json | 83 ++++++++ .../scenarios/011-coalesce-terminal-data.json | 124 ++++++++++++ .../scenarios/012-coalescing-disabled.json | 123 ++++++++++++ .../013-no-coalesce-across-origin.json | 146 ++++++++++++++ .../scenarios/014-coalesce-chat-delta.json | 110 +++++++++++ .../scenarios/015-send-blocks-on-window.json | 151 ++++++++++++++ .../scenarios/016-outbound-fragmenting.json | 73 +++++++ .../017-tcp-data-send-coalescing.json | 110 +++++++++++ .../018-resubscribe-keeps-counters.json | 183 +++++++++++++++++ .../019-violation-oversized-message.json | 57 ++++++ .../020-violation-credit-exceeds-sent.json | 57 ++++++ .../021-violation-frame-channel-mismatch.json | 83 ++++++++ .../022-violation-raised-limits.json | 48 +++++ .../023-unsubscribe-then-resubscribe.json | 186 ++++++++++++++++++ .../024-message-larger-than-send-buffer.json | 81 ++++++++ types/test-cases/windowed-delivery/sizes.json | 77 ++++++++ types/windowed-delivery-corpus.test.ts | 69 +++++++ 28 files changed, 2802 insertions(+), 5 deletions(-) create mode 100644 types/test-cases/windowed-delivery/README.md create mode 100644 types/test-cases/windowed-delivery/scenarios/001-accept-snapshot-credit.json create mode 100644 types/test-cases/windowed-delivery/scenarios/002-declined-falls-back.json create mode 100644 types/test-cases/windowed-delivery/scenarios/003-not-opted-in.json create mode 100644 types/test-cases/windowed-delivery/scenarios/004-per-subscription-overrides.json create mode 100644 types/test-cases/windowed-delivery/scenarios/005-result-and-data-back-to-back.json create mode 100644 types/test-cases/windowed-delivery/scenarios/006-manual-release.json create mode 100644 types/test-cases/windowed-delivery/scenarios/007-spec-fragments.json create mode 100644 types/test-cases/windowed-delivery/scenarios/008-fragmented-larger-than-window.json create mode 100644 types/test-cases/windowed-delivery/scenarios/009-unknown-counts-exempt-does-not.json create mode 100644 types/test-cases/windowed-delivery/scenarios/010-empty-and-unknown-snapshot-state.json create mode 100644 types/test-cases/windowed-delivery/scenarios/011-coalesce-terminal-data.json create mode 100644 types/test-cases/windowed-delivery/scenarios/012-coalescing-disabled.json create mode 100644 types/test-cases/windowed-delivery/scenarios/013-no-coalesce-across-origin.json create mode 100644 types/test-cases/windowed-delivery/scenarios/014-coalesce-chat-delta.json create mode 100644 types/test-cases/windowed-delivery/scenarios/015-send-blocks-on-window.json create mode 100644 types/test-cases/windowed-delivery/scenarios/016-outbound-fragmenting.json create mode 100644 types/test-cases/windowed-delivery/scenarios/017-tcp-data-send-coalescing.json create mode 100644 types/test-cases/windowed-delivery/scenarios/018-resubscribe-keeps-counters.json create mode 100644 types/test-cases/windowed-delivery/scenarios/019-violation-oversized-message.json create mode 100644 types/test-cases/windowed-delivery/scenarios/020-violation-credit-exceeds-sent.json create mode 100644 types/test-cases/windowed-delivery/scenarios/021-violation-frame-channel-mismatch.json create mode 100644 types/test-cases/windowed-delivery/scenarios/022-violation-raised-limits.json create mode 100644 types/test-cases/windowed-delivery/scenarios/023-unsubscribe-then-resubscribe.json create mode 100644 types/test-cases/windowed-delivery/scenarios/024-message-larger-than-send-buffer.json create mode 100644 types/test-cases/windowed-delivery/sizes.json create mode 100644 types/windowed-delivery-corpus.test.ts 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/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); + } + } + }); + } +}); From d7f01eeced2a51e94cf634ca7ae6d81f292e7796 Mon Sep 17 00:00:00 2001 From: Connor Peet Date: Wed, 14 Oct 2026 20:46:52 -0700 Subject: [PATCH 3/8] feat(typescript): add opt-in windowed delivery to the TypeScript client Subscriptions can opt into windowed delivery. Credit follows the consumer, sends wait for the host window, queued data is coalesced, and managed hosts keep windowed subscriptions across reconnects. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- clients/typescript/README.md | 208 ++++ clients/typescript/src/client/client.ts | 421 ++++++- clients/typescript/src/client/error.ts | 9 + clients/typescript/src/client/events.ts | 43 +- clients/typescript/src/client/flow-control.ts | 1028 +++++++++++++++++ .../src/client/hosts/host-client-handle.ts | 31 +- clients/typescript/src/client/hosts/index.ts | 1 + clients/typescript/src/client/hosts/multi.ts | 45 +- .../typescript/src/client/hosts/runtime.ts | 169 ++- .../src/client/hosts/state-mirror.ts | 11 + clients/typescript/src/client/hosts/types.ts | 5 + clients/typescript/src/client/index.ts | 4 + .../src/client/managed-subscriptions.ts | 77 +- clients/typescript/src/client/state-mirror.ts | 24 + clients/typescript/src/client/transport.ts | 9 +- clients/typescript/src/ws/transport.ts | 44 +- .../test/windowed-delivery-corpus.test.ts | 212 ++++ .../test/windowing-hosts-lifecycle.test.ts | 110 ++ .../typescript/test/windowing-hosts.test.ts | 254 ++++ .../test/windowing-lifecycle.test.ts | 230 ++++ clients/typescript/test/windowing.test.ts | 414 +++++++ clients/typescript/test/ws-transport.test.ts | 34 + ...20261014-typescript-windowed-delivery.json | 7 + 23 files changed, 3317 insertions(+), 73 deletions(-) create mode 100644 clients/typescript/src/client/flow-control.ts create mode 100644 clients/typescript/test/windowed-delivery-corpus.test.ts create mode 100644 clients/typescript/test/windowing-hosts-lifecycle.test.ts create mode 100644 clients/typescript/test/windowing-hosts.test.ts create mode 100644 clients/typescript/test/windowing-lifecycle.test.ts create mode 100644 clients/typescript/test/windowing.test.ts create mode 100644 docs/.changes/20261014-typescript-windowed-delivery.json diff --git a/clients/typescript/README.md b/clients/typescript/README.md index e7ea7371f..f2ae7de9c 100644 --- a/clients/typescript/README.md +++ b/clients/typescript/README.md @@ -118,6 +118,214 @@ 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. + +### 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/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-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" + ] +} From 504e89a897eb2934aee0a1c484d6ac1e92cabf5d Mon Sep 17 00:00:00 2001 From: Connor Peet Date: Wed, 14 Oct 2026 20:46:52 -0700 Subject: [PATCH 4/8] feat(rust): add opt-in windowed delivery to the Rust client Subscriptions can opt into windowed delivery. Credit follows the consumer, sends wait for the host window, queued data is coalesced, and managed hosts keep windowed subscriptions across reconnects. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- clients/rust/Cargo.lock | 7 + clients/rust/README.md | 15 + clients/rust/crates/ahp-ws/README.md | 15 + clients/rust/crates/ahp-ws/src/transport.rs | 22 + clients/rust/crates/ahp-ws/tests/windowing.rs | 84 + clients/rust/crates/ahp/Cargo.toml | 3 +- clients/rust/crates/ahp/README.md | 143 ++ clients/rust/crates/ahp/src/client.rs | 703 +++++- clients/rust/crates/ahp/src/hosts/factory.rs | 20 + clients/rust/crates/ahp/src/hosts/mod.rs | 5 +- clients/rust/crates/ahp/src/hosts/multi.rs | 40 + clients/rust/crates/ahp/src/hosts/runtime.rs | 259 +- clients/rust/crates/ahp/src/hosts/types.rs | 33 + clients/rust/crates/ahp/src/lib.rs | 6 + .../crates/ahp/src/multi_host_state_mirror.rs | 13 + clients/rust/crates/ahp/src/transport.rs | 48 + clients/rust/crates/ahp/src/windowing.rs | 2200 +++++++++++++++++ .../rust/crates/ahp/src/windowing_tests.rs | 434 ++++ .../rust/crates/ahp/tests/windowing_corpus.rs | 707 ++++++ .../crates/ahp/tests/windowing_lifecycle.rs | 369 +++ .../20261014-rust-windowed-delivery.json | 7 + 21 files changed, 5053 insertions(+), 80 deletions(-) create mode 100644 clients/rust/crates/ahp-ws/tests/windowing.rs create mode 100644 clients/rust/crates/ahp/src/windowing.rs create mode 100644 clients/rust/crates/ahp/src/windowing_tests.rs create mode 100644 clients/rust/crates/ahp/tests/windowing_corpus.rs create mode 100644 clients/rust/crates/ahp/tests/windowing_lifecycle.rs create mode 100644 docs/.changes/20261014-rust-windowed-delivery.json 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..77c020b4d 100644 --- a/clients/rust/crates/ahp/README.md +++ b/clients/rust/crates/ahp/README.md @@ -71,3 +71,146 @@ 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`. 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/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" + ] +} From 8d90538dad4c0f4a18d788a862e2c209b570eda5 Mon Sep 17 00:00:00 2001 From: Connor Peet Date: Wed, 14 Oct 2026 20:46:52 -0700 Subject: [PATCH 5/8] feat(go): add opt-in windowed delivery to the Go client Subscriptions can opt into windowed delivery. Credit follows the consumer, sends wait for the host window, queued data is coalesced, and managed hosts keep windowed subscriptions across reconnects. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- clients/go/README.md | 211 ++++ clients/go/ahp/client.go | 239 ++++- clients/go/ahp/coalescing.go | 115 +++ clients/go/ahp/delivery.go | 338 ++++++ clients/go/ahp/hosts/hosts.go | 81 +- .../go/ahp/hosts/managed_lifecycle_test.go | 126 +++ clients/go/ahp/hosts/managed_streams_test.go | 37 + clients/go/ahp/hosts/managed_subscription.go | 348 +++++++ .../go/ahp/hosts/managed_subscription_test.go | 267 +++++ clients/go/ahp/transport.go | 51 + clients/go/ahp/windowing.go | 965 ++++++++++++++++++ clients/go/ahp/windowing_corpus_test.go | 467 +++++++++ clients/go/ahp/windowing_lifecycle.go | 153 +++ clients/go/ahp/windowing_lifecycle_test.go | 360 +++++++ clients/go/ahp/windowing_reconnect.go | 154 +++ clients/go/ahp/windowing_scheduler.go | 548 ++++++++++ .../windowing_subscription_lifecycle_test.go | 129 +++ clients/go/ahp/windowing_test.go | 517 ++++++++++ clients/go/ahpws/transport.go | 14 +- clients/go/ahpws/transport_test.go | 66 ++ .../20261014-go-windowed-delivery.json | 7 + 21 files changed, 5121 insertions(+), 72 deletions(-) create mode 100644 clients/go/ahp/coalescing.go create mode 100644 clients/go/ahp/delivery.go create mode 100644 clients/go/ahp/hosts/managed_lifecycle_test.go create mode 100644 clients/go/ahp/hosts/managed_streams_test.go create mode 100644 clients/go/ahp/hosts/managed_subscription.go create mode 100644 clients/go/ahp/hosts/managed_subscription_test.go create mode 100644 clients/go/ahp/windowing.go create mode 100644 clients/go/ahp/windowing_corpus_test.go create mode 100644 clients/go/ahp/windowing_lifecycle.go create mode 100644 clients/go/ahp/windowing_lifecycle_test.go create mode 100644 clients/go/ahp/windowing_reconnect.go create mode 100644 clients/go/ahp/windowing_scheduler.go create mode 100644 clients/go/ahp/windowing_subscription_lifecycle_test.go create mode 100644 clients/go/ahp/windowing_test.go create mode 100644 clients/go/ahpws/transport_test.go create mode 100644 docs/.changes/20261014-go-windowed-delivery.json diff --git a/clients/go/README.md b/clients/go/README.md index 13562b870..614d8ccad 100644 --- a/clients/go/README.md +++ b/clients/go/README.md @@ -60,6 +60,217 @@ 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. + +### 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/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/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" + ] +} From 5c5e202be22937368bca10a4c20469673ee9ec5e Mon Sep 17 00:00:00 2001 From: Connor Peet Date: Wed, 14 Oct 2026 20:46:53 -0700 Subject: [PATCH 6/8] feat(swift): add opt-in windowed delivery to the Swift client Subscriptions can opt into windowed delivery. Credit follows the consumer, sends wait for the host window, queued data is coalesced, and managed hosts keep windowed subscriptions across reconnects. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- clients/swift/AgentHostProtocol/README.md | 135 +++ .../AgentHostProtocolClient/AHPClient.swift | 460 ++++++++- .../AHPClientConfig.swift | 12 + .../AHPClientEvents.swift | 3 +- .../AHPStateMirror.swift | 8 + .../Hosts/HostClientHandle.swift | 16 + .../Hosts/HostConfig.swift | 6 + .../Hosts/HostRuntime.swift | 188 +++- .../Hosts/HostShared.swift | 3 + .../Hosts/MultiHostClient.swift | 24 +- .../MultiHostStateMirror.swift | 8 + .../NWConnectionWebSocketTransport.swift | 4 + .../URLSessionWebSocketTransport.swift | 19 +- .../AgentHostProtocolClient/Windowing.swift | 519 +++++++++++ .../WindowingEngine.swift | 875 ++++++++++++++++++ .../WindowingTransport.swift | 58 ++ .../NWConnectionWebSocketTransportTests.swift | 24 + .../WindowedDeliveryCorpusTests.swift | 302 ++++++ .../WindowingLifecycleTests.swift | 394 ++++++++ .../WindowingTests.swift | 461 +++++++++ .../20261014-swift-windowed-delivery.json | 7 + 21 files changed, 3462 insertions(+), 64 deletions(-) create mode 100644 clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/Windowing.swift create mode 100644 clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/WindowingEngine.swift create mode 100644 clients/swift/AgentHostProtocol/Sources/AgentHostProtocolClient/WindowingTransport.swift create mode 100644 clients/swift/AgentHostProtocol/Tests/AgentHostProtocolClientTests/WindowedDeliveryCorpusTests.swift create mode 100644 clients/swift/AgentHostProtocol/Tests/AgentHostProtocolClientTests/WindowingLifecycleTests.swift create mode 100644 clients/swift/AgentHostProtocol/Tests/AgentHostProtocolClientTests/WindowingTests.swift create mode 100644 docs/.changes/20261014-swift-windowed-delivery.json diff --git a/clients/swift/AgentHostProtocol/README.md b/clients/swift/AgentHostProtocol/README.md index bab640157..96331578f 100644 --- a/clients/swift/AgentHostProtocol/README.md +++ b/clients/swift/AgentHostProtocol/README.md @@ -78,6 +78,141 @@ 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. + ## 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/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" + ] +} From 75b10020c6a22e0fba8d34bc06edb24a32db29db Mon Sep 17 00:00:00 2001 From: Connor Peet Date: Wed, 14 Oct 2026 20:46:53 -0700 Subject: [PATCH 7/8] feat(dotnet): add opt-in windowed delivery to the .NET client Subscriptions can opt into windowed delivery. Credit follows the consumer, sends wait for the host window, queued data is coalesced, and managed hosts keep windowed subscriptions across reconnects. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- clients/dotnet/README.md | 128 +++ clients/dotnet/TELEMETRY.md | 30 + .../Transport/ITransportMessageLimits.cs | 10 + .../dotnet/src/AgentHostProtocol/AhpClient.cs | 145 ++- .../Hosts/HostClientHandle.cs | 12 + .../Hosts/IMultiHostWindowedClient.cs | 24 + .../Hosts/MultiHostClient.cs | 151 ++- .../Hosts/MultiHostWindowed.cs | 101 ++ .../dotnet/src/AgentHostProtocol/Telemetry.cs | 48 + .../AgentHostProtocol/WebSocketTransport.cs | 4 +- .../src/AgentHostProtocol/WindowOptions.cs | 135 +++ .../WindowedDeliveryEngine.cs | 981 ++++++++++++++++++ .../AgentHostProtocol/WindowedSubscription.cs | 334 ++++++ .../AgentHostProtocol.AotSmoke/Program.cs | 61 ++ .../WindowedClientTests.cs | 341 ++++++ .../WindowedDeliveryCorpusTests.cs | 314 ++++++ .../WindowedLifecycleTests.cs | 554 ++++++++++ .../20261014-dotnet-windowed-delivery.json | 7 + 18 files changed, 3330 insertions(+), 50 deletions(-) create mode 100644 clients/dotnet/src/AgentHostProtocol.Abstractions/Transport/ITransportMessageLimits.cs create mode 100644 clients/dotnet/src/AgentHostProtocol/Hosts/IMultiHostWindowedClient.cs create mode 100644 clients/dotnet/src/AgentHostProtocol/Hosts/MultiHostWindowed.cs create mode 100644 clients/dotnet/src/AgentHostProtocol/WindowOptions.cs create mode 100644 clients/dotnet/src/AgentHostProtocol/WindowedDeliveryEngine.cs create mode 100644 clients/dotnet/src/AgentHostProtocol/WindowedSubscription.cs create mode 100644 clients/dotnet/tests/AgentHostProtocol.Tests/WindowedClientTests.cs create mode 100644 clients/dotnet/tests/AgentHostProtocol.Tests/WindowedDeliveryCorpusTests.cs create mode 100644 clients/dotnet/tests/AgentHostProtocol.Tests/WindowedLifecycleTests.cs create mode 100644 docs/.changes/20261014-dotnet-windowed-delivery.json diff --git a/clients/dotnet/README.md b/clients/dotnet/README.md index b5d3839b8..784c7f9e4 100644 --- a/clients/dotnet/README.md +++ b/clients/dotnet/README.md @@ -87,6 +87,134 @@ 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. + ## 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/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/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" + ] +} From b31a9926acc5de396e09a55a91a8223e61cb00cb Mon Sep 17 00:00:00 2001 From: Connor Peet Date: Thu, 15 Oct 2026 09:59:07 -0700 Subject: [PATCH 8/8] docs(clients): add windowed delivery examples Add a basic and an advanced windowing example to the TypeScript, Rust, Go, Swift, and .NET clients. Each example runs against an in-memory demo host that follows the window rules, so it needs no server. The basic example shows the host pausing at a full window and resuming on credit. The advanced example shows a slow renderer, built-in coalescing, a custom coalescer that keeps the latest terminal resize, and a paste that waits for the subscription to become writable. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- clients/dotnet/AgentHostProtocol.slnx | 2 + clients/dotnet/README.md | 22 +- clients/dotnet/examples/Shared/DemoSupport.cs | 37 ++ .../examples/Shared/InMemoryTransport.cs | 62 ++++ .../examples/Shared/WindowingDemoHost.cs | 324 ++++++++++++++++ .../examples/WindowingBackpressure/Program.cs | 170 +++++++++ .../WindowingBackpressure.csproj | 21 ++ .../dotnet/examples/WindowingBasic/Program.cs | 86 +++++ .../WindowingBasic/WindowingBasic.csproj | 21 ++ clients/go/README.md | 21 ++ .../go/examples/internal/demohost/demohost.go | 348 +++++++++++++++++ clients/go/examples/internal/demohost/pipe.go | 107 ++++++ .../examples/windowing_backpressure/main.go | 230 ++++++++++++ clients/go/examples/windowing_basic/main.go | 110 ++++++ clients/rust/crates/ahp/README.md | 20 + .../examples/support/windowing_demo_host.rs | 351 ++++++++++++++++++ .../ahp/examples/windowing_backpressure.rs | 257 +++++++++++++ .../crates/ahp/examples/windowing_basic.rs | 132 +++++++ clients/swift/AgentHostProtocol/README.md | 20 + clients/swift/Examples/Package.swift | 37 ++ .../Sources/WindowingBackpressure/main.swift | 163 ++++++++ .../Sources/WindowingBasic/main.swift | 79 ++++ .../Sources/WindowingDemoHost/DemoHost.swift | 226 +++++++++++ clients/typescript/README.md | 21 ++ clients/typescript/examples/tsconfig.json | 8 + .../examples/windowing-backpressure.ts | 159 ++++++++ .../typescript/examples/windowing-basic.ts | 74 ++++ .../examples/windowing-demo-host.ts | 187 ++++++++++ clients/typescript/package.json | 1 + 29 files changed, 3295 insertions(+), 1 deletion(-) create mode 100644 clients/dotnet/examples/Shared/DemoSupport.cs create mode 100644 clients/dotnet/examples/Shared/InMemoryTransport.cs create mode 100644 clients/dotnet/examples/Shared/WindowingDemoHost.cs create mode 100644 clients/dotnet/examples/WindowingBackpressure/Program.cs create mode 100644 clients/dotnet/examples/WindowingBackpressure/WindowingBackpressure.csproj create mode 100644 clients/dotnet/examples/WindowingBasic/Program.cs create mode 100644 clients/dotnet/examples/WindowingBasic/WindowingBasic.csproj create mode 100644 clients/go/examples/internal/demohost/demohost.go create mode 100644 clients/go/examples/internal/demohost/pipe.go create mode 100644 clients/go/examples/windowing_backpressure/main.go create mode 100644 clients/go/examples/windowing_basic/main.go create mode 100644 clients/rust/crates/ahp/examples/support/windowing_demo_host.rs create mode 100644 clients/rust/crates/ahp/examples/windowing_backpressure.rs create mode 100644 clients/rust/crates/ahp/examples/windowing_basic.rs create mode 100644 clients/swift/Examples/Package.swift create mode 100644 clients/swift/Examples/Sources/WindowingBackpressure/main.swift create mode 100644 clients/swift/Examples/Sources/WindowingBasic/main.swift create mode 100644 clients/swift/Examples/Sources/WindowingDemoHost/DemoHost.swift create mode 100644 clients/typescript/examples/tsconfig.json create mode 100644 clients/typescript/examples/windowing-backpressure.ts create mode 100644 clients/typescript/examples/windowing-basic.ts create mode 100644 clients/typescript/examples/windowing-demo-host.ts 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 784c7f9e4..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 @@ -215,6 +216,25 @@ 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/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/go/README.md b/clients/go/README.md index 614d8ccad..b1c68ee97 100644 --- a/clients/go/README.md +++ b/clients/go/README.md @@ -140,6 +140,27 @@ 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 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/rust/crates/ahp/README.md b/clients/rust/crates/ahp/README.md index 77c020b4d..71e895193 100644 --- a/clients/rust/crates/ahp/README.md +++ b/clients/rust/crates/ahp/README.md @@ -214,3 +214,23 @@ 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/swift/AgentHostProtocol/README.md b/clients/swift/AgentHostProtocol/README.md index 96331578f..0eb03f095 100644 --- a/clients/swift/AgentHostProtocol/README.md +++ b/clients/swift/AgentHostProtocol/README.md @@ -213,6 +213,26 @@ 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/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 f2ae7de9c..e2aafe3ff 100644 --- a/clients/typescript/README.md +++ b/clients/typescript/README.md @@ -146,6 +146,27 @@ unknown notifications, without extending the legacy `SubscriptionEvent` union. 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 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" },