@ably/ai-transport 0.0.1 → 0.1.0

package/src/core/transport/server-transport.ts

@@ -18,7 +18,6 @@ import {
   HEADER_CANCEL_CLIENT_ID,
   HEADER_CANCEL_OWN,
   HEADER_CANCEL_TURN_ID,
-  HEADER_MSG_ID,
 } from '../../constants.js';
 import { ErrorCode } from '../../errors.js';
 import type { Logger } from '../../logger.js';
@@ -32,7 +31,8 @@ import type {
   AddMessagesResult,
   CancelFilter,
   CancelRequest,
-  MessageWithHeaders,
+  EventsNode,
+  MessageNode,
   NewTurnOptions,
   ServerTransport,
   ServerTransportOptions,
@@ -50,10 +50,27 @@ interface RegisteredTurn {
   turnId: string;
   clientId: string;
   controller: AbortController;
+  /** Composite signal that fires when either the internal controller or the external signal aborts. */
+  signal: AbortSignal;
   onCancel?: (request: CancelRequest) => Promise<boolean>;
   onError?: (error: Ably.ErrorInfo) => void;
 }
 
+// ---------------------------------------------------------------------------
+// Internal state machines
+// ---------------------------------------------------------------------------
+
+enum ServerTransportState {
+  READY = 'ready',
+  CLOSED = 'closed',
+}
+
+enum TurnState {
+  INITIALIZED = 'initialized',
+  STARTED = 'started',
+  ENDED = 'ended',
+}
+
 // ---------------------------------------------------------------------------
 // Implementation
 // ---------------------------------------------------------------------------
@@ -69,6 +86,10 @@ class DefaultServerTransport<TEvent, TMessage> implements ServerTransport<TEvent
   private readonly _channelListener: (msg: Ably.InboundMessage) => void;
   private readonly _attachPromise: Promise<void>;
 
+  private _state = ServerTransportState.READY;
+  private _hasAttachedOnce: boolean;
+  private readonly _onChannelStateChange: Ably.channelEventCallback;
+
   constructor(options: ServerTransportOptions<TEvent, TMessage>) {
     this._channel = options.channel;
     this._codec = options.codec;
@@ -98,6 +119,19 @@ class DefaultServerTransport<TEvent, TMessage> implements ServerTransport<TEvent
       },
     );
 
+    // Spec: AIT-ST12, AIT-ST12a
+    // Listen for channel state changes that break message continuity. The
+    // server only consumes cancel messages from the channel, so losing one
+    // is survivable — but the developer needs to know so they can decide
+    // whether to abort in-flight work. _hasAttachedOnce is seeded from the
+    // channel's current state so pre-attached channels are handled correctly;
+    // it distinguishes the initial attach from a genuine discontinuity.
+    this._hasAttachedOnce = this._channel.state === 'attached';
+    this._onChannelStateChange = (stateChange: Ably.ChannelStateChange) => {
+      this._handleChannelStateChange(stateChange);
+    };
+    this._channel.on(this._onChannelStateChange);
+
     this._logger?.debug('DefaultServerTransport(); transport created');
   }
 
@@ -113,8 +147,11 @@ class DefaultServerTransport<TEvent, TMessage> implements ServerTransport<TEvent
 
   // Spec: AIT-ST11
   close(): void {
+    if (this._state === ServerTransportState.CLOSED) return;
+    this._state = ServerTransportState.CLOSED;
     this._logger?.trace('DefaultServerTransport.close();');
     this._channel.unsubscribe(EVENT_CANCEL, this._channelListener);
+    this._channel.off(this._onChannelStateChange);
     for (const reg of this._registeredTurns.values()) {
       reg.controller.abort();
     }
@@ -143,10 +180,11 @@ class DefaultServerTransport<TEvent, TMessage> implements ServerTransport<TEvent
     return [];
   }
 
-  // Spec: AIT-ST8, AIT-ST9
+  // Spec: AIT-ST8, AIT-ST8a, AIT-ST8b, AIT-ST8c, AIT-ST8d, AIT-ST9, AIT-ST9a
   private async _handleCancelMessage(msg: Ably.InboundMessage): Promise<void> {
     const headers = getHeaders(msg);
 
+    // Spec: AIT-ST8a, AIT-ST8b, AIT-ST8c, AIT-ST8d
     const filter: CancelFilter = {};
     if (headers[HEADER_CANCEL_TURN_ID]) {
       filter.turnId = headers[HEADER_CANCEL_TURN_ID];
@@ -203,6 +241,50 @@ class DefaultServerTransport<TEvent, TMessage> implements ServerTransport<TEvent
     }
   }
 
+  // -------------------------------------------------------------------------
+  // Channel state change handler
+  // -------------------------------------------------------------------------
+
+  // Spec: AIT-ST12, AIT-ST12a
+  private _handleChannelStateChange(stateChange: Ably.ChannelStateChange): void {
+    if (this._state === ServerTransportState.CLOSED) return;
+
+    const { current, resumed } = stateChange;
+
+    // Track the initial attach so we don't treat it as a discontinuity
+    if (current === 'attached' && !this._hasAttachedOnce) {
+      this._hasAttachedOnce = true;
+      return;
+    }
+
+    // Continuity-breaking states:
+    // - FAILED, SUSPENDED, DETACHED: no more messages expected (or gap)
+    // - ATTACHED with resumed: false (UPDATE): messages were lost
+    const continuityLost =
+      current === 'failed' || current === 'suspended' || current === 'detached' || (current === 'attached' && !resumed);
+
+    if (!continuityLost) return;
+
+    this._logger?.error('DefaultServerTransport._handleChannelStateChange(); channel continuity lost', {
+      current,
+      resumed,
+      previous: stateChange.previous,
+    });
+
+    const err = new Ably.ErrorInfo(
+      `unable to deliver cancel messages; channel continuity lost (${current}${current === 'attached' ? ', resumed: false' : ''})`,
+      ErrorCode.ChannelContinuityLost,
+      500,
+      stateChange.reason,
+    );
+
+    // Transport-level notification only: continuity loss is not scoped to any
+    // turn. Per-turn onError handlers are reserved for errors from that turn's
+    // own operations (publish failures, encoder errors). Developers that need
+    // per-turn reaction can iterate active turns from the transport handler.
+    this._onError?.(err);
+  }
+
   // -------------------------------------------------------------------------
   // Channel subscription handler
   // -------------------------------------------------------------------------
@@ -248,17 +330,23 @@ class DefaultServerTransport<TEvent, TMessage> implements ServerTransport<TEvent
       onError: turnOnError,
       parent: turnParent,
       forkOf: turnForkOf,
+      signal: externalSignal,
     } = turnOpts;
 
     const controller = new AbortController();
-    let started = false;
-    let ended = false;
+    let state = TurnState.INITIALIZED;
+
+    // Compose the internal controller signal with the external signal (e.g.
+    // req.signal) so platform-level cancellation (request cancellation, function
+    // timeout) aborts the turn through the same path as Ably cancel messages.
+    const signal = externalSignal ? AbortSignal.any([controller.signal, externalSignal]) : controller.signal;
 
-    // Register immediately so early cancels can fire the abort signal.
+    // Spec: AIT-ST3a — register immediately so early cancels can fire the abort signal.
     const registration: RegisteredTurn = {
       turnId,
       clientId: turnClientId ?? '',
       controller,
+      signal,
       onCancel,
       onError: turnOnError,
     };
@@ -278,25 +366,29 @@ class DefaultServerTransport<TEvent, TMessage> implements ServerTransport<TEvent
         return turnId;
       },
       get abortSignal() {
-        return controller.signal;
+        return signal;
       },
 
-      // Spec: AIT-ST4
+      // Spec: AIT-ST4, AIT-ST4a, AIT-ST4b
       start: async (): Promise<void> => {
         logger?.trace('Turn.start();', { turnId });
 
-        if (controller.signal.aborted) {
+        // Spec: AIT-ST4a
+        if (signal.aborted) {
           throw new Ably.ErrorInfo(
             `unable to start turn; turn ${turnId} was cancelled before start()`,
             ErrorCode.InvalidArgument,
             400,
           );
         }
-        if (started) return;
-        started = true;
+        if (state !== TurnState.INITIALIZED) return;
+        state = TurnState.STARTED;
 
         try {
-          await turnManager.startTurn(turnId, turnClientId, controller);
+          await turnManager.startTurn(turnId, turnClientId, controller, {
+            parent: turnParent,
+            forkOf: turnForkOf,
+          });
         } catch (error) {
           const errInfo = new Ably.ErrorInfo(
             `unable to publish turn-start for turn ${turnId}; ${error instanceof Error ? error.message : String(error)}`,
@@ -305,21 +397,17 @@ class DefaultServerTransport<TEvent, TMessage> implements ServerTransport<TEvent
             error instanceof Ably.ErrorInfo ? error : undefined,
           );
           logger?.error('Turn.start(); failed to publish turn-start', { turnId });
-          turnOnError?.(errInfo);
           throw errInfo;
         }
 
         logger?.debug('Turn.start(); turn started', { turnId });
       },
 
-      // Spec: AIT-ST5
-      addMessages: async (
-        inputs: MessageWithHeaders<TMessage>[],
-        opts?: AddMessageOptions,
-      ): Promise<AddMessagesResult> => {
-        logger?.trace('Turn.addMessages();', { turnId, count: inputs.length });
+      // Spec: AIT-ST5, AIT-ST5a, AIT-ST5b, AIT-ST5c
+      addMessages: async (nodes: MessageNode<TMessage>[], opts?: AddMessageOptions): Promise<AddMessagesResult> => {
+        logger?.trace('Turn.addMessages();', { turnId, count: nodes.length });
 
-        if (!started) {
+        if (state === TurnState.INITIALIZED) {
           throw new Ably.ErrorInfo(
             `unable to add messages; start() must be called before addMessages() (turn ${turnId})`,
             ErrorCode.InvalidArgument,
@@ -330,48 +418,104 @@ class DefaultServerTransport<TEvent, TMessage> implements ServerTransport<TEvent
 
         const msgIds: string[] = [];
 
-        for (const input of inputs) {
-          const msgId = crypto.randomUUID();
+        try {
+          for (const node of nodes) {
+            // Build transport headers from the node's typed fields, then merge
+            // any extra headers from the node (e.g. domain-specific headers).
+            const headers = mergeHeaders(
+              buildTransportHeaders({
+                role: 'user',
+                turnId,
+                msgId: node.msgId,
+                turnClientId: opts?.clientId,
+                parent: node.parentId ?? turnParent,
+                forkOf: node.forkOf ?? turnForkOf,
+              }),
+              node.headers,
+            );
+
+            const encoder = codec.createEncoder(channel, {
+              extras: { headers },
+              onMessage,
+            });
 
-          // Transport headers are the defaults; per-message headers from the
-          // client override them. This lets the client's x-ably-msg-id pass
-          // through for optimistic reconciliation with client inserts.
-          const headers = mergeHeaders(
-            buildTransportHeaders({
-              role: 'user',
-              turnId,
-              msgId,
-              turnClientId: opts?.clientId,
-              // Per-operation options override turn-level defaults
-              parent: opts?.parent === undefined ? (turnParent ?? undefined) : (opts.parent ?? undefined),
-              forkOf: opts?.forkOf ?? turnForkOf,
-            }),
-            input.headers,
+            await encoder.writeMessages([node.message], opts?.clientId ? { clientId: opts.clientId } : undefined);
+
+            msgIds.push(node.msgId);
+          }
+        } catch (error) {
+          const errInfo = new Ably.ErrorInfo(
+            `unable to publish messages for turn ${turnId}; ${error instanceof Error ? error.message : String(error)}`,
+            ErrorCode.TurnLifecycleError,
+            500,
+            error instanceof Ably.ErrorInfo ? error : undefined,
           );
+          logger?.error('Turn.addMessages(); publish failed', { turnId });
+          throw errInfo;
+        }
 
-          const encoder = codec.createEncoder(channel, {
-            extras: { headers },
-            onMessage,
-          });
+        logger?.debug('Turn.addMessages(); messages published', { turnId, count: nodes.length });
+        return { msgIds };
+      },
 
-          await encoder.writeMessages([input.message], opts?.clientId ? { clientId: opts.clientId } : undefined);
+      // Spec: AIT-ST5c
+      addEvents: async (nodes: EventsNode<TEvent>[]): Promise<void> => {
+        logger?.trace('Turn.addEvents();', { turnId, count: nodes.length });
 
-          // Capture the effective msg-id after input.headers may have overridden it.
-          msgIds.push(headers[HEADER_MSG_ID] ?? msgId);
+        if (state === TurnState.INITIALIZED) {
+          throw new Ably.ErrorInfo(
+            `unable to add events; start() must be called before addEvents() (turn ${turnId})`,
+            ErrorCode.InvalidArgument,
+            400,
+          );
         }
+        await attachPromise;
 
-        logger?.debug('Turn.addMessages(); messages published', { turnId, count: inputs.length });
-        return { msgIds };
+        const turnOwnerClientId = turnManager.getClientId(turnId);
+
+        try {
+          for (const node of nodes) {
+            const headers = buildTransportHeaders({
+              role: 'assistant',
+              turnId,
+              msgId: node.msgId,
+              turnClientId: turnOwnerClientId,
+              amend: node.msgId,
+            });
+
+            const encoder = codec.createEncoder(channel, {
+              extras: { headers },
+              onMessage,
+            });
+
+            for (const event of node.events) {
+              await encoder.writeEvent(event);
+            }
+
+            await encoder.close();
+          }
+        } catch (error) {
+          const errInfo = new Ably.ErrorInfo(
+            `unable to publish events for turn ${turnId}; ${error instanceof Error ? error.message : String(error)}`,
+            ErrorCode.TurnLifecycleError,
+            500,
+            error instanceof Ably.ErrorInfo ? error : undefined,
+          );
+          logger?.error('Turn.addEvents(); publish failed', { turnId });
+          throw errInfo;
+        }
+
+        logger?.debug('Turn.addEvents(); events published', { turnId, count: nodes.length });
       },
 
-      // Spec: AIT-ST6
+      // Spec: AIT-ST6, AIT-ST6a, AIT-ST6b, AIT-ST6b1, AIT-ST6b2, AIT-ST6b3, AIT-ST6b4, AIT-ST6c
       streamResponse: async (
         stream: ReadableStream<TEvent>,
-        streamOpts?: StreamResponseOptions,
+        streamOpts?: StreamResponseOptions<TEvent>,
       ): Promise<StreamResult> => {
         logger?.trace('Turn.streamResponse();', { turnId });
 
-        if (!started) {
+        if (state === TurnState.INITIALIZED) {
           throw new Ably.ErrorInfo(
             `unable to stream response; start() must be called before streamResponse() (turn ${turnId})`,
             ErrorCode.InvalidArgument,
@@ -380,17 +524,16 @@ class DefaultServerTransport<TEvent, TMessage> implements ServerTransport<TEvent
         }
         await attachPromise;
 
-        const signal = turnManager.getSignal(turnId);
         const turnOwnerClientId = turnManager.getClientId(turnId);
 
         // Per-operation parent overrides the turn-level default.
-        const assistantParent =
-          streamOpts?.parent === undefined ? (turnParent ?? undefined) : (streamOpts.parent ?? undefined);
+        const assistantParent = streamOpts?.parent === undefined ? turnParent : streamOpts.parent;
 
+        const msgId = crypto.randomUUID();
         const defaultHeaders = buildTransportHeaders({
           role: 'assistant',
           turnId,
-          msgId: crypto.randomUUID(),
+          msgId,
           turnClientId: turnOwnerClientId,
           parent: assistantParent,
           forkOf: streamOpts?.forkOf ?? turnForkOf,
@@ -398,27 +541,39 @@ class DefaultServerTransport<TEvent, TMessage> implements ServerTransport<TEvent
         const encoder = codec.createEncoder(channel, {
           extras: { headers: defaultHeaders },
           onMessage,
+          messageId: msgId,
         });
 
-        const result = await pipeStream(stream, encoder, signal, onAbort, logger);
+        const result = await pipeStream(stream, encoder, signal, onAbort, streamOpts?.resolveWriteOptions, logger);
+
+        if (result.error) {
+          const errInfo = new Ably.ErrorInfo(
+            `unable to stream response for turn ${turnId}; ${result.error.message}`,
+            ErrorCode.StreamError,
+            500,
+            result.error instanceof Ably.ErrorInfo ? result.error : undefined,
+          );
+          logger?.error('Turn.streamResponse(); stream error', { turnId });
+          turnOnError?.(errInfo);
+        }
 
         logger?.debug('Turn.streamResponse(); stream finished', { turnId, reason: result.reason });
         return result;
       },
 
-      // Spec: AIT-ST7
+      // Spec: AIT-ST7, AIT-ST7a, AIT-ST7b
       end: async (reason: TurnEndReason): Promise<void> => {
         logger?.trace('Turn.end();', { turnId, reason });
 
-        if (!started) {
+        if (state === TurnState.INITIALIZED) {
           throw new Ably.ErrorInfo(
             `unable to end turn; start() must be called before end() (turn ${turnId})`,
             ErrorCode.InvalidArgument,
             400,
           );
         }
-        if (ended) return;
-        ended = true;
+        if (state === TurnState.ENDED) return;
+        state = TurnState.ENDED;
 
         try {
           await turnManager.endTurn(turnId, reason);
@@ -430,7 +585,6 @@ class DefaultServerTransport<TEvent, TMessage> implements ServerTransport<TEvent
             error instanceof Ably.ErrorInfo ? error : undefined,
           );
           logger?.error('Turn.end(); failed to publish turn-end', { turnId });
-          turnOnError?.(errInfo);
           throw errInfo;
         } finally {
           registeredTurns.delete(turnId);

package/src/core/transport/stream-router.ts

@@ -2,7 +2,8 @@
  * Client-side stream routing.
  *
  * Maintains a map of turnId to ReadableStreamController. Routes decoded events
- * to the correct stream. Closes streams on terminal events or explicit close.
+ * to the correct stream. Closes streams on terminal events, explicit close, or
+ * error.
  */
 
 import * as Ably from 'ably';
@@ -19,8 +20,10 @@ import type { TurnEntry } from './types.js';
 export interface StreamRouter<TEvent> {
   /** Register a new stream for a turnId. Returns the ReadableStream the consumer reads from. */
   createStream(turnId: string): ReadableStream<TEvent>;
-  /** Close the stream for a turnId. Returns true if a stream was closed. */
+  /** Close the stream for a turnId. Returns true if a stream existed. */
   closeStream(turnId: string): boolean;
+  /** Error the stream for a turnId. The consumer's reader will reject with the given error. Returns true if a stream existed. */
+  errorStream(turnId: string, error: Ably.ErrorInfo): boolean;
   /** Enqueue an event to the correct stream. Returns true if routed successfully. */
   route(turnId: string, event: TEvent): boolean;
   /** Whether a specific turnId has an active stream. */
@@ -73,7 +76,22 @@ class DefaultStreamRouter<TEvent> implements StreamRouter<TEvent> {
     try {
       turn.controller.close();
     } catch {
-      /* already closed */
+      /* consumer cancelled the stream */
+    }
+    this._turns.delete(turnId);
+    return true;
+  }
+
+  // Spec: AIT-CT14c
+  errorStream(turnId: string, error: Ably.ErrorInfo): boolean {
+    const turn = this._turns.get(turnId);
+    if (!turn) return false;
+
+    this._logger.debug('StreamRouter.errorStream(); erroring stream', { turnId });
+    try {
+      turn.controller.error(error);
+    } catch {
+      /* consumer cancelled the stream */
     }
     this._turns.delete(turnId);
     return true;