SSE auto-reconnect for axons; AbortSignal for poll-based waits (#765)

Co-authored-by: Cursor Agent <cursoragent@cursor.com>
Co-authored-by: albert <alb-rl@users.noreply.github.com>

Author: alb-rl

src/core.ts

@@ -51,7 +51,8 @@ export type { _Array as Array, _Record as Record };
 
 type PromiseOrValue<T> = T | Promise<T>;
 
-type APIResponseProps = {
+/** @internal Response metadata paired with each HTTP request; used by streaming helpers. */
+export type APIResponseProps = {
   response: Response;
   options: FinalRequestOptions;
   controller: AbortController;
@@ -130,6 +131,11 @@ export class APIPromise<T> extends Promise<T> {
     );
   }
 
+  /** @internal Same promise backing {@link asResponse}; includes the full request context. */
+  _getResponseProps(): Promise<APIResponseProps> {
+    return this.responsePromise;
+  }
+
   /**
    * Gets the raw `Response` instance instead of parsing the response
    * data.
@@ -189,6 +195,16 @@ export class APIPromise<T> extends Promise<T> {
   }
 }
 
+/**
+ * @internal Promise for a {@link Stream} built after the HTTP response is available
+ * (e.g. SSE with reconnect). Preserves {@link APIPromise} helpers like {@link APIPromise.asResponse}.
+ */
+export class StreamBackedAPIPromise<T> extends APIPromise<T> {
+  constructor(responseProps: Promise<APIResponseProps>, getData: () => Promise<T>) {
+    super(responseProps, () => getData());
+  }
+}
+
 export abstract class APIClient {
   baseURL: string;
   #baseURLOverridden: boolean;

src/lib/polling.ts

@@ -10,6 +10,8 @@ export interface PollingOptions<T> {
   maxAttempts?: number;
   /** Optional timeout for the entire polling operation (in milliseconds) */
   timeoutMs?: number;
+  /** Optional AbortSignal to cancel polling and in-flight delays between attempts. */
+  signal?: AbortSignal | null | undefined;
   /**
    * Condition to check if polling should stop
    * Return true when the condition is met and polling should stop
@@ -196,9 +198,29 @@ export class LongPollAbortError extends Error {
 }
 
 /**
- * Delay execution for specified milliseconds
+ * Delay execution for specified milliseconds, aborting with {@link LongPollAbortError} when `signal` is aborted.
  */
-const delay = (ms: number): Promise<void> => new Promise((resolve) => setTimeout(resolve, ms));
+function abortableDelay(
+  ms: number,
+  signal: AbortSignal | null | undefined,
+  lastResult: unknown,
+): Promise<void> {
+  return new Promise((resolve, reject) => {
+    if (signal?.aborted) {
+      reject(new LongPollAbortError('Polling aborted', lastResult));
+      return;
+    }
+    const timer = setTimeout(() => {
+      signal?.removeEventListener('abort', onAbort);
+      resolve();
+    }, ms);
+    const onAbort = () => {
+      clearTimeout(timer);
+      reject(new LongPollAbortError('Polling aborted', lastResult));
+    };
+    signal?.addEventListener('abort', onAbort, { once: true });
+  });
+}
 
 /**
  * Generic polling function that handles polling logic with configurable options
@@ -212,11 +234,23 @@ export async function poll<T>(
   pollingRequest: () => Promise<T>,
   options: PollingOptions<T> = {},
 ): Promise<T> {
-  const { initialDelayMs, pollingIntervalMs, maxAttempts, timeoutMs, shouldStop, onPollingAttempt, onError } =
-    {
-      ...DEFAULT_OPTIONS,
-      ...options,
-    };
+  const {
+    initialDelayMs,
+    pollingIntervalMs,
+    maxAttempts,
+    timeoutMs,
+    shouldStop,
+    onPollingAttempt,
+    onError,
+    signal,
+  } = {
+    ...DEFAULT_OPTIONS,
+    ...options,
+  };
+
+  if (signal?.aborted) {
+    throw new LongPollAbortError('Polling aborted', undefined);
+  }
 
   if (initialDelayMs !== undefined && initialDelayMs < 0) {
     throw new Error('initialDelayMs must be non-negative');
@@ -253,10 +287,35 @@ export async function poll<T>(
   const raceTimeout = <R>(promise: Promise<R>): Promise<R> =>
     timeoutPromise ? Promise.race([promise, timeoutPromise]) : promise;
 
+  const raceAbort = <R>(promise: Promise<R>): Promise<R> => {
+    if (!signal) return promise;
+    const abortSignal = signal;
+    if (abortSignal.aborted) {
+      return Promise.reject(new LongPollAbortError('Polling aborted', lastResult));
+    }
+    return new Promise((resolve, reject) => {
+      const onAbort = () => {
+        abortSignal.removeEventListener('abort', onAbort);
+        reject(new LongPollAbortError('Polling aborted', lastResult));
+      };
+      abortSignal.addEventListener('abort', onAbort);
+      promise.then(
+        (value) => {
+          abortSignal.removeEventListener('abort', onAbort);
+          resolve(value);
+        },
+        (err) => {
+          abortSignal.removeEventListener('abort', onAbort);
+          reject(err);
+        },
+      );
+    });
+  };
+
   try {
     let result: T;
     try {
-      result = await raceTimeout(initialRequest());
+      result = await raceAbort(raceTimeout(initialRequest()));
     } catch (error) {
       if (onError && error instanceof APIError) {
         result = onError(error);
@@ -275,7 +334,7 @@ export async function poll<T>(
       return result;
     }
 
-    await raceTimeout(delay(initialDelayMs!));
+    await raceAbort(raceTimeout(abortableDelay(initialDelayMs!, signal, lastResult)));
 
     let attempts = 0;
 
@@ -284,7 +343,7 @@ export async function poll<T>(
       ++attempts;
 
       try {
-        result = await raceTimeout(pollingRequest());
+        result = await raceAbort(raceTimeout(pollingRequest()));
       } catch (error) {
         if (onError && error instanceof APIError) {
           result = onError(error);
@@ -304,7 +363,7 @@ export async function poll<T>(
         throw new MaxAttemptsExceededError(`Polling exceeded maximum attempts (${maxAttempts})`, result);
       }
 
-      await raceTimeout(delay(pollingIntervalMs!));
+      await raceAbort(raceTimeout(abortableDelay(pollingIntervalMs!, signal, lastResult)));
     }
 
     // This should only be reachable if maxAttempts is defined

src/lib/streaming-reconnection.ts

@@ -1,35 +1,58 @@
 import { Stream } from '../streaming';
+import { APIPromise, StreamBackedAPIPromise, type APIResponseProps } from '../core';
 
 /**
  * Wraps a stream with automatic reconnection on timeout.
+ * Returns an {@link APIPromise} so callers can use {@link APIPromise.asResponse} and
+ * {@link APIPromise.withResponse} like other streaming endpoints.
  */
-export async function withStreamAutoReconnect<Item>(
-  streamCreator: (offset: number | undefined) => Promise<Stream<Item>>,
+export function withStreamAutoReconnect<Item>(
+  streamCreator: (offset: number | undefined) => APIPromise<Stream<Item>>,
   getOffset: (item: Item) => number | undefined,
-): Promise<Stream<Item>> {
-  let lastOffset: number | undefined = undefined;
-  let currentStream = await streamCreator(lastOffset);
+): StreamBackedAPIPromise<Stream<Item>> {
+  let firstRequest: APIPromise<Stream<Item>> | undefined;
+  const ensureFirst = () => (firstRequest ??= streamCreator(undefined));
 
-  async function* createReconnectingIterator(): AsyncIterator<Item> {
-    while (true) {
-      try {
-        for await (const item of currentStream) {
-          if (getOffset(item) !== undefined) {
-            lastOffset = getOffset(item);
+  // Defer the first HTTP request until something awaits this promise (avoids eager
+  // connection attempts and unhandled rejections when the caller only attaches later).
+  const responsePropsPromise = new Promise<APIResponseProps>((resolve, reject) => {
+    queueMicrotask(() => {
+      ensureFirst()._getResponseProps().then(resolve, reject);
+    });
+  });
+
+  let dataPromiseMemo: Promise<Stream<Item>> | undefined;
+  const getDataPromise = () => {
+    if (!dataPromiseMemo) {
+      dataPromiseMemo = (async () => {
+        let lastOffset: number | undefined = undefined;
+        let currentStream = await ensureFirst();
+
+        async function* createReconnectingIterator(): AsyncIterator<Item> {
+          while (true) {
+            try {
+              for await (const item of currentStream) {
+                if (getOffset(item) !== undefined) {
+                  lastOffset = getOffset(item);
+                }
+                yield item;
+              }
+              return; // Stream completed normally
+            } catch (error) {
+              if ((error as any)?.status === 408) {
+                currentStream = await streamCreator(lastOffset);
+                continue;
+              }
+              throw error; // Not a timeout, rethrow
+            }
           }
-          yield item;
         }
-        return; // Stream completed normally
-      } catch (error) {
-        if ((error as any)?.status === 408) {
-          // Reconnect with the last known offset
-          currentStream = await streamCreator(lastOffset);
-          continue;
-        }
-        throw error; // Not a timeout, rethrow
-      }
+
+        return new Stream(createReconnectingIterator, currentStream.controller);
+      })();
     }
-  }
+    return dataPromiseMemo;
+  };
 
-  return new Stream(createReconnectingIterator, currentStream.controller);
+  return new StreamBackedAPIPromise(responsePropsPromise, getDataPromise);
 }

src/resources/axons.ts

@@ -5,6 +5,7 @@ import { isRequestOptions } from '../core';
 import { APIPromise } from '../core';
 import * as Core from '../core';
 import { Stream } from '../streaming';
+import { withStreamAutoReconnect } from '@runloop/api-client/lib/streaming-reconnection';
 
 export class Axons extends APIResource {
   /**
@@ -49,19 +50,36 @@ export class Axons extends APIResource {
 
   /**
    * [Beta] Subscribe to an axon event stream via server-sent events.
+   * On idle timeout (408), reconnects with `after_sequence` derived from the last
+   * received event (internal to {@link withStreamAutoReconnect}).
    */
   subscribeSse(id: string, options?: Core.RequestOptions): APIPromise<Stream<AxonEventView>> {
-    const defaultHeaders = {
-      Accept: 'text/event-stream',
-    };
     const mergedOptions: Core.RequestOptions = {
-      headers: defaultHeaders,
       ...options,
+      headers: {
+        Accept: 'text/event-stream',
+        ...options?.headers,
+      },
+    };
+    const { query: userQuery, ...restMerged } = mergedOptions;
+    const getStream: (afterSequence: number | undefined) => APIPromise<Stream<AxonEventView>> = (
+      afterSequence,
+    ) => {
+      const base =
+        userQuery && typeof userQuery === 'object' && !Array.isArray(userQuery) ?
+          { ...(userQuery as Record<string, string | undefined>) }
+        : {};
+      const query =
+        afterSequence !== undefined ? { ...base, after_sequence: afterSequence.toString() }
+        : Object.keys(base).length > 0 ? base
+        : undefined;
+      return this._client.get(`/v1/axons/${id}/subscribe/sse`, {
+        ...restMerged,
+        ...(query ? { query } : {}),
+        stream: true,
+      }) as APIPromise<Stream<AxonEventView>>;
     };
-    return this._client.get(`/v1/axons/${id}/subscribe/sse`, {
-      ...mergedOptions,
-      stream: true,
-    }) as APIPromise<Stream<AxonEventView>>;
+    return withStreamAutoReconnect(getStream, (item) => item.sequence);
   }
 }
 

src/resources/axons/axons.ts

@@ -19,6 +19,7 @@ import {
 } from './sql';
 import { AxonsCursorIDPage, type AxonsCursorIDPageParams } from '../../pagination';
 import { Stream } from '../../streaming';
+import { withStreamAutoReconnect } from '@runloop/api-client/lib/streaming-reconnection';
 
 export class Axons extends APIResource {
   sql: SqlAPI.Sql = new SqlAPI.Sql(this._client);
@@ -76,6 +77,8 @@ export class Axons extends APIResource {
 
   /**
    * [Beta] Subscribe to an axon event stream via server-sent events.
+   * On idle timeout (408), reconnects with `after_sequence` derived from the last
+   * received event (internal to {@link withStreamAutoReconnect}).
    */
   subscribeSse(id: string, options?: Core.RequestOptions): APIPromise<Stream<AxonEventView>> {
     const mergedOptions: Core.RequestOptions = {
@@ -85,10 +88,25 @@ export class Axons extends APIResource {
         ...options?.headers,
       },
     };
-    return this._client.get(`/v1/axons/${id}/subscribe/sse`, {
-      ...mergedOptions,
-      stream: true,
-    }) as APIPromise<Stream<AxonEventView>>;
+    const { query: userQuery, ...restMerged } = mergedOptions;
+    const getStream: (afterSequence: number | undefined) => APIPromise<Stream<AxonEventView>> = (
+      afterSequence,
+    ) => {
+      const base =
+        userQuery && typeof userQuery === 'object' && !Array.isArray(userQuery) ?
+          { ...(userQuery as Record<string, string | undefined>) }
+        : {};
+      const query =
+        afterSequence !== undefined ? { ...base, after_sequence: afterSequence.toString() }
+        : Object.keys(base).length > 0 ? base
+        : undefined;
+      return this._client.get(`/v1/axons/${id}/subscribe/sse`, {
+        ...restMerged,
+        ...(query ? { query } : {}),
+        stream: true,
+      }) as APIPromise<Stream<AxonEventView>>;
+    };
+    return withStreamAutoReconnect(getStream, (item) => item.sequence);
   }
 }