@trpc/server 11.0.0-alpha-tmp-subscription-connection-state.488 → 11.0.0-alpha-tmp-12-06-react.665

package/src/unstable-core-do-not-import/stream/sse.ts

@@ -1,63 +1,27 @@
+import { Unpromise } from '../../vendor/unpromise';
 import { getTRPCErrorFromUnknown } from '../error/TRPCError';
-import type { ValidateShape } from '../types';
-import { run } from '../utils';
-import type { ConsumerOnError } from './jsonl';
-import { createTimeoutPromise } from './utils/createDeferred';
-import { createReadableStream } from './utils/createReadableStream';
+import { isAbortError } from '../http/isAbortError';
+import type { MaybePromise } from '../types';
+import { identity, run } from '../utils';
+import type { EventSourceLike } from './sse.types';
+import type { inferTrackedOutput } from './tracked';
+import { isTrackedEnvelope } from './tracked';
+import { takeWithGrace, withMaxDuration } from './utils/asyncIterable';
+import { makeAsyncResource } from './utils/disposable';
+import { readableStreamFrom } from './utils/readableStreamFrom';
+import {
+  disposablePromiseTimerResult,
+  timerResource,
+} from './utils/timerResource';
+import { PING_SYM, withPing } from './utils/withPing';
 
 type Serialize = (value: any) => any;
 type Deserialize = (value: any) => any;
 
-/**
- * Server-sent Event Message
- * @see https://html.spec.whatwg.org/multipage/server-sent-events.html
- * @public
- */
-export interface SSEMessage {
-  /**
-   * The data field of the message - this can be anything
-   */
-  data: unknown;
-  /**
-   * The id for this message
-   * Passing this id will allow the client to resume the connection from this point if the connection is lost
-   * @see https://html.spec.whatwg.org/multipage/server-sent-events.html#the-last-event-id-header
-   */
-  id: string;
-}
-
-const sseSymbol = Symbol('SSEMessageEnvelope');
-export type SSEMessageEnvelope<TData> = [typeof sseSymbol, TData];
-
-/**
- * Produce a typed server-sent event message
- */
-export function sse<TData extends SSEMessage>(
-  event: ValidateShape<TData, SSEMessage>,
-): SSEMessageEnvelope<TData> {
-  if (event.id === '') {
-    // This could be removed by using different event names for `yield sse(x)`-emitted events and `yield y`-emitted events
-    throw new Error(
-      '`id` must not be an empty string as empty string is the same as not setting the id at all',
-    );
-  }
-  return [sseSymbol, event as TData];
-}
-
-export function isSSEMessageEnvelope<TData extends SSEMessage>(
-  value: unknown,
-): value is SSEMessageEnvelope<TData> {
-  return Array.isArray(value) && value[0] === sseSymbol;
-}
-
-export type SerializedSSEvent = Omit<SSEMessage, 'data'> & {
-  data?: string;
-};
-
 /**
  * @internal
  */
-export interface PingOptions {
+export interface SSEPingOptions {
   /**
    * Enable ping comments sent from the server
    * @default false
@@ -70,14 +34,22 @@ export interface PingOptions {
   intervalMs?: number;
 }
 
-export interface SSEStreamProducerOptions {
+export interface SSEClientOptions {
+  /**
+   * Timeout and reconnect after inactivity in milliseconds
+   * @default undefined
+   */
+  reconnectAfterInactivityMs?: number;
+}
+
+export interface SSEStreamProducerOptions<TValue = unknown> {
   serialize?: Serialize;
-  data: AsyncIterable<unknown>;
+  data: AsyncIterable<TValue>;
+
   maxDepth?: number;
-  ping?: PingOptions;
+  ping?: SSEPingOptions;
   /**
    * Maximum duration in milliseconds for the request before ending the stream
-   * Only useful for serverless runtimes
    * @default undefined
    */
   maxDurationMs?: number;
@@ -87,109 +59,128 @@ export interface SSEStreamProducerOptions {
    * @default false
    */
   emitAndEndImmediately?: boolean;
+  formatError?: (opts: { error: unknown }) => unknown;
+  /**
+   * Client-specific options - these will be sent to the client as part of the first message
+   * @default {}
+   */
+  client?: SSEClientOptions;
 }
 
-type SSEvent = Partial<
-  SSEMessage & {
-    comment: string;
-    event: string;
-  }
->;
+const PING_EVENT = 'ping';
+const SERIALIZED_ERROR_EVENT = 'serialized-error';
+const CONNECTED_EVENT = 'connected';
+
+interface SSEvent {
+  id?: string;
+  data: unknown;
+  comment?: string;
+  event?: string;
+}
 /**
  *
  * @see https://html.spec.whatwg.org/multipage/server-sent-events.html
  */
-export function sseStreamProducer(opts: SSEStreamProducerOptions) {
-  const stream = createReadableStream<SSEvent>();
-  stream.controller.enqueue({
-    comment: 'connected',
-  });
+export function sseStreamProducer<TValue = unknown>(
+  opts: SSEStreamProducerOptions<TValue>,
+) {
+  const { serialize = identity } = opts;
 
-  const { serialize = (v) => v } = opts;
-
-  const ping: Required<PingOptions> = {
+  const ping: Required<SSEPingOptions> = {
     enabled: opts.ping?.enabled ?? false,
     intervalMs: opts.ping?.intervalMs ?? 1000,
   };
+  const client: SSEClientOptions = opts.client ?? {};
 
-  run(async () => {
-    const iterator = opts.data[Symbol.asyncIterator]();
-
-    const closedPromise = stream.cancelledPromise.then(() => 'closed' as const);
-    const maxDurationPromise = createTimeoutPromise(
-      opts.maxDurationMs ?? Infinity,
-      'maxDuration' as const,
+  if (
+    ping.enabled &&
+    client.reconnectAfterInactivityMs &&
+    ping.intervalMs > client.reconnectAfterInactivityMs
+  ) {
+    throw new Error(
+      `Ping interval must be less than client reconnect interval to prevent unnecessary reconnection - ping.intervalMs: ${ping.intervalMs} client.reconnectAfterInactivityMs: ${client.reconnectAfterInactivityMs}`,
     );
+  }
 
-    let nextPromise = iterator.next();
+  async function* generator(): AsyncIterable<SSEvent, void> {
+    yield {
+      event: CONNECTED_EVENT,
+      data: JSON.stringify(client),
+    };
 
-    while (true) {
-      const pingPromise = createTimeoutPromise(
-        ping.enabled ? ping.intervalMs : Infinity,
-        'ping' as const,
-      );
-      const next = await Promise.race([
-        nextPromise.catch(getTRPCErrorFromUnknown),
-        pingPromise.promise,
-        closedPromise,
-        maxDurationPromise.promise,
-      ]);
-
-      pingPromise.clear();
-      if (next === 'closed') {
-        break;
-      }
-      if (next === 'maxDuration') {
-        break;
-      }
+    type TIteratorValue = Awaited<TValue> | typeof PING_SYM;
 
-      if (next === 'ping') {
-        stream.controller.enqueue({
-          comment: 'ping',
-        });
-        continue;
-      }
+    let iterable: AsyncIterable<TValue | typeof PING_SYM> = opts.data;
 
-      if (next instanceof Error) {
-        stream.controller.error(next);
-        break;
-      }
-      if (next.done) {
-        break;
-      }
+    if (opts.emitAndEndImmediately) {
+      iterable = takeWithGrace(iterable, {
+        count: 1,
+        gracePeriodMs: 1,
+      });
+    }
 
-      const value = next.value;
+    if (
+      opts.maxDurationMs &&
+      opts.maxDurationMs > 0 &&
+      opts.maxDurationMs !== Infinity
+    ) {
+      iterable = withMaxDuration(iterable, {
+        maxDurationMs: opts.maxDurationMs,
+      });
+    }
 
-      const chunk: SSEvent = isSSEMessageEnvelope(value)
-        ? { ...value[1] }
-        : {
-            data: value,
-          };
-      if ('data' in chunk) {
-        chunk.data = JSON.stringify(serialize(chunk.data));
-      }
+    if (ping.enabled && ping.intervalMs !== Infinity && ping.intervalMs > 0) {
+      iterable = withPing(iterable, ping.intervalMs);
+    }
 
-      stream.controller.enqueue(chunk);
+    // We need those declarations outside the loop for garbage collection reasons. If they were
+    // declared inside, they would not be freed until the next value is present.
+    let value: null | TIteratorValue;
+    let chunk: null | SSEvent;
 
-      if (opts.emitAndEndImmediately) {
-        // end the stream in the next tick so that we can send a few more events from the queue
-        setTimeout(maxDurationPromise.resolve, 1);
+    for await (value of iterable) {
+      if (value === PING_SYM) {
+        yield { event: PING_EVENT, data: '' };
+        continue;
       }
 
-      nextPromise = iterator.next();
+      chunk = isTrackedEnvelope(value)
+        ? { id: value[0], data: value[1] }
+        : { data: value };
+
+      chunk.data = JSON.stringify(serialize(chunk.data));
+
+      yield chunk;
+
+      // free up references for garbage collection
+      value = null;
+      chunk = null;
     }
-    maxDurationPromise.clear();
-    await iterator.return?.();
+  }
+
+  async function* generatorWithErrorHandling(): AsyncIterable<SSEvent, void> {
     try {
-      stream.controller.close();
-    } catch {}
-  }).catch((error) => {
-    return stream.controller.error(error);
-  });
+      yield* generator();
+    } catch (cause) {
+      if (isAbortError(cause)) {
+        // ignore abort errors, send any other errors
+        return;
+      }
+      // `err` must be caused by `opts.data`, `JSON.stringify` or `serialize`.
+      // So, a user error in any case.
+      const error = getTRPCErrorFromUnknown(cause);
+      const data = opts.formatError?.({ error }) ?? null;
+      yield {
+        event: SERIALIZED_ERROR_EVENT,
+        data: JSON.stringify(serialize(data)),
+      };
+    }
+  }
+  const stream = readableStreamFrom(generatorWithErrorHandling());
 
-  return stream.readable.pipeThrough(
-    new TransformStream<SSEvent, string>({
-      transform(chunk, controller) {
+  return stream.pipeThrough(
+    new TransformStream({
+      transform(chunk, controller: TransformStreamDefaultController<string>) {
         if ('event' in chunk) {
           controller.enqueue(`event: ${chunk.event}\n`);
         }
@@ -207,77 +198,244 @@ export function sseStreamProducer(opts: SSEStreamProducerOptions) {
     }),
   );
 }
-export type inferSSEOutput<TData> = TData extends SSEMessageEnvelope<
-  infer $Data
->
-  ? $Data
-  : TData;
+
+interface ConsumerStreamResultBase<TConfig extends ConsumerConfig> {
+  eventSource: InstanceType<TConfig['EventSource']> | null;
+}
+
+interface ConsumerStreamResultData<TConfig extends ConsumerConfig>
+  extends ConsumerStreamResultBase<TConfig> {
+  type: 'data';
+  data: inferTrackedOutput<TConfig['data']>;
+}
+
+interface ConsumerStreamResultError<TConfig extends ConsumerConfig>
+  extends ConsumerStreamResultBase<TConfig> {
+  type: 'serialized-error';
+  error: TConfig['error'];
+}
+
+interface ConsumerStreamResultConnecting<TConfig extends ConsumerConfig>
+  extends ConsumerStreamResultBase<TConfig> {
+  type: 'connecting';
+  event: EventSourceLike.EventOf<TConfig['EventSource']> | null;
+}
+interface ConsumerStreamResultTimeout<TConfig extends ConsumerConfig>
+  extends ConsumerStreamResultBase<TConfig> {
+  type: 'timeout';
+  ms: number;
+}
+interface ConsumerStreamResultPing<TConfig extends ConsumerConfig>
+  extends ConsumerStreamResultBase<TConfig> {
+  type: 'ping';
+}
+
+interface ConsumerStreamResultConnected<TConfig extends ConsumerConfig>
+  extends ConsumerStreamResultBase<TConfig> {
+  type: 'connected';
+  options: SSEClientOptions;
+}
+
+type ConsumerStreamResult<TConfig extends ConsumerConfig> =
+  | ConsumerStreamResultData<TConfig>
+  | ConsumerStreamResultError<TConfig>
+  | ConsumerStreamResultConnecting<TConfig>
+  | ConsumerStreamResultTimeout<TConfig>
+  | ConsumerStreamResultPing<TConfig>
+  | ConsumerStreamResultConnected<TConfig>;
+
+export interface SSEStreamConsumerOptions<TConfig extends ConsumerConfig> {
+  url: () => MaybePromise<string>;
+  init: () =>
+    | MaybePromise<EventSourceLike.InitDictOf<TConfig['EventSource']>>
+    | undefined;
+  signal: AbortSignal;
+  deserialize?: Deserialize;
+  EventSource: TConfig['EventSource'];
+}
+
+interface ConsumerConfig {
+  data: unknown;
+  error: unknown;
+  EventSource: EventSourceLike.AnyConstructor;
+}
+
+async function withTimeout<T>(opts: {
+  promise: Promise<T>;
+  timeoutMs: number;
+  onTimeout: () => Promise<NoInfer<T>>;
+}): Promise<T> {
+  using timeoutPromise = timerResource(opts.timeoutMs);
+  const res = await Unpromise.race([opts.promise, timeoutPromise.start()]);
+
+  if (res === disposablePromiseTimerResult) {
+    return await opts.onTimeout();
+  }
+  return res;
+}
+
 /**
  * @see https://html.spec.whatwg.org/multipage/server-sent-events.html
  */
-
-export function sseStreamConsumer<TData>(opts: {
-  from: EventSource;
-  onError?: ConsumerOnError;
-  deserialize?: Deserialize;
-}): AsyncIterable<inferSSEOutput<TData>> {
+export function sseStreamConsumer<TConfig extends ConsumerConfig>(
+  opts: SSEStreamConsumerOptions<TConfig>,
+): AsyncIterable<ConsumerStreamResult<TConfig>> {
   const { deserialize = (v) => v } = opts;
-  const eventSource = opts.from;
 
-  const stream = createReadableStream<MessageEvent>();
+  let clientOptions: SSEClientOptions = {};
 
-  const transform = new TransformStream<MessageEvent, inferSSEOutput<TData>>({
-    async transform(chunk, controller) {
-      const def: Partial<SSEMessage> = {
-        data: deserialize(JSON.parse(chunk.data)),
-      };
+  const signal = opts.signal;
 
-      if (chunk.lastEventId) {
-        def.id = chunk.lastEventId;
-      }
-      controller.enqueue(def as inferSSEOutput<TData>);
-    },
-  });
+  let _es: InstanceType<TConfig['EventSource']> | null = null;
 
-  eventSource.addEventListener('message', (msg) => {
-    stream.controller.enqueue(msg);
-  });
-  eventSource.addEventListener('error', (cause) => {
-    if (eventSource.readyState === EventSource.CLOSED) {
-      stream.controller.error(cause);
-    }
-  });
+  const createStream = () =>
+    new ReadableStream<ConsumerStreamResult<TConfig>>({
+      async start(controller) {
+        const [url, init] = await Promise.all([opts.url(), opts.init()]);
+        const eventSource = (_es = new opts.EventSource(
+          url,
+          init,
+        ) as InstanceType<TConfig['EventSource']>);
 
-  const readable = stream.readable.pipeThrough(transform);
-  return {
-    [Symbol.asyncIterator]() {
-      const reader = readable.getReader();
-
-      const iterator: AsyncIterator<inferSSEOutput<TData>> = {
-        async next() {
-          const value = await reader.read();
-          if (value.done) {
-            return {
-              value: undefined,
-              done: true,
-            };
+        controller.enqueue({
+          type: 'connecting',
+          eventSource: _es,
+          event: null,
+        });
+
+        eventSource.addEventListener(CONNECTED_EVENT, (_msg) => {
+          const msg = _msg as EventSourceLike.MessageEvent;
+
+          const options: SSEClientOptions = JSON.parse(msg.data);
+
+          clientOptions = options;
+          controller.enqueue({
+            type: 'connected',
+            options,
+            eventSource,
+          });
+        });
+
+        eventSource.addEventListener(SERIALIZED_ERROR_EVENT, (_msg) => {
+          const msg = _msg as EventSourceLike.MessageEvent;
+
+          controller.enqueue({
+            type: 'serialized-error',
+            error: deserialize(JSON.parse(msg.data)),
+            eventSource,
+          });
+        });
+        eventSource.addEventListener(PING_EVENT, () => {
+          controller.enqueue({
+            type: 'ping',
+            eventSource,
+          });
+        });
+        eventSource.addEventListener('error', (event) => {
+          if (eventSource.readyState === EventSource.CLOSED) {
+            controller.error(event);
+          } else {
+            controller.enqueue({
+              type: 'connecting',
+              eventSource,
+              event,
+            });
           }
-          return {
-            value: value.value,
-            done: false,
+        });
+        eventSource.addEventListener('message', (_msg) => {
+          const msg = _msg as EventSourceLike.MessageEvent;
+
+          const chunk = deserialize(JSON.parse(msg.data));
+
+          const def: SSEvent = {
+            data: chunk,
           };
+          if (msg.lastEventId) {
+            def.id = msg.lastEventId;
+          }
+          controller.enqueue({
+            type: 'data',
+            data: def as inferTrackedOutput<TConfig['data']>,
+            eventSource,
+          });
+        });
+
+        const onAbort = () => {
+          controller.close();
+          eventSource.close();
+        };
+        if (signal.aborted) {
+          onAbort();
+        } else {
+          signal.addEventListener('abort', onAbort);
+        }
+      },
+      cancel() {
+        _es?.close();
+      },
+    });
+
+  const getStreamResource = () => {
+    let stream = createStream();
+    let reader = stream.getReader();
+
+    async function dispose() {
+      await reader.cancel();
+      _es = null;
+    }
+
+    return makeAsyncResource(
+      {
+        read() {
+          return reader.read();
         },
-        async return() {
-          reader.releaseLock();
-          return {
-            value: undefined,
-            done: true,
-          };
+        async recreate() {
+          await dispose();
+
+          stream = createStream();
+          reader = stream.getReader();
         },
-      };
-      return iterator;
-    },
+      },
+      dispose,
+    );
   };
+
+  return run(async function* () {
+    await using stream = getStreamResource();
+
+    while (true) {
+      let promise = stream.read();
+
+      const timeoutMs = clientOptions.reconnectAfterInactivityMs;
+      if (timeoutMs) {
+        promise = withTimeout({
+          promise,
+          timeoutMs,
+          onTimeout: async () => {
+            const res: Awaited<typeof promise> = {
+              value: {
+                type: 'timeout',
+                ms: timeoutMs,
+                eventSource: _es,
+              },
+              done: false,
+            };
+            // Close and release old reader
+            await stream.recreate();
+
+            return res;
+          },
+        });
+      }
+
+      const result = await promise;
+
+      if (result.done) {
+        return result.value;
+      }
+      yield result.value;
+    }
+  });
 }
 
 export const sseHeaders = {

package/src/unstable-core-do-not-import/stream/sse.types.ts

@@ -0,0 +1,44 @@
+/* eslint-disable @typescript-eslint/no-namespace */
+
+/**
+ * @internal
+ */
+export namespace EventSourceLike {
+  export interface InitDict {
+    withCredentials?: boolean;
+  }
+
+  export interface MessageEvent extends Event {
+    data: any;
+    lastEventId?: string;
+  }
+  export interface Event {}
+
+  type EventSourceListenerLike = (event: Event) => void;
+
+  export type AnyConstructorLike<TInit extends InitDict> = new (
+    url: string,
+    eventSourceInitDict?: TInit,
+  ) => Instance;
+
+  export interface Instance {
+    readonly CLOSED: number;
+    readonly CONNECTING: number;
+    readonly OPEN: number;
+
+    addEventListener(type: string, listener: EventSourceListenerLike): void;
+    removeEventListener(type: string, listener: EventSourceListenerLike): void;
+    close: () => void;
+
+    readyState: number;
+  }
+
+  export type AnyConstructor = AnyConstructorLike<any>;
+
+  export type ListenerOf<T extends AnyConstructor> = Parameters<
+    InstanceType<T>['addEventListener']
+  >[1];
+  export type EventOf<T extends AnyConstructor> = Parameters<ListenerOf<T>>[0];
+  export type InitDictOf<T extends AnyConstructor> =
+    ConstructorParameters<T>[1];
+}