chore(PLA-4098): merge WebSocket gateway changes into canary

.changeset/long-cougars-play.md

@@ -0,0 +1,5 @@
+---
+"@knocklabs/client": minor
+---
+
+Connect sockets via WebSocket gateway

.changeset/ripe-pandas-eat.md

@@ -0,0 +1,15 @@
+---
+"@knocklabs/client": minor
+---
+
+Add `maxRetries` option to `ApiRequestConfig`
+
+When making a request to the Knock API using `ApiClient`'s `makeRequest` function, the `maxRetries` option of `ApiRequestConfig` can now be used to set the maximum number of times the request should be retried when a retryable error occurs. To disable retries altogether, set `maxRetries` to 0.
+
+```js
+const result = await knock.client().makeRequest({
+  method: "GET",
+  url: "/v1/users",
+  maxRetries: 0,
+});
+```

packages/client/src/api.ts

@@ -1,6 +1,10 @@
 import { Socket } from "phoenix";
 
-import { exponentialBackoffFullJitter } from "./helpers";
+import {
+  createGatewayTransport,
+  fetchSocketConnectUrl,
+} from "./gatewayTransport";
+import { cacheUntilExpired, exponentialBackoffFullJitter } from "./helpers";
 import { PageVisibilityManager } from "./pageVisibility";
 
 // Ceiling for socket reconnect backoff. This is intentionally much larger than
@@ -19,6 +23,10 @@ const STABLE_CONNECTION_THRESHOLD_MS = 30_000;
 // counting.
 const MAX_UNSTABLE_CONNECTION_COUNT = 15;
 
+// The max number of times a request is retried after a retryable failure,
+// unless overridden by `ApiRequestConfig.maxRetries`.
+const DEFAULT_MAX_REQUEST_RETRIES = 3;
+
 type ApiClientOptions = {
   host: string;
   apiKey: string;
@@ -44,6 +52,8 @@ export type ApiRequestConfig = {
   data?: unknown;
   headers?: HeadersInit;
   signal?: AbortSignal;
+  /** Max number of retries for retryable failures. Set this to 0 to disable retries. */
+  maxRetries?: number;
 };
 
 type FetchClient = (
@@ -117,7 +127,21 @@ class ApiClient {
     });
 
     if (typeof window !== "undefined" && typeof WebSocket !== "undefined") {
+      // Connect sockets via the WebSocket gateway API endpoint, which returns a signed ws:// or
+      // wss:// connect URL. This signed URL is cached and used to reconnect until it expires.
+      const connectUrlCache = cacheUntilExpired((signal: AbortSignal) =>
+        // If this fetch rejects, the error is handled by the gateway transport, and Phoenix
+        // considers it a failed connection attempt
+        fetchSocketConnectUrl({ apiClient: this, signal }),
+      );
+      const transport = createGatewayTransport({
+        baseTransport: WebSocket,
+        fetchConnectUrl: connectUrlCache.get,
+        onConnectUrlInvalid: connectUrlCache.invalidate,
+      });
+
       this.socket = new Socket(`${this.host.replace("http", "ws")}/ws/v1`, {
+        transport: transport,
         params: {
           user_token: this.userToken,
           api_key: this.apiKey,
@@ -249,9 +273,11 @@ class ApiClient {
   private async requestWithRetries(req: ApiRequestConfig) {
     let lastError: unknown;
 
+    const { maxRetries = DEFAULT_MAX_REQUEST_RETRIES } = req;
+
     // Sequential retry loop: each attempt awaits in order and returns early on
     // success or a non-retryable error, so it doesn't reduce to an array method.
-    for (let attempt = 0; attempt <= 3; attempt++) {
+    for (let attempt = 0; attempt <= maxRetries; attempt++) {
       let retryAfterMs = 0;
 
       try {
@@ -279,7 +305,7 @@ class ApiClient {
         }
       }
 
-      if (attempt < 3) {
+      if (attempt < maxRetries) {
         await this.delay(this.getRetryDelay(attempt + 1, retryAfterMs));
       }
     }

packages/client/src/gatewayTransport.ts

@@ -0,0 +1,272 @@
+import type ApiClient from "./api";
+import { appendQueryString, getQueryString } from "./helpers";
+
+/**
+ * The signature that Phoenix expects socket transports to conform to, based on the native browser
+ * `WebSocket`. Only fields and methods that Phoenix relies on are included here.
+ */
+export type WebSocketLike = Pick<
+  WebSocket,
+  | "readyState"
+  | "bufferedAmount"
+  | "binaryType"
+  | "onopen"
+  | "onclose"
+  | "onmessage"
+  | "onerror"
+  | "send"
+  | "close"
+>;
+
+export type TransportConstructor = new (
+  url: string,
+  protocols?: string | string[],
+) => WebSocketLike;
+
+export type GatewayConnectResult = {
+  signedUrl: string;
+  expiresAt: number;
+};
+
+export type GatewayTransportOptions = {
+  /**
+   * The underlying transport constructor. An instance of this transport will be created once a
+   * connect URL has been fetched.
+   */
+  baseTransport: TransportConstructor;
+  /**
+   * Resolves the signed WebSocket URL to connect to. Rejects on failure.
+   */
+  fetchConnectUrl: (signal: AbortSignal) => Promise<GatewayConnectResult>;
+  /**
+   * Called when a connection made with the fetched URL closed before it ever
+   * opened, which most likely means the URL was rejected (though a network blip
+   * could look the same).
+   */
+  onConnectUrlInvalid?: () => void;
+};
+
+/**
+ * All possible socket states. Copied from Phoenix, based on the `readyState` property of
+ * `WebSocket`.
+ * @see {@link https://github.com/phoenixframework/phoenix/blob/1bae3a639f84cd97ccb6b984a8ef4fcdc7aa48ae/assets/js/phoenix/constants.js#L5}
+ */
+export const SOCKET_STATES = {
+  connecting: 0,
+  open: 1,
+  closing: 2,
+  closed: 3,
+} as const;
+
+// WebSocket close event status codes.
+// See https://www.rfc-editor.org/info/rfc6455/#section-7.4.1
+const CLOSE_CODE_NORMAL = 1000;
+const CLOSE_CODE_ABNORMAL = 1006;
+
+/**
+ * Requests a signed WebSocket URL from the WebSocket gateway API endpoint. Rejects on any HTTP
+ * error or malformed body so the transport treats it like a failed connection attempt.
+ */
+export async function fetchSocketConnectUrl({
+  apiClient,
+  signal,
+}: {
+  apiClient: ApiClient;
+  signal: AbortSignal;
+}): Promise<GatewayConnectResult> {
+  const result = await apiClient.makeRequest({
+    method: "POST",
+    url: "/v1/sockets/connect",
+    signal,
+    // Set maxRetries to 0, so that the API client's own retry logic is never used. Phoenix
+    // handles retries automatically via the socket's reconnect backoff.
+    maxRetries: 0,
+  });
+
+  if (result.statusCode !== "ok") {
+    throw (
+      result.error ??
+      new Error(`Socket connect request failed with status ${result.status}`)
+    );
+  }
+
+  const body = result.body as
+    | Partial<{ url: string; expires_at: string }>
+    | undefined;
+
+  const expiresAt =
+    typeof body?.expires_at === "string" ? Date.parse(body.expires_at) : NaN;
+
+  if (typeof body?.url !== "string" || !body.url || Number.isNaN(expiresAt)) {
+    throw new Error("Invalid socket connect response");
+  }
+
+  return { signedUrl: body.url, expiresAt };
+}
+
+export function createGatewayTransport({
+  baseTransport,
+  fetchConnectUrl,
+  onConnectUrlInvalid,
+}: GatewayTransportOptions): TransportConstructor {
+  // Use a named class, since Phoenix logs using the class name
+  return class GatewayTransport implements WebSocketLike {
+    binaryType: BinaryType = "blob";
+
+    onopen: ((event: Event) => void) | null = null;
+    onclose: ((event: CloseEvent) => void) | null = null;
+    onmessage: ((event: MessageEvent) => void) | null = null;
+    onerror: ((event: unknown) => void) | null = null;
+
+    private inner: WebSocketLike | null = null;
+    private state: number = SOCKET_STATES.connecting;
+    private readonly abort = new AbortController();
+    private readonly unsignedQueryString: string;
+    private opened = false;
+    private closeRequested = false;
+    private settled = false;
+    private readonly timeoutId: ReturnType<typeof setTimeout>;
+
+    constructor(
+      url: string,
+      private readonly protocols?: string | string[],
+    ) {
+      // The base URL passed in by Phoenix contains some query params we want to preserve, including
+      // vsn (protocol version to send on connect), api_key, user_token, and branch_slug. This query
+      // string is considered "unsigned" because it lacks an auth_token query param.
+      this.unsignedQueryString = getQueryString(url);
+
+      // Wait for the caller to finish setting up callbacks and properties. This pattern was copied
+      // from Phoenix's LongPoll transport.
+      this.timeoutId = setTimeout(() => this.start(), 0);
+    }
+
+    get readyState(): number {
+      return this.inner ? this.inner.readyState : this.state;
+    }
+
+    get bufferedAmount(): number {
+      return this.inner ? this.inner.bufferedAmount : 0;
+    }
+
+    send(data: string | ArrayBufferLike | ArrayBufferView | Blob): void {
+      if (!this.inner) {
+        // Phoenix only sends when a connection is open, so prevent calls to send before a
+        // connection is established.
+        throw new Error(
+          "GatewayTransport: cannot send before the connection is open",
+        );
+      }
+      this.inner.send(data);
+    }
+
+    close(code?: number, reason?: string): void {
+      if (this.closeRequested || this.settled) {
+        return;
+      }
+
+      this.closeRequested = true;
+
+      // In case start() has not run yet, clear the timeout used to invoke it in the constructor
+      clearTimeout(this.timeoutId);
+
+      if (this.inner) {
+        this.inner.close(code, reason);
+        return;
+      }
+
+      // The socket is closing before the fetch completed. This is still considered a clean close.
+      this.abort.abort();
+      this.state = SOCKET_STATES.closed;
+      this.emitClose(CLOSE_CODE_NORMAL, "", true);
+    }
+
+    private async start(): Promise<void> {
+      // Phoenix may connect() and disconnect() in the same tick; teardown then calls close() on us
+      // before this runs. This prevents a fetch that would be discarded.
+      if (this.closeRequested) {
+        return;
+      }
+
+      let result: GatewayConnectResult;
+      try {
+        result = await fetchConnectUrl(this.abort.signal);
+      } catch (error) {
+        if (this.closeRequested) {
+          // Ignore this error; the socket was already closed cleanly.
+          return;
+        }
+        // An unexepected error occurred while fetching
+        this.fail(error);
+        return;
+      }
+
+      // The socket was closed while a fetch was in-progress
+      if (this.closeRequested) {
+        return;
+      }
+
+      const connectUrl = appendQueryString(
+        result.signedUrl,
+        this.unsignedQueryString,
+      );
+
+      try {
+        this.inner = this.initBaseTransport(connectUrl);
+      } catch (error) {
+        // The base transport's constructor threw an exception. We never even attempted a
+        // connection, so onclose won't fire. The connect URL must be invalid.
+        onConnectUrlInvalid?.();
+        this.fail(error);
+      }
+    }
+
+    private initBaseTransport(url: string): WebSocketLike {
+      const inner = new baseTransport(url, this.protocols);
+
+      inner.binaryType = this.binaryType;
+      inner.onopen = (ev) => {
+        this.opened = true;
+        this.onopen?.(ev);
+      };
+      inner.onmessage = (ev) => this.onmessage?.(ev);
+      inner.onerror = (ev) => this.onerror?.(ev);
+      inner.onclose = (ev) => {
+        if (!this.opened && !this.closeRequested) {
+          // A connection was never opened and this close event was unexpected.
+          // The signed connect URL is likely invalid.
+          onConnectUrlInvalid?.();
+        }
+        this.settled = true;
+        this.onclose?.(ev);
+      };
+
+      return inner;
+    }
+
+    private fail(error: unknown): void {
+      this.state = SOCKET_STATES.closed;
+      this.onerror?.(error);
+      this.emitClose(CLOSE_CODE_ABNORMAL, "gateway_connect_failed", false);
+    }
+
+    private emitClose(code: number, reason: string, wasClean: boolean): void {
+      // Check this.settled to prevent triggering onclose more than once. For example, this could
+      // occur if the transport is closed manually inside of an onerror callback.
+      if (this.settled) {
+        return;
+      }
+
+      this.settled = true;
+
+      // Mirrors Phoenix's LongPoll.close: CloseEvent exists in browsers.
+      // Fall back to a plain object elsewhere (e.g. React Native).
+      const opts = { code, reason, wasClean };
+      if (typeof CloseEvent !== "undefined") {
+        this.onclose?.(new CloseEvent("close", opts));
+      } else {
+        this.onclose?.(opts as CloseEvent);
+      }
+    }
+  };
+}