@ably/ai-transport 0.1.0 → 0.3.0

package/src/core/transport/run-manager.ts

@@ -0,0 +1,243 @@
+/**
+ * Server-side run state management and lifecycle event publishing.
+ *
+ * Owns the authoritative run lifecycle. Tracks active runs with their
+ * AbortControllers and clientIds. Publishes run-start, run-resume, run-suspend, and
+ * run-end events on the Ably channel so all clients can react to run
+ * state changes.
+ */
+
+import type * as Ably from 'ably';
+
+import { EVENT_RUN_END, EVENT_RUN_RESUME, EVENT_RUN_START, EVENT_RUN_SUSPEND } from '../../constants.js';
+import type { Logger } from '../../logger.js';
+import { buildLifecycleHeaders } from './headers.js';
+import type { RunEndReason } from './types.js';
+
+/**
+ * Per-invocation metadata carried on a run's opening lifecycle event. A
+ * continuation (re-entering an existing run) sets `continuation` and omits the
+ * structural `parent` / `forkOf` / `regenerates` fields.
+ */
+interface StartRunMetadata {
+  /** Structural parent codec-message-id (fresh run-start only). */
+  parent?: string;
+  /** Forked user-prompt codec-message-id for an edit (fresh run-start only). */
+  forkOf?: string;
+  /** Regenerated assistant codec-message-id (fresh run-start only). */
+  regenerates?: string;
+  /** Agent-minted invocation id, carried on the lifecycle event. */
+  invocationId?: string;
+  /** ClientId of the triggering input event. */
+  inputClientId?: string;
+  /** Codec-message-id of the triggering input event. */
+  inputCodecMessageId?: string;
+  /** When true, publish `ai-run-resume` (re-entry) instead of `ai-run-start`. */
+  continuation?: boolean;
+}
+
+// ---------------------------------------------------------------------------
+// Interface
+// ---------------------------------------------------------------------------
+
+/** Manages active runs and publishes run lifecycle events on the channel. */
+export interface RunManager {
+  /**
+   * Register a run and publish its opening lifecycle event. Publishes
+   * `ai-run-start` for a fresh run, or `ai-run-resume` when `metadata.continuation`
+   * is set (a subsequent invocation re-entering an existing run). A resume omits
+   * the structural `parent` / `forkOf` / `regenerates` headers — the original
+   * run-start owns the run's structure.
+   */
+  startRun(runId: string, clientId?: string, controller?: AbortController, metadata?: StartRunMetadata): Promise<void>;
+  /**
+   * Suspend a run. Publishes run-suspend on the channel and drops the run's
+   * active-run entry — the agent process terminates on suspend, so there is no
+   * live AbortController to retain. A cancel arriving during suspension is a
+   * no-op; the resuming invocation re-registers the run via {@link startRun}.
+   * Carries the same per-invocation attribution as {@link endRun}
+   * (`inputClientId`, `inputCodecMessageId`), since a suspend is the terminal
+   * event of the suspending invocation just as run-end is of an ending one.
+   */
+  suspendRun(runId: string, invocationId?: string, inputClientId?: string, inputCodecMessageId?: string): Promise<void>;
+  /**
+   * End a run. Publishes run-end on the channel (stamping `reason` as the
+   * run-reason header) and drops the run's active-run entry. Carries the same
+   * per-invocation attribution as {@link suspendRun} (`invocationId`,
+   * `inputClientId`, `inputCodecMessageId`), since run-end is the terminal event
+   * of the ending invocation. When `reason` is `'error'` and an `error` is
+   * supplied, its `code` and `message` are additionally stamped as the
+   * `error-code` / `error-message` headers — a codec-agnostic baseline failure
+   * detail for consumers; omitting `error` publishes a bare `reason: 'error'`.
+   */
+  endRun(
+    runId: string,
+    reason: RunEndReason,
+    invocationId?: string,
+    inputClientId?: string,
+    inputCodecMessageId?: string,
+    error?: Ably.ErrorInfo,
+  ): Promise<void>;
+  /** Get the clientId that owns a run. */
+  getClientId(runId: string): string | undefined;
+  /** Cancel all active runs and clear state. */
+  close(): void;
+}
+
+// ---------------------------------------------------------------------------
+// Internal state
+// ---------------------------------------------------------------------------
+
+interface ActiveRunEntry {
+  controller: AbortController;
+  clientId: string;
+}
+
+// ---------------------------------------------------------------------------
+// Implementation
+// ---------------------------------------------------------------------------
+
+class DefaultRunManager implements RunManager {
+  private readonly _channel: Ably.RealtimeChannel;
+  private readonly _logger: Logger | undefined;
+  private readonly _activeRuns = new Map<string, ActiveRunEntry>();
+
+  constructor(channel: Ably.RealtimeChannel, logger?: Logger) {
+    this._channel = channel;
+    this._logger = logger?.withContext({ component: 'RunManager' });
+  }
+
+  async startRun(
+    runId: string,
+    clientId?: string,
+    externalController?: AbortController,
+    metadata?: StartRunMetadata,
+  ): Promise<void> {
+    this._logger?.trace('DefaultRunManager.startRun();', { runId, clientId });
+
+    const controller = externalController ?? new AbortController();
+    const resolvedClientId = clientId ?? '';
+    this._activeRuns.set(runId, { controller, clientId: resolvedClientId });
+
+    // A continuation re-enters an already-started run: publish `ai-run-resume`
+    // rather than `ai-run-start`. Resume is a pure re-entry signal — the
+    // original run-start already established the run's structure, so the
+    // parent / forkOf / regenerates metadata is NOT re-stamped here (doing so
+    // would point the run at content within itself). The agent learned this is
+    // a continuation from the run-id on the triggering input; the re-entry is
+    // conveyed to clients by the event name, not a header echo. The
+    // invocation-id / input attribution headers are carried on both.
+    const continuation = metadata?.continuation === true;
+
+    const headers = buildLifecycleHeaders({
+      runId,
+      runClientId: resolvedClientId,
+      parent: continuation ? undefined : metadata?.parent,
+      forkOf: continuation ? undefined : metadata?.forkOf,
+      regenerates: continuation ? undefined : metadata?.regenerates,
+      invocationId: metadata?.invocationId,
+      inputClientId: metadata?.inputClientId,
+      inputCodecMessageId: metadata?.inputCodecMessageId,
+    });
+
+    await this._channel.publish({
+      name: continuation ? EVENT_RUN_RESUME : EVENT_RUN_START,
+      extras: { ai: { transport: headers } },
+    });
+
+    this._logger?.debug('DefaultRunManager.startRun(); run started', { runId });
+  }
+
+  async suspendRun(
+    runId: string,
+    invocationId?: string,
+    inputClientId?: string,
+    inputCodecMessageId?: string,
+  ): Promise<void> {
+    this._logger?.trace('DefaultRunManager.suspendRun();', { runId });
+    await this._publishTerminal(EVENT_RUN_SUSPEND, runId, { invocationId, inputClientId, inputCodecMessageId });
+    this._logger?.debug('DefaultRunManager.suspendRun(); run suspended', { runId });
+  }
+
+  async endRun(
+    runId: string,
+    reason: RunEndReason,
+    invocationId?: string,
+    inputClientId?: string,
+    inputCodecMessageId?: string,
+    error?: Ably.ErrorInfo,
+  ): Promise<void> {
+    this._logger?.trace('DefaultRunManager.endRun();', { runId, reason });
+    // Stamp error detail only for a terminal error the agent chose to surface
+    // (AIT-ST6b4: explicit, never automatic). error-code / error-message are
+    // generic transport headers, so any codec or consumer can read them.
+    const errorAttribution = reason === 'error' && error ? { errorCode: error.code, errorMessage: error.message } : {};
+    await this._publishTerminal(EVENT_RUN_END, runId, {
+      reason,
+      invocationId,
+      inputClientId,
+      inputCodecMessageId,
+      ...errorAttribution,
+    });
+    this._logger?.debug('DefaultRunManager.endRun(); run ended', { runId, reason });
+  }
+
+  /**
+   * Publish a run's terminal lifecycle event (run-suspend or run-end) and drop
+   * its active-run entry. Both events are the suspending/ending invocation's
+   * terminal signal, carrying the same per-invocation correlation; they differ
+   * only by event name and the run-reason header (run-end). Publishes BEFORE
+   * dropping local state so a publish failure leaves the run in the active set.
+   * @param eventName - The lifecycle event to publish (run-suspend or run-end).
+   * @param runId - The run being suspended or ended.
+   * @param attribution - Per-invocation correlation and the terminal reason.
+   * @param attribution.reason - Terminal reason; set for run-end, omitted for run-suspend.
+   * @param attribution.invocationId - The invocation's id.
+   * @param attribution.inputClientId - ClientId of the triggering input event.
+   * @param attribution.inputCodecMessageId - Codec-message-id of the triggering input event.
+   * @param attribution.errorCode - Numeric error code; set for run-end only when a terminal error is surfaced.
+   * @param attribution.errorMessage - Error message; paired with errorCode.
+   */
+  private async _publishTerminal(
+    eventName: string,
+    runId: string,
+    attribution: {
+      reason?: RunEndReason;
+      invocationId?: string;
+      inputClientId?: string;
+      inputCodecMessageId?: string;
+      errorCode?: number;
+      errorMessage?: string;
+    },
+  ): Promise<void> {
+    const resolvedClientId = this._activeRuns.get(runId)?.clientId ?? '';
+    const headers = buildLifecycleHeaders({ runId, runClientId: resolvedClientId, ...attribution });
+    await this._channel.publish({ name: eventName, extras: { ai: { transport: headers } } });
+    this._activeRuns.delete(runId);
+  }
+
+  getClientId(runId: string): string | undefined {
+    return this._activeRuns.get(runId)?.clientId;
+  }
+
+  close(): void {
+    this._logger?.trace('DefaultRunManager.close();', { activeRuns: this._activeRuns.size });
+    for (const state of this._activeRuns.values()) {
+      state.controller.abort();
+    }
+    this._activeRuns.clear();
+  }
+}
+
+// ---------------------------------------------------------------------------
+// Factory
+// ---------------------------------------------------------------------------
+
+/**
+ * Create a run manager bound to the given channel.
+ * @param channel - The Ably channel to publish lifecycle events on.
+ * @param logger - Optional logger for diagnostic output.
+ * @returns A new {@link RunManager} instance.
+ */
+export const createRunManager = (channel: Ably.RealtimeChannel, logger?: Logger): RunManager =>
+  new DefaultRunManager(channel, logger);

package/src/core/transport/session-support.ts

@@ -0,0 +1,96 @@
+/**
+ * Shared lifecycle plumbing for the client and agent sessions.
+ *
+ * Both `DefaultClientSession` and `DefaultAgentSession` gate their writes on
+ * `connect()` having run, detach their channel best-effort on close, and react
+ * to channel continuity loss with the same detection rule and error shape.
+ * These helpers own that common machinery so the two sessions cannot drift on
+ * the connection guard, the detach-swallow behaviour, or — most importantly —
+ * the continuity-loss predicate, which encodes channel protocol semantics
+ * (Spec AIT-CT19 / AIT-ST12). Each session keeps its own divergent reaction to
+ * continuity loss (the client emits; the agent aborts runs and swaps its Tree).
+ */
+
+import * as Ably from 'ably';
+
+import { ErrorCode } from '../../errors.js';
+import type { Logger } from '../../logger.js';
+
+/**
+ * Resolve a session's connect guard: return the in-flight/settled connect
+ * promise, or reject with `InvalidArgument` when `connect()` has not been
+ * called. Callers `await` the result before any write.
+ * @param connectPromise - The session's connect promise, or `undefined` when not yet connected.
+ * @param method - The method name being guarded, for the error message.
+ * @returns The connect promise.
+ * @throws {Ably.ErrorInfo} `InvalidArgument` when `connectPromise` is `undefined`.
+ */
+export const requireConnected = async (connectPromise: Promise<void> | undefined, method: string): Promise<void> => {
+  if (!connectPromise) {
+    throw new Ably.ErrorInfo(
+      `unable to ${method}; connect() must be called before ${method}()`,
+      ErrorCode.InvalidArgument,
+      400,
+    );
+  }
+  return connectPromise;
+};
+
+/**
+ * Detach the session's channel on close, best-effort. `connect()` subscribes
+ * (which implicitly attaches), so a detach is only attempted when `connect()`
+ * ran. A detach failure (e.g. the channel is already FAILED) must not throw out
+ * of `close()`, so it is swallowed and logged at debug.
+ * @param channel - The session's channel.
+ * @param connectPromise - The session's connect promise; detach is skipped when `undefined`.
+ * @param logger - Logger for the swallowed-failure debug line, or `undefined`.
+ * @param component - The owning class name, used as the log message prefix.
+ */
+export const bestEffortDetach = async (
+  channel: Ably.RealtimeChannel,
+  connectPromise: Promise<void> | undefined,
+  logger: Logger | undefined,
+  component: string,
+): Promise<void> => {
+  if (connectPromise === undefined) return;
+  try {
+    await channel.detach();
+  } catch (error) {
+    logger?.debug(`${component}.close(); channel detach failed`, { error });
+  }
+};
+
+/**
+ * Whether a channel state change breaks message continuity:
+ * - FAILED, SUSPENDED, DETACHED — no more messages expected (or a gap)
+ * - ATTACHED with `resumed: false` (an UPDATE) — messages were lost
+ *
+ * The initial attach (ATTACHED with no prior attach) is the caller's concern
+ * and is not handled here.
+ * @param stateChange - The channel state change to classify.
+ * @returns True when continuity was lost.
+ */
+export const isContinuityLost = (stateChange: Ably.ChannelStateChange): boolean => {
+  const { current, resumed } = stateChange;
+  return (
+    current === 'failed' || current === 'suspended' || current === 'detached' || (current === 'attached' && !resumed)
+  );
+};
+
+/**
+ * Build the `ChannelContinuityLost` error for a continuity-breaking state
+ * change, attaching the state change's `reason` as `cause`.
+ * @param stateChange - The continuity-breaking state change.
+ * @param verb - The operation that can no longer proceed, for the
+ *   `unable to <verb>; ...` message (e.g. "deliver events", "continue").
+ * @returns The continuity-loss error.
+ */
+export const continuityLostError = (stateChange: Ably.ChannelStateChange, verb: string): Ably.ErrorInfo => {
+  const { current } = stateChange;
+  return new Ably.ErrorInfo(
+    `unable to ${verb}; channel continuity lost (${current}${current === 'attached' ? ', resumed: false' : ''})`,
+    ErrorCode.ChannelContinuityLost,
+    500,
+    stateChange.reason,
+  );
+};