From 2e1452419f3b4c461aeb6858e78d2c3bf65e11dd Mon Sep 17 00:00:00 2001 From: Joshua Mouch Date: Wed, 26 Aug 2026 21:05:15 -0400 Subject: [PATCH] feat(typescript): manage shared subscriptions --- clients/typescript/README.md | 41 ++- clients/typescript/src/client/index.ts | 9 + .../src/client/managed-subscriptions.ts | 330 ++++++++++++++++++ .../test/managed-subscriptions.test.ts | 220 ++++++++++++ clients/typescript/tsconfig.json | 2 +- ...0826-typescript-managed-subscriptions.json | 5 + 6 files changed, 605 insertions(+), 2 deletions(-) create mode 100644 clients/typescript/src/client/managed-subscriptions.ts create mode 100644 clients/typescript/test/managed-subscriptions.test.ts create mode 100644 docs/.changes/20260826-typescript-managed-subscriptions.json diff --git a/clients/typescript/README.md b/clients/typescript/README.md index d8b85a4c2..e7ea7371f 100644 --- a/clients/typescript/README.md +++ b/clients/typescript/README.md @@ -14,7 +14,7 @@ The package exposes four subpath exports: | Import path | What it gives you | |---|---| | `@microsoft/agent-host-protocol` | Wire types, actions, commands, reducers, version constants. No I/O. | -| `@microsoft/agent-host-protocol/client` | `AhpClient`, `Subscription`, `AhpStateMirror`, the `AhpTransport` interface, `InMemoryTransport`, and the error taxonomy. | +| `@microsoft/agent-host-protocol/client` | `AhpClient`, `Subscription`, `ManagedSubscriptionManager`, `AhpStateMirror`, the `AhpTransport` interface, `InMemoryTransport`, and the error taxonomy. | | `@microsoft/agent-host-protocol/hosts` | `MultiHostClient`, `HostClientHandle`, `ReconnectPolicy`, `ClientIdStore` (with `InMemoryClientIdStore`), `MultiHostStateMirror`, and the `Host*Error` family. Builds on `/client` to manage one or more host connections with reconnect, generation-checked handles, and fan-in events. | | `@microsoft/agent-host-protocol/ws` | `WebSocketTransport` — an `AhpTransport` implementation backed by the global `WebSocket`. | @@ -79,6 +79,45 @@ class MyTransport implements AhpTransport { `InMemoryTransport.pair()` returns two connected halves that exchange text frames — handy for unit tests that don't need a real socket. +## Shared subscription ownership + +`ManagedSubscriptionManager` coalesces concurrent consumers of the same URI +onto one wire-level subscription. Each named lease receives an independent +event iterator. Disposing the last lease sends `unsubscribe`; failed subscriptions are +cleaned up so the next acquire makes a fresh request. + +```ts +import type { SessionState } from '@microsoft/agent-host-protocol'; +import { + ManagedSubscriptionManager, + type SubscriptionEvent, +} from '@microsoft/agent-host-protocol/client'; + +const subscriptions = new ManagedSubscriptionManager(client); +let consume: Promise; +{ + using lease = subscriptions.acquire( + sessionUri, + 'SessionEditor', + ); + + const { snapshot } = await lease.subscription.ready; + if (snapshot) mirror.applySnapshot(snapshot); + consume = (async () => { + for await (const event of lease.events) { + if (event.type === 'action') mirror.apply(event.params); + } + })(); + + // Use the subscription. Leaving this block disposes the lease. +} +await consume; +``` + +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. + ## Reducers and state mirror The reducer functions (`rootReducer`, `sessionReducer`, diff --git a/clients/typescript/src/client/index.ts b/clients/typescript/src/client/index.ts index 3e8dc1a75..e264977bc 100644 --- a/clients/typescript/src/client/index.ts +++ b/clients/typescript/src/client/index.ts @@ -25,3 +25,12 @@ export type { TransportErrorKind } from './error.js'; export { InMemoryTransport } from './transport.js'; export type { AhpTransport, JsonRpcMessage, TransportFrame } from './transport.js'; export { AhpStateMirror } from './state-mirror.js'; +export { ManagedSubscriptionManager } from './managed-subscriptions.js'; +export type { + ManagedSubscription, + ManagedSubscriptionHolder, + ManagedSubscriptionInfo, + ManagedSubscriptionLease, + ManagedSubscriptionStatus, + ManagedSubscribeResult, +} from './managed-subscriptions.js'; diff --git a/clients/typescript/src/client/managed-subscriptions.ts b/clients/typescript/src/client/managed-subscriptions.ts new file mode 100644 index 000000000..437f8d777 --- /dev/null +++ b/clients/typescript/src/client/managed-subscriptions.ts @@ -0,0 +1,330 @@ +/** + * Shared, reference-counted ownership for AHP subscriptions. + * + * @module client/managed-subscriptions + */ + +import type { SubscribeResult } from '../types/common/commands.js'; +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'; + +/** Lifecycle state of a shared managed subscription. */ +export type ManagedSubscriptionStatus = 'pending' | 'active' | 'failed' | 'closed'; + +/** A named owner retaining a shared subscription. */ +export interface ManagedSubscriptionHolder { + readonly owner: string; + readonly count: number; +} + +/** Read-only inspection data for one managed subscription. */ +export interface ManagedSubscriptionInfo { + readonly uri: URI; + readonly status: ManagedSubscriptionStatus; + readonly refCount: number; + readonly holders: readonly ManagedSubscriptionHolder[]; +} + +/** + * A subscribe result whose snapshot state has been narrowed by the consumer. + * + * The wire owner remains the generated {@link SubscribeResult}; this type only + * gives handwritten client code the same caller-selected state typing that + * VS Code's subscription manager exposes for its resource kinds. + */ +export type ManagedSubscribeResult< + TState extends Snapshot['state'] = Snapshot['state'], +> = Omit & { + readonly snapshot?: Omit & { readonly state: TState }; +}; + +/** + * Shared state for one wire-level subscription. + * + * The object remains valid for the lifetime of every lease that acquired it. + * Await {@link ManagedSubscription.ready} for the initial subscribe result; + * events received during that round-trip are retained by each lease's event + * iterator and delivered afterward in wire order. + */ +export interface ManagedSubscription< + TState extends Snapshot['state'] = Snapshot['state'], +> { + readonly uri: URI; + readonly ready: Promise>; + readonly status: ManagedSubscriptionStatus; + /** Initial subscribe result, once {@link ready} has resolved. */ + readonly result: ManagedSubscribeResult | undefined; + /** Subscribe failure, when {@link status} is `failed`. */ + readonly error: Error | undefined; +} + +class ManagedSubscriptionState implements ManagedSubscription { + readonly uri: URI; + readonly ready: Promise; + private statusValue: ManagedSubscriptionStatus = 'pending'; + private resultValue: SubscribeResult | undefined; + private errorValue: Error | undefined; + private readonly resolveReady: (result: SubscribeResult) => void; + private readonly rejectReady: (error: Error) => void; + + constructor(uri: URI) { + this.uri = uri; + let resolveReady!: (result: SubscribeResult) => void; + let rejectReady!: (error: Error) => void; + this.ready = new Promise((resolve, reject) => { + resolveReady = resolve; + rejectReady = reject; + }); + // A consumer may prefer status/error inspection over awaiting `ready`. + // Observe the rejection here so that path never creates an unhandled one. + void this.ready.catch(() => undefined); + this.resolveReady = resolveReady; + this.rejectReady = rejectReady; + } + + get status(): ManagedSubscriptionStatus { + return this.statusValue; + } + + get result(): SubscribeResult | undefined { + return this.resultValue; + } + + get error(): Error | undefined { + return this.errorValue; + } + + activate(result: SubscribeResult): void { + if (this.statusValue !== 'pending') return; + this.resultValue = result; + this.statusValue = 'active'; + this.resolveReady(result); + } + + fail(error: Error): void { + if (this.statusValue !== 'pending') return; + this.errorValue = error; + this.statusValue = 'failed'; + this.rejectReady(error); + } + + close(message = 'managed subscription closed'): void { + if (this.statusValue === 'failed' || this.statusValue === 'closed') return; + if (this.statusValue === 'pending') { + this.rejectReady(new ClientClosedError(message)); + } + this.statusValue = 'closed'; + } +} + +/** + * One holder's lease on a shared managed subscription. + * + * {@link events} is attached before the wire-level `subscribe` request is + * sent. Disposal is idempotent; disposing the final lease sends the protocol + * `unsubscribe` notification. + */ +export interface ManagedSubscriptionLease< + TState extends Snapshot['state'] = Snapshot['state'], + TEvent extends SubscriptionEvent = SubscriptionEvent, +> extends Disposable { + readonly subscription: ManagedSubscription; + readonly events: AsyncIterableIterator; +} + +interface Entry { + readonly uri: URI; + readonly optionsKey: string; + readonly subscription: ManagedSubscriptionState; + readonly events: AsyncBroadcastQueue; + readonly holders: Map; + source?: Subscription; +} + +/** + * Owns one wire-level subscription per URI and shares it across named leases. + * + * Concurrent acquires coalesce onto the same request. A failed request is + * removed from the manager after its local attachment is cleaned up, so the + * next acquire deterministically makes a fresh request. Releasing the last + * lease tears down both local fan-out and the server subscription. + */ +export class ManagedSubscriptionManager { + private readonly client: AhpClient; + private readonly eventBuffer: number; + private readonly entries = new Map(); + private nextHolderId = 1; + private closed = false; + + constructor(client: AhpClient, options: { eventBuffer?: number } = {}) { + this.client = client; + const buffer = options.eventBuffer ?? 4096; + this.eventBuffer = buffer >= 1 ? Math.floor(buffer) : 1; + } + + /** + * Acquire a named lease for `uri`. + * + * The first acquire starts the wire request. Later acquires share its result + * and event fan-out. All holders for a URI must use the same subscribe + * options; conflicting options throw rather than silently changing the + * already-active server subscription. + */ + acquire< + TState extends Snapshot['state'] = Snapshot['state'], + TEvent extends SubscriptionEvent = SubscriptionEvent, + >( + uri: URI, + owner: string, + options: SubscribeOptions = {}, + ): ManagedSubscriptionLease { + if (this.closed) { + throw new ClientClosedError('managed subscription manager closed'); + } + + const optionsKey = subscriptionOptionsKey(options); + let entry = this.entries.get(uri); + if (entry && entry.optionsKey !== optionsKey) { + throw new TypeError(`subscription options for "${uri}" differ from the active subscription`); + } + + if (!entry) { + entry = { + uri, + optionsKey, + subscription: new ManagedSubscriptionState(uri), + events: new AsyncBroadcastQueue(this.eventBuffer), + holders: new Map(), + }; + this.entries.set(uri, entry); + } + + const lease = this.createLease(entry, owner); + if (entry.holders.size === 1) { + void this.start(entry, options); + } + return lease; + } + + /** Current managed subscription without acquiring another lease. */ + get( + uri: URI, + ): ManagedSubscription | undefined { + return this.entries.get(uri)?.subscription as ManagedSubscription | undefined; + } + + /** Active subscription URIs, in deterministic lexical order. */ + currentSubscriptionUris(): URI[] { + return [...this.entries.keys()].sort(); + } + + /** Read-only lifecycle and ownership snapshot for diagnostics. */ + activeSubscriptions(): ManagedSubscriptionInfo[] { + return [...this.entries.values()] + .sort((a, b) => a.uri.localeCompare(b.uri)) + .map(entry => ({ + uri: entry.uri, + status: entry.subscription.status, + refCount: entry.holders.size, + holders: summarizeHolders(entry.holders), + })); + } + + /** Release every managed subscription and reject future acquires. */ + async close(): Promise { + if (this.closed) return; + this.closed = true; + const entries = [...this.entries.values()]; + this.entries.clear(); + await Promise.all(entries.map(entry => this.disposeEntry( + entry, + 'managed subscription manager closed', + ))); + } + + private createLease< + TState extends Snapshot['state'], + TEvent extends SubscriptionEvent, + >(entry: Entry, owner: string): ManagedSubscriptionLease { + const holderId = this.nextHolderId++; + entry.holders.set(holderId, owner); + const events = entry.events.reader(); + let released = false; + + return { + subscription: entry.subscription as ManagedSubscription, + events: events as AsyncIterableIterator, + [Symbol.dispose]: () => { + if (released) return; + released = true; + void events.return?.(); + entry.holders.delete(holderId); + if (entry.holders.size === 0 && this.entries.get(entry.uri) === entry) { + this.entries.delete(entry.uri); + void this.disposeEntry(entry, 'managed subscription released before ready'); + } + }, + }; + } + + private async start(entry: Entry, options: SubscribeOptions): Promise { + try { + 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); + } catch (cause) { + if (this.entries.get(entry.uri) !== entry) return; + this.entries.delete(entry.uri); + const error = cause instanceof Error ? cause : new Error(String(cause)); + entry.subscription.fail(error); + entry.events.close(); + await this.client.unsubscribe(entry.uri); + } + } + + private async pump(entry: Entry, source: Subscription): Promise { + try { + for await (const event of source) { + entry.events.publish(event); + } + } finally { + if (this.entries.get(entry.uri) === entry) { + this.entries.delete(entry.uri); + entry.subscription.close('managed subscription event stream closed'); + entry.events.close(); + } + } + } + + private async disposeEntry(entry: Entry, message: string): Promise { + entry.subscription.close(message); + entry.events.close(); + await this.client.unsubscribe(entry.uri); + await entry.source?.close(); + } +} + +function subscriptionOptionsKey(options: SubscribeOptions): string { + return JSON.stringify({ + maxLatencyMs: options.delivery?.maxLatencyMs ?? null, + turns: options.view?.turns ?? null, + }); +} + +function summarizeHolders(holders: ReadonlyMap): ManagedSubscriptionHolder[] { + const counts = new Map(); + for (const owner of holders.values()) { + counts.set(owner, (counts.get(owner) ?? 0) + 1); + } + return [...counts.entries()] + .map(([owner, count]) => ({ owner, count })) + .sort((a, b) => b.count - a.count || a.owner.localeCompare(b.owner)); +} diff --git a/clients/typescript/test/managed-subscriptions.test.ts b/clients/typescript/test/managed-subscriptions.test.ts new file mode 100644 index 000000000..a318e03f6 --- /dev/null +++ b/clients/typescript/test/managed-subscriptions.test.ts @@ -0,0 +1,220 @@ +import test from 'node:test'; +import assert from 'node:assert/strict'; + +import { + AhpClient, + ClientClosedError, + InMemoryTransport, + ManagedSubscriptionManager, + RpcError, + type AhpTransport, + type SubscriptionEvent, +} from '../src/client/index.js'; +import type { ActionEnvelope, StateAction } from '../src/types/common/actions.js'; +import { ActionType } from '../src/types/common/actions.js'; +import type { JsonRpcNotification, JsonRpcRequest } from '../src/types/common/messages.js'; +import type { SubscribeResult } from '../src/types/common/commands.js'; +import type { SessionState } from '../src/types/channels-session/state.js'; + +const URI = 'ahp-session:/managed'; + +async function readRequest(server: AhpTransport): Promise { + const frame = await server.recv(); + assert.ok(frame && frame.kind === 'text', 'expected a text request'); + return JSON.parse(frame.text) as JsonRpcRequest; +} + +async function readNotification(server: AhpTransport): Promise { + const frame = await server.recv(); + assert.ok(frame && frame.kind === 'text', 'expected a text notification'); + return JSON.parse(frame.text) as JsonRpcNotification; +} + +function result(uri = URI, fromSeq = 0): SubscribeResult { + return { + snapshot: { + resource: uri, + state: {} as NonNullable['state'], + fromSeq, + }, + }; +} + +function reply(server: AhpTransport, id: number, value: unknown): void { + server.send(JSON.stringify({ jsonrpc: '2.0', id, result: value })); +} + +function replyError(server: AhpTransport, id: number, code: number, message: string): void { + server.send(JSON.stringify({ jsonrpc: '2.0', id, error: { code, message } })); +} + +function pushAction(server: AhpTransport, serverSeq: number): void { + const envelope: ActionEnvelope = { + channel: URI, + serverSeq, + action: { + type: ActionType.SessionTitleChanged, + title: `title-${serverSeq}`, + } as unknown as StateAction, + origin: null, + }; + server.send(JSON.stringify({ jsonrpc: '2.0', method: 'action', params: envelope })); +} + +test('coalesces holders and unsubscribes only after the final disposal', async () => { + const [clientTransport, server] = InMemoryTransport.pair(); + const client = new AhpClient(clientTransport); + client.connect(); + const manager = new ManagedSubscriptionManager(client); + + const first = manager.acquire(URI, 'SessionView'); + const second = manager.acquire(URI, 'SessionView'); + assert.equal(first.subscription, second.subscription); + assert.deepEqual(manager.activeSubscriptions(), [{ + uri: URI, + status: 'pending', + refCount: 2, + holders: [{ owner: 'SessionView', count: 2 }], + }]); + + const request = await readRequest(server); + assert.equal(request.method, 'subscribe'); + reply(server, request.id, result()); + await first.subscription.ready; + + pushAction(server, 1); + const [firstEvent, secondEvent] = await Promise.all([ + first.events.next(), + second.events.next(), + ]); + assert.equal(firstEvent.value?.type, 'action'); + assert.equal(secondEvent.value?.type, 'action'); + + first[Symbol.dispose](); + assert.equal(manager.activeSubscriptions()[0]?.refCount, 1); + second[Symbol.dispose](); + + const notification = await readNotification(server); + assert.equal(notification.method, 'unsubscribe'); + assert.deepEqual(manager.activeSubscriptions(), []); + await client.shutdown(); +}); + +test('retains actions received during the initial snapshot round-trip', async () => { + const [clientTransport, server] = InMemoryTransport.pair(); + const client = new AhpClient(clientTransport); + client.connect(); + const manager = new ManagedSubscriptionManager(client); + const lease = manager.acquire(URI, 'StateMirror'); + + const request = await readRequest(server); + pushAction(server, 7); + reply(server, request.id, result(URI, 6)); + + const initial = await lease.subscription.ready; + assert.equal(initial.snapshot?.fromSeq, 6); + const next = await lease.events.next(); + assert.equal(next.done, false); + assert.equal(next.value?.type, 'action'); + if (next.value?.type === 'action') { + assert.equal(next.value.params.serverSeq, 7); + } + + lease[Symbol.dispose](); + await readNotification(server); + await client.shutdown(); +}); + +test('a released pending acquire cannot overwrite a replacement', async () => { + const [clientTransport, server] = InMemoryTransport.pair(); + const client = new AhpClient(clientTransport); + client.connect(); + const manager = new ManagedSubscriptionManager(client); + + const abandoned = manager.acquire(URI, 'Preview'); + const firstRequest = await readRequest(server); + abandoned[Symbol.dispose](); + await assert.rejects(abandoned.subscription.ready, ClientClosedError); + assert.equal((await readNotification(server)).method, 'unsubscribe'); + + const replacement = manager.acquire(URI, 'Editor'); + const secondRequest = await readRequest(server); + reply(server, firstRequest.id, result(URI, 1)); + reply(server, secondRequest.id, result(URI, 2)); + + const replacementResult = await replacement.subscription.ready; + assert.equal(replacementResult.snapshot?.fromSeq, 2); + assert.equal(manager.get(URI), replacement.subscription); + + replacement[Symbol.dispose](); + await readNotification(server); + await client.shutdown(); +}); + +test('cleans up a failed request so the next acquire retries', async () => { + const [clientTransport, server] = InMemoryTransport.pair(); + const client = new AhpClient(clientTransport); + client.connect(); + const manager = new ManagedSubscriptionManager(client); + + const failed = manager.acquire(URI, 'SessionView'); + const firstRequest = await readRequest(server); + replyError(server, firstRequest.id, -32_001, 'not ready'); + await assert.rejects(failed.subscription.ready, RpcError); + assert.equal(failed.subscription.status, 'failed'); + assert.equal((await readNotification(server)).method, 'unsubscribe'); + assert.equal(manager.get(URI), undefined); + + const retry = manager.acquire(URI, 'SessionView'); + const secondRequest = await readRequest(server); + reply(server, secondRequest.id, result(URI, 3)); + assert.equal((await retry.subscription.ready).snapshot?.fromSeq, 3); + assert.equal(retry.subscription.status, 'active'); + + failed[Symbol.dispose](); + assert.equal(manager.get(URI), retry.subscription); + retry[Symbol.dispose](); + await readNotification(server); + await client.shutdown(); +}); + +test('rejects conflicting options for an already managed URI', async () => { + const [clientTransport, server] = InMemoryTransport.pair(); + const client = new AhpClient(clientTransport); + client.connect(); + const manager = new ManagedSubscriptionManager(client); + const lease = manager.acquire(URI, 'Immediate', { delivery: { maxLatencyMs: 0 } }); + + assert.throws( + () => manager.acquire(URI, 'Buffered', { delivery: { maxLatencyMs: 100 } }), + (error: unknown) => error instanceof TypeError + && /differ from the active subscription/.test(error.message), + ); + + const request = await readRequest(server); + reply(server, request.id, result()); + await lease.subscription.ready; + lease[Symbol.dispose](); + await readNotification(server); + await client.shutdown(); +}); + +test('close releases all channels and rejects future acquires', async () => { + const [clientTransport, server] = InMemoryTransport.pair(); + const client = new AhpClient(clientTransport); + client.connect(); + const manager = new ManagedSubscriptionManager(client); + const lease = manager.acquire(URI, 'Window'); + + const request = await readRequest(server); + reply(server, request.id, result()); + await lease.subscription.ready; + + const closing = manager.close(); + assert.equal((await readNotification(server)).method, 'unsubscribe'); + await closing; + assert.equal(lease.subscription.status, 'closed'); + assert.throws(() => manager.acquire(URI, 'Window'), ClientClosedError); + lease[Symbol.dispose](); + await client.shutdown(); +}); diff --git a/clients/typescript/tsconfig.json b/clients/typescript/tsconfig.json index cd44d97f7..e3445a763 100644 --- a/clients/typescript/tsconfig.json +++ b/clients/typescript/tsconfig.json @@ -3,7 +3,7 @@ "target": "ES2022", "module": "NodeNext", "moduleResolution": "NodeNext", - "lib": ["ES2022", "DOM"], + "lib": ["ES2022", "ESNext.Disposable", "DOM"], "strict": true, "declaration": true, "declarationMap": true, diff --git a/docs/.changes/20260826-typescript-managed-subscriptions.json b/docs/.changes/20260826-typescript-managed-subscriptions.json new file mode 100644 index 000000000..0dc0b9211 --- /dev/null +++ b/docs/.changes/20260826-typescript-managed-subscriptions.json @@ -0,0 +1,5 @@ +{ + "type": "added", + "message": "`ManagedSubscriptionManager` provides typed shared leases, initial event buffering, disposable last-holder unsubscribe, and retry-safe failure cleanup.", + "targets": ["typescript"] +}