@replit/river 0.21.1 → 0.23.0

package/dist/transport/impls/ws/client.cjs

@@ -1,9 +1,7 @@
 "use strict";
-var __create = Object.create;
 var __defProp = Object.defineProperty;
 var __getOwnPropDesc = Object.getOwnPropertyDescriptor;
 var __getOwnPropNames = Object.getOwnPropertyNames;
-var __getProtoOf = Object.getPrototypeOf;
 var __hasOwnProp = Object.prototype.hasOwnProperty;
 var __export = (target, all) => {
   for (var name in all)
@@ -17,14 +15,6 @@ var __copyProps = (to, from, except, desc) => {
   }
   return to;
 };
-var __toESM = (mod, isNodeMode, target) => (target = mod != null ? __create(__getProtoOf(mod)) : {}, __copyProps(
-  // If the importer is in node compatibility mode or this is not an ESM
-  // file that has been converted to a CommonJS file using a Babel-
-  // compatible transform (i.e. "__esModule" has not been set), then set
-  // "default" to the CommonJS "module.exports" for node compatibility.
-  isNodeMode || !mod || !mod.__esModule ? __defProp(target, "default", { value: mod, enumerable: true }) : target,
-  mod
-));
 var __toCommonJS = (mod) => __copyProps(__defProp({}, "__esModule", { value: true }), mod);
 
 // transport/impls/ws/client.ts
@@ -33,11 +23,9 @@ __export(client_exports, {
   WebSocketClientTransport: () => WebSocketClientTransport
 });
 module.exports = __toCommonJS(client_exports);
-var import_agnostic_ws2 = __toESM(require("agnostic-ws"), 1);
 
 // transport/transport.ts
 var import_value = require("@sinclair/typebox/value");
-var import_api2 = require("@opentelemetry/api");
 
 // transport/message.ts
 var import_typebox = require("@sinclair/typebox");
@@ -156,17 +144,71 @@ var EventDispatcher = class {
 
 // transport/session.ts
 var import_nanoid2 = require("nanoid");
+
+// tracing/index.ts
+var import_api = require("@opentelemetry/api");
+
+// package.json
+var version = "0.23.0";
+
+// tracing/index.ts
+function getPropagationContext(ctx) {
+  const tracing = {
+    traceparent: "",
+    tracestate: ""
+  };
+  import_api.propagation.inject(ctx, tracing);
+  return tracing;
+}
+function createSessionTelemetryInfo(session, propagationCtx) {
+  const ctx = propagationCtx ? import_api.propagation.extract(import_api.context.active(), propagationCtx) : import_api.context.active();
+  const span = tracer.startSpan(
+    `session ${session.id}`,
+    {
+      attributes: {
+        component: "river",
+        "river.session.id": session.id,
+        "river.session.to": session.to,
+        "river.session.from": session.from
+      }
+    },
+    ctx
+  );
+  return { span, ctx };
+}
+function createConnectionTelemetryInfo(connection, sessionSpan) {
+  const ctx = import_api.trace.setSpan(import_api.context.active(), sessionSpan);
+  const span = tracer.startSpan(
+    `connection ${connection.id}`,
+    {
+      attributes: {
+        component: "river",
+        "river.connection.id": connection.id
+      },
+      links: [{ context: sessionSpan.spanContext() }]
+    },
+    ctx
+  );
+  return { span, ctx };
+}
+var tracer = import_api.trace.getTracer("river", version);
+var tracing_default = tracer;
+
+// transport/session.ts
+var import_api2 = require("@opentelemetry/api");
 var nanoid2 = (0, import_nanoid2.customAlphabet)("1234567890abcdefghijklmnopqrstuvxyz", 6);
 var unsafeId = () => nanoid2();
 var Connection = class {
-  debugId;
+  id;
+  telemetry;
   constructor() {
-    this.debugId = `conn-${unsafeId()}`;
+    this.id = `conn-${nanoid2(12)}`;
   }
 };
 var Session = class {
   codec;
   options;
+  telemetry;
   /**
    * The buffer of messages that have been sent but not yet acknowledged.
    */
@@ -186,12 +228,6 @@ var Session = class {
    * for this session.
    */
   advertisedSessionId;
-  /**
-   * The metadata for this session, as parsed from the handshake.
-   *
-   * Will only ever be populated on the server side.
-   */
-  metadata;
   /**
    * Number of messages we've sent along this session (excluding handshake and acks)
    */
@@ -213,7 +249,7 @@ var Session = class {
    * The interval for sending heartbeats.
    */
   heartbeat;
-  constructor(conn, from, to, options) {
+  constructor(conn, from, to, options, propagationCtx) {
     this.id = `session-${nanoid2(12)}`;
     this.options = options;
     this.from = from;
@@ -225,13 +261,14 @@ var Session = class {
       () => this.sendHeartbeat(),
       options.heartbeatIntervalMs
     );
+    this.telemetry = createSessionTelemetryInfo(this, propagationCtx);
   }
   get loggingMetadata() {
     return {
       clientId: this.from,
       connectedTo: this.to,
       sessionId: this.id,
-      connId: this.connection?.debugId
+      connId: this.connection?.id
     };
   }
   /**
@@ -276,6 +313,7 @@ var Session = class {
           `closing connection to ${this.to} due to inactivity (missed ${misses} heartbeats which is ${missDuration}ms)`,
           this.loggingMetadata
         );
+        this.telemetry.span.addEvent("closing connection due to inactivity");
         this.closeStaleConnection();
       }
       return;
@@ -297,21 +335,25 @@ var Session = class {
   sendBufferedMessages(conn) {
     log?.info(`resending ${this.sendBuffer.length} buffered messages`, {
       ...this.loggingMetadata,
-      connId: conn.debugId
+      connId: conn.id
     });
     for (const msg of this.sendBuffer) {
       log?.debug(`resending msg`, {
         ...this.loggingMetadata,
         fullTransportMessage: msg,
-        connId: conn.debugId
+        connId: conn.id
       });
       const ok = conn.send(this.codec.toBuffer(msg));
       if (!ok) {
         const errMsg = `failed to send buffered message to ${this.to} (sus, this is a fresh connection)`;
+        conn.telemetry?.span.setStatus({
+          code: import_api2.SpanStatusCode.ERROR,
+          message: errMsg
+        });
         log?.error(errMsg, {
           ...this.loggingMetadata,
           fullTransportMessage: msg,
-          connId: conn.debugId,
+          connId: conn.id,
           tags: ["invariant-violation"]
         });
         conn.close();
@@ -394,11 +436,6 @@ var Session = class {
   }
 };
 
-// tracing/index.ts
-var import_api = require("@opentelemetry/api");
-var tracer = import_api.trace.getTracer("river");
-var tracing_default = tracer;
-
 // util/stringify.ts
 function coerceErrorString(err) {
   if (err instanceof Error) {
@@ -531,6 +568,7 @@ var NaiveJsonCodec = {
 };
 
 // transport/transport.ts
+var import_api3 = require("@opentelemetry/api");
 var defaultTransportOptions = {
   heartbeatIntervalMs: 1e3,
   heartbeatsUntilDead: 2,
@@ -609,17 +647,22 @@ var Transport = class {
       status: "connect",
       conn
     });
+    conn.telemetry = createConnectionTelemetryInfo(
+      conn,
+      session.telemetry.span
+    );
     if (isReconnect) {
       session.replaceWithNewConnection(conn);
       log?.info(`reconnected to ${connectedTo}`, session.loggingMetadata);
     }
   }
-  createSession(to, conn) {
+  createSession(to, conn, propagationCtx) {
     const session = new Session(
       conn,
       this.clientId,
       to,
-      this.options
+      this.options,
+      propagationCtx
     );
     this.sessions.set(session.to, session);
     this.eventDispatcher.dispatchEvent("sessionStatus", {
@@ -628,11 +671,11 @@ var Transport = class {
     });
     return session;
   }
-  getOrCreateSession(to, conn, sessionId) {
+  getOrCreateSession(to, conn, sessionId, propagationCtx) {
     let session = this.sessions.get(to);
     let isReconnect = session !== void 0;
     if (session?.advertisedSessionId !== void 0 && sessionId !== void 0 && session.advertisedSessionId !== sessionId) {
-      log?.warn(
+      log?.info(
         `session for ${to} already exists but has a different session id (expected: ${session.advertisedSessionId}, got: ${sessionId}), creating a new one`,
         session.loggingMetadata
       );
@@ -641,7 +684,7 @@ var Transport = class {
       session = void 0;
     }
     if (!session) {
-      session = this.createSession(to, conn);
+      session = this.createSession(to, conn, propagationCtx);
       log?.info(
         `no session for ${to}, created a new one`,
         session.loggingMetadata
@@ -654,6 +697,7 @@ var Transport = class {
   }
   deleteSession(session) {
     session.close();
+    session.telemetry.span.end();
     this.sessions.delete(session.to);
     log?.info(
       `session ${session.id} disconnect from ${session.to}`,
@@ -670,12 +714,16 @@ var Transport = class {
    * @param connectedTo The peer we are connected to.
    */
   onDisconnect(conn, session) {
+    conn.telemetry?.span.end();
     this.eventDispatcher.dispatchEvent("connectionStatus", {
       status: "disconnect",
       conn
     });
     session.connection = void 0;
-    session.beginGrace(() => this.deleteSession(session));
+    session.beginGrace(() => {
+      session.telemetry.span.addEvent("session grace period expired");
+      this.deleteSession(session);
+    });
   }
   /**
    * Parses a message from a Uint8Array into a {@link OpaqueTransportMessage}.
@@ -735,6 +783,10 @@ var Transport = class {
           tags: ["invariant-violation"]
         });
         this.protocolError(ProtocolError.MessageOrderingViolated, errMsg);
+        session.telemetry.span.setStatus({
+          code: import_api3.SpanStatusCode.ERROR,
+          message: "message order violated"
+        });
         session.close();
       }
       return;
@@ -845,6 +897,10 @@ var ClientTransport = class extends Transport {
    * tests or a special case where you don't want to reconnect.
    */
   reconnectOnConnectionDrop = true;
+  /**
+   * Optional handshake options for this client.
+   */
+  handshakeExtensions;
   constructor(clientId, providedOptions) {
     super(clientId, providedOptions);
     this.options = {
@@ -854,6 +910,9 @@ var ClientTransport = class extends Transport {
     this.inflightConnectionPromises = /* @__PURE__ */ new Map();
     this.retryBudget = new LeakyBucketRateLimit(this.options);
   }
+  extendHandshake(options) {
+    this.handshakeExtensions = options;
+  }
   handleConnection(conn, to) {
     if (this.state !== "open")
       return;
@@ -862,7 +921,7 @@ var ClientTransport = class extends Transport {
       if (!session) {
         log?.warn(
           `connection to ${to} timed out waiting for handshake, closing`,
-          { clientId: this.clientId, connectedTo: to, connId: conn.debugId }
+          { clientId: this.clientId, connectedTo: to, connId: conn.id }
         );
         conn.close();
       }
@@ -880,6 +939,10 @@ var ClientTransport = class extends Transport {
       conn.addDataListener((data2) => {
         const parsed = this.parseMsg(data2);
         if (!parsed) {
+          conn.telemetry?.span.setStatus({
+            code: import_api3.SpanStatusCode.ERROR,
+            message: "message parse failure"
+          });
           conn.close();
           return;
         }
@@ -902,6 +965,10 @@ var ClientTransport = class extends Transport {
       }
     });
     conn.addErrorListener((err) => {
+      conn.telemetry?.span.setStatus({
+        code: import_api3.SpanStatusCode.ERROR,
+        message: "connection error"
+      });
       log?.warn(`error in connection to ${to}: ${coerceErrorString(err)}`, {
         ...session?.loggingMetadata,
         clientId: this.clientId,
@@ -912,6 +979,10 @@ var ClientTransport = class extends Transport {
   receiveHandshakeResponseMessage(data, conn) {
     const parsed = this.parseMsg(data);
     if (!parsed) {
+      conn.telemetry?.span.setStatus({
+        code: import_api3.SpanStatusCode.ERROR,
+        message: "non-transport message"
+      });
       this.protocolError(
         ProtocolError.HandshakeFailed,
         "received non-transport message"
@@ -919,6 +990,10 @@ var ClientTransport = class extends Transport {
       return false;
     }
     if (!import_value.Value.Check(ControlMessageHandshakeResponseSchema, parsed.payload)) {
+      conn.telemetry?.span.setStatus({
+        code: import_api3.SpanStatusCode.ERROR,
+        message: "invalid handshake response"
+      });
       log?.warn(`received invalid handshake resp`, {
         clientId: this.clientId,
         connectedTo: parsed.from,
@@ -931,7 +1006,11 @@ var ClientTransport = class extends Transport {
       return false;
     }
     if (!parsed.payload.status.ok) {
-      log?.warn(`received invalid handshake resp`, {
+      conn.telemetry?.span.setStatus({
+        code: import_api3.SpanStatusCode.ERROR,
+        message: "handshake rejected"
+      });
+      log?.warn(`received handshake rejection`, {
         clientId: this.clientId,
         connectedTo: parsed.from,
         fullTransportMessage: parsed
@@ -961,142 +1040,94 @@ var ClientTransport = class extends Transport {
    * @param to The client ID of the node to connect to.
    */
   async connect(to) {
-    return tracing_default.startActiveSpan(
-      "connect",
-      {
-        attributes: {
-          component: "river",
-          "span.kind": "client"
-        },
-        kind: import_api2.SpanKind.CLIENT
-      },
-      async (span) => {
-        try {
-          await this.connectAttempt(to);
-        } catch (e) {
-          if (e instanceof Error) {
-            span.recordException(e);
-          } else {
-            span.recordException(coerceErrorString(e));
-          }
-          span.setStatus({ code: import_api2.SpanStatusCode.ERROR });
-        } finally {
-          span.end();
-        }
+    const canProceedWithConnection = () => this.state === "open";
+    if (!canProceedWithConnection()) {
+      log?.info(
+        `transport state is no longer open, cancelling attempt to connect to ${to}`,
+        { clientId: this.clientId, connectedTo: to }
+      );
+      return;
+    }
+    let reconnectPromise = this.inflightConnectionPromises.get(to);
+    if (!reconnectPromise) {
+      const budgetConsumed = this.retryBudget.getBudgetConsumed(to);
+      if (!this.retryBudget.hasBudget(to)) {
+        const errMsg = `tried to connect to ${to} but retry budget exceeded (more than ${budgetConsumed} attempts in the last ${this.retryBudget.totalBudgetRestoreTime}ms)`;
+        log?.warn(errMsg, { clientId: this.clientId, connectedTo: to });
+        this.protocolError(ProtocolError.RetriesExceeded, errMsg);
+        return;
       }
-    );
-  }
-  async connectAttempt(to, attempt = 0) {
-    const retry = await tracing_default.startActiveSpan(
-      "connect",
-      {
-        attributes: {
-          component: "river",
-          "river.attempt": attempt,
-          "span.kind": "client"
-        },
-        kind: import_api2.SpanKind.CLIENT
-      },
-      async (span) => {
+      let sleep = Promise.resolve();
+      const backoffMs = this.retryBudget.getBackoffMs(to);
+      if (backoffMs > 0) {
+        sleep = new Promise((resolve) => setTimeout(resolve, backoffMs));
+      }
+      log?.info(`attempting connection to ${to} (${backoffMs}ms backoff)`, {
+        clientId: this.clientId,
+        connectedTo: to
+      });
+      this.retryBudget.consumeBudget(to);
+      reconnectPromise = tracing_default.startActiveSpan("connect", async (span) => {
         try {
-          const canProceedWithConnection = () => this.state === "open";
+          span.addEvent("backoff", { backoffMs });
+          await sleep;
           if (!canProceedWithConnection()) {
-            log?.info(
-              `transport state is no longer open, cancelling attempt to connect to ${to}`,
-              { clientId: this.clientId, connectedTo: to }
-            );
-            return false;
+            throw new Error("transport state is no longer open");
           }
-          let reconnectPromise = this.inflightConnectionPromises.get(to);
-          if (!reconnectPromise) {
-            const budgetConsumed = this.retryBudget.getBudgetConsumed(to);
-            if (!this.retryBudget.hasBudget(to)) {
-              const errMsg = `tried to connect to ${to} but retry budget exceeded (more than ${budgetConsumed} attempts in the last ${this.retryBudget.totalBudgetRestoreTime}ms)`;
-              log?.warn(errMsg, { clientId: this.clientId, connectedTo: to });
-              this.protocolError(ProtocolError.RetriesExceeded, errMsg);
-              return false;
-            }
-            let sleep = Promise.resolve();
-            const backoffMs = this.retryBudget.getBackoffMs(to);
-            if (backoffMs > 0) {
-              sleep = new Promise((resolve) => setTimeout(resolve, backoffMs));
-            }
-            log?.info(
-              `attempting connection to ${to} (${backoffMs}ms backoff)`,
-              {
-                clientId: this.clientId,
-                connectedTo: to
-              }
-            );
-            this.retryBudget.consumeBudget(to);
-            reconnectPromise = sleep.then(() => {
-              if (!canProceedWithConnection()) {
-                throw new Error("transport state is no longer open");
-              }
-            }).then(() => this.createNewOutgoingConnection(to)).then((conn) => {
-              if (!canProceedWithConnection()) {
-                log?.info(
-                  `transport state is no longer open, closing pre-handshake connection to ${to}`,
-                  {
-                    clientId: this.clientId,
-                    connectedTo: to,
-                    connId: conn.debugId
-                  }
-                );
-                conn.close();
-                throw new Error("transport state is no longer open");
-              }
-              return this.sendHandshake(to, conn).then((ok) => {
-                if (!ok) {
-                  conn.close();
-                  throw new Error("failed to send handshake");
-                }
-                return conn;
-              });
-            });
-            this.inflightConnectionPromises.set(to, reconnectPromise);
-          } else {
+          span.addEvent("connecting");
+          const conn = await this.createNewOutgoingConnection(to);
+          if (!canProceedWithConnection()) {
             log?.info(
-              `attempting connection to ${to} (reusing previous attempt)`,
+              `transport state is no longer open, closing pre-handshake connection to ${to}`,
               {
                 clientId: this.clientId,
-                connectedTo: to
+                connectedTo: to,
+                connId: conn.id
               }
             );
+            conn.close();
+            throw new Error("transport state is no longer open");
           }
-          try {
-            await reconnectPromise;
-          } catch (error) {
-            this.inflightConnectionPromises.delete(to);
-            const errStr = coerceErrorString(error);
-            if (!this.reconnectOnConnectionDrop || !canProceedWithConnection()) {
-              log?.warn(`connection to ${to} failed (${errStr})`, {
-                clientId: this.clientId,
-                connectedTo: to
-              });
-            } else {
-              log?.warn(`connection to ${to} failed (${errStr}), retrying`, {
-                clientId: this.clientId,
-                connectedTo: to
-              });
-              return true;
-            }
+          span.addEvent("sending handshake");
+          const ok = await this.sendHandshake(to, conn);
+          if (!ok) {
+            conn.close();
+            throw new Error("failed to send handshake");
           }
-        } catch (e) {
-          if (e instanceof Error) {
-            span.recordException(e);
-          } else {
-            span.recordException(coerceErrorString(e));
-          }
-          span.setStatus({ code: import_api2.SpanStatusCode.ERROR });
+          return conn;
+        } catch (err) {
+          const errStr = coerceErrorString(err);
+          span.recordException(errStr);
+          span.setStatus({ code: import_api3.SpanStatusCode.ERROR });
+          throw err;
         } finally {
           span.end();
         }
-        return false;
+      });
+      this.inflightConnectionPromises.set(to, reconnectPromise);
+    } else {
+      log?.info(`attempting connection to ${to} (reusing previous attempt)`, {
+        clientId: this.clientId,
+        connectedTo: to
+      });
+    }
+    try {
+      await reconnectPromise;
+    } catch (error) {
+      this.inflightConnectionPromises.delete(to);
+      const errStr = coerceErrorString(error);
+      if (!this.reconnectOnConnectionDrop || !canProceedWithConnection()) {
+        log?.warn(`connection to ${to} failed (${errStr})`, {
+          clientId: this.clientId,
+          connectedTo: to
+        });
+      } else {
+        log?.warn(`connection to ${to} failed (${errStr}), retrying`, {
+          clientId: this.clientId,
+          connectedTo: to
+        });
+        return this.connect(to);
       }
-    );
-    if (retry) {
-      return this.connectAttempt(to, attempt + 1);
     }
   }
   deleteSession(session) {
@@ -1104,13 +1135,11 @@ var ClientTransport = class extends Transport {
     super.deleteSession(session);
   }
   async sendHandshake(to, conn) {
-    const tracing = { traceparent: "", tracestate: "" };
-    import_api2.propagation.inject(import_api2.context.active(), tracing);
     let metadata;
-    if (this.options.handshake) {
-      metadata = await this.options.handshake.get();
-      if (!import_value.Value.Check(this.options.handshake.schema, metadata)) {
-        log?.error(`handshake metadata did not match schema`, {
+    if (this.handshakeExtensions) {
+      metadata = await this.handshakeExtensions.construct();
+      if (!import_value.Value.Check(this.handshakeExtensions.schema, metadata)) {
+        log?.error(`constructed handshake metadata did not match schema`, {
           clientId: this.clientId,
           connectedTo: to,
           tags: ["invariant-violation"]
@@ -1119,6 +1148,10 @@ var ClientTransport = class extends Transport {
           ProtocolError.HandshakeFailed,
           "handshake metadata did not match schema"
         );
+        conn.telemetry?.span.setStatus({
+          code: import_api3.SpanStatusCode.ERROR,
+          message: "handshake meta mismatch"
+        });
         return false;
       }
     }
@@ -1128,7 +1161,7 @@ var ClientTransport = class extends Transport {
       to,
       session.id,
       metadata,
-      tracing
+      getPropagationContext(session.telemetry.ctx)
     );
     log?.debug(`sending handshake request to ${to}`, {
       clientId: this.clientId,
@@ -1144,13 +1177,31 @@ var ClientTransport = class extends Transport {
 };
 
 // transport/impls/ws/connection.ts
-var import_agnostic_ws = __toESM(require("agnostic-ws"), 1);
 var WebSocketConnection = class extends Connection {
+  errorCb = null;
+  closeCb = null;
   ws;
   constructor(ws) {
     super();
     this.ws = ws;
     this.ws.binaryType = "arraybuffer";
+    let didError = false;
+    this.ws.onerror = () => {
+      didError = true;
+    };
+    this.ws.onclose = ({ code, reason }) => {
+      if (didError && this.errorCb) {
+        this.errorCb(
+          new Error(
+            `websocket closed with code and reason: ${code} - ${reason}`
+          )
+        );
+        return;
+      }
+      if (this.closeCb) {
+        this.closeCb();
+      }
+    };
   }
   addDataListener(cb) {
     this.ws.onmessage = (msg) => cb(msg.data);
@@ -1159,13 +1210,13 @@ var WebSocketConnection = class extends Connection {
     this.ws.onmessage = null;
   }
   addCloseListener(cb) {
-    this.ws.onclose = cb;
+    this.closeCb = cb;
   }
   addErrorListener(cb) {
-    this.ws.onerror = (err) => cb(err.error);
+    this.errorCb = cb;
   }
   send(payload) {
-    if (this.ws.readyState === import_agnostic_ws.default.OPEN) {
+    if (this.ws.readyState === this.ws.OPEN) {
       this.ws.send(payload);
       return true;
     } else {
@@ -1182,7 +1233,7 @@ var WebSocketClientTransport = class extends ClientTransport {
   /**
    * A function that returns a Promise that resolves to a websocket URL.
    */
-  urlGetter;
+  wsGetter;
   /**
    * Creates a new WebSocketClientTransport instance.
    * @param wsGetter A function that returns a Promise that resolves to a WebSocket instance.
@@ -1190,53 +1241,39 @@ var WebSocketClientTransport = class extends ClientTransport {
    * @param serverId The ID of the server this transport is connecting to.
    * @param providedOptions An optional object containing configuration options for the transport.
    */
-  constructor(urlGetter, clientId, providedOptions) {
+  constructor(wsGetter, clientId, providedOptions) {
     super(clientId, providedOptions);
-    this.urlGetter = urlGetter;
+    this.wsGetter = wsGetter;
   }
   async createNewOutgoingConnection(to) {
-    const wsRes = await new Promise((resolve) => {
-      log?.info(`establishing a new websocket to ${to}`, {
-        clientId: this.clientId,
-        connectedTo: to
-      });
-      Promise.resolve(this.urlGetter(to)).then((url) => new import_agnostic_ws2.default(url)).then((ws) => {
-        if (ws.readyState === import_agnostic_ws2.default.OPEN) {
-          resolve({ ws });
-          return;
-        }
-        if (ws.readyState === import_agnostic_ws2.default.CLOSING || ws.readyState === import_agnostic_ws2.default.CLOSED) {
-          resolve({ err: "ws is closing or closed" });
-          return;
-        }
-        ws.onopen = () => {
-          resolve({ ws });
-        };
-        ws.onclose = (evt) => {
-          resolve({ err: evt.reason });
-        };
-        ws.onerror = (evt) => {
-          const err = evt.error;
-          resolve({
-            err: `${err.name}: ${err.message}`
-          });
-        };
-      }).catch((e) => {
-        const reason = e instanceof Error ? e.message : "unknown reason";
-        resolve({ err: `couldn't get a new websocket: ${reason}` });
-      });
+    log?.info(`establishing a new websocket to ${to}`, {
+      clientId: this.clientId,
+      connectedTo: to
     });
-    if ("ws" in wsRes) {
-      const conn = new WebSocketConnection(wsRes.ws);
-      log?.info(`raw websocket to ${to} ok, starting handshake`, {
-        clientId: this.clientId,
-        connectedTo: to
-      });
-      this.handleConnection(conn, to);
-      return conn;
-    } else {
-      throw new Error(wsRes.err);
-    }
+    const ws = await this.wsGetter(to);
+    await new Promise((resolve, reject) => {
+      if (ws.readyState === ws.OPEN) {
+        resolve();
+        return;
+      }
+      if (ws.readyState === ws.CLOSING || ws.readyState === ws.CLOSED) {
+        reject(new Error("ws is closing or closed"));
+        return;
+      }
+      ws.onopen = () => {
+        resolve();
+      };
+      ws.onclose = (evt) => {
+        reject(new Error(evt.reason));
+      };
+    });
+    const conn = new WebSocketConnection(ws);
+    log?.info(`raw websocket to ${to} ok, starting handshake`, {
+      clientId: this.clientId,
+      connectedTo: to
+    });
+    this.handleConnection(conn, to);
+    return conn;
   }
 };
 // Annotate the CommonJS export names for ESM import in node: