@replit/river 0.25.1 → 0.25.2

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

@@ -212,7 +212,7 @@ var defaultTransportOptions = {
   codec: NaiveJsonCodec
 };
 var defaultConnectionRetryOptions = {
-  baseIntervalMs: 250,
+  baseIntervalMs: 150,
   maxJitterMs: 200,
   maxBackoffMs: 32e3,
   attemptBudgetCapacity: 5,
@@ -229,17 +229,17 @@ var defaultServerTransportOptions = {
 // transport/rateLimit.ts
 var LeakyBucketRateLimit = class {
   budgetConsumed;
-  intervalHandles;
+  intervalHandle;
   options;
   constructor(options) {
     this.options = options;
-    this.budgetConsumed = /* @__PURE__ */ new Map();
-    this.intervalHandles = /* @__PURE__ */ new Map();
+    this.budgetConsumed = 0;
   }
-  getBackoffMs(user) {
-    if (!this.budgetConsumed.has(user))
+  getBackoffMs() {
+    if (this.getBudgetConsumed() === 0) {
       return 0;
-    const exponent = Math.max(0, this.getBudgetConsumed(user) - 1);
+    }
+    const exponent = Math.max(0, this.getBudgetConsumed() - 1);
     const jitter = Math.floor(Math.random() * this.options.maxJitterMs);
     const backoffMs = Math.min(
       this.options.baseIntervalMs * 2 ** exponent,
@@ -250,50 +250,46 @@ var LeakyBucketRateLimit = class {
   get totalBudgetRestoreTime() {
     return this.options.budgetRestoreIntervalMs * this.options.attemptBudgetCapacity;
   }
-  consumeBudget(user) {
-    this.stopLeak(user);
-    this.budgetConsumed.set(user, this.getBudgetConsumed(user) + 1);
+  consumeBudget() {
+    this.stopLeak();
+    this.budgetConsumed = this.getBudgetConsumed() + 1;
   }
-  getBudgetConsumed(user) {
-    return this.budgetConsumed.get(user) ?? 0;
+  getBudgetConsumed() {
+    return this.budgetConsumed;
   }
-  hasBudget(user) {
-    return this.getBudgetConsumed(user) < this.options.attemptBudgetCapacity;
+  hasBudget() {
+    return this.getBudgetConsumed() < this.options.attemptBudgetCapacity;
   }
-  startRestoringBudget(user) {
-    if (this.intervalHandles.has(user)) {
+  startRestoringBudget() {
+    if (this.intervalHandle) {
       return;
     }
     const restoreBudgetForUser = () => {
-      const currentBudget = this.budgetConsumed.get(user);
+      const currentBudget = this.budgetConsumed;
       if (!currentBudget) {
-        this.stopLeak(user);
+        this.stopLeak();
         return;
       }
       const newBudget = currentBudget - 1;
       if (newBudget === 0) {
-        this.budgetConsumed.delete(user);
         return;
       }
-      this.budgetConsumed.set(user, newBudget);
+      this.budgetConsumed = newBudget;
     };
-    const intervalHandle = setInterval(
+    this.intervalHandle = setInterval(
       restoreBudgetForUser,
       this.options.budgetRestoreIntervalMs
     );
-    this.intervalHandles.set(user, intervalHandle);
   }
-  stopLeak(user) {
-    if (!this.intervalHandles.has(user)) {
+  stopLeak() {
+    if (!this.intervalHandle) {
       return;
     }
-    clearInterval(this.intervalHandles.get(user));
-    this.intervalHandles.delete(user);
+    clearInterval(this.intervalHandle);
+    this.intervalHandle = void 0;
   }
   close() {
-    for (const user of this.intervalHandles.keys()) {
-      this.stopLeak(user);
-    }
+    this.stopLeak();
   }
 };
 
@@ -440,7 +436,7 @@ var CommonSession = class extends StateMachineState {
   from;
   options;
   log;
-  constructor(from, options, log) {
+  constructor({ from, options, log }) {
     super();
     this.from = from;
     this.options = options;
@@ -481,8 +477,9 @@ var IdentifiedSession = class extends CommonSession {
    */
   ack;
   sendBuffer;
-  constructor(id, from, to, seq, ack, sendBuffer, telemetry, options, log) {
-    super(from, options, log);
+  constructor(props) {
+    const { id, to, seq, ack, sendBuffer, telemetry, log } = props;
+    super(props);
     this.id = id;
     this.to = to;
     this.seq = seq;
@@ -537,23 +534,23 @@ var SessionConnecting = class extends IdentifiedSession {
   connPromise;
   listeners;
   connectionTimeout;
-  constructor(connPromise, listeners, ...args) {
-    super(...args);
-    this.connPromise = connPromise;
-    this.listeners = listeners;
+  constructor(props) {
+    super(props);
+    this.connPromise = props.connPromise;
+    this.listeners = props.listeners;
     this.connectionTimeout = setTimeout(() => {
-      listeners.onConnectionTimeout();
+      this.listeners.onConnectionTimeout();
     }, this.options.connectionTimeoutMs);
-    connPromise.then(
+    this.connPromise.then(
       (conn) => {
         if (this._isConsumed)
           return;
-        listeners.onConnectionEstablished(conn);
+        this.listeners.onConnectionEstablished(conn);
       },
       (err) => {
         if (this._isConsumed)
           return;
-        listeners.onConnectionFailed(err);
+        this.listeners.onConnectionFailed(err);
       }
     );
   }
@@ -579,9 +576,9 @@ var SessionNoConnection = class extends IdentifiedSession {
   state = "NoConnection" /* NoConnection */;
   listeners;
   gracePeriodTimeout;
-  constructor(listeners, ...args) {
-    super(...args);
-    this.listeners = listeners;
+  constructor(props) {
+    super(props);
+    this.listeners = props.listeners;
     this.gracePeriodTimeout = setTimeout(() => {
       this.listeners.onSessionGracePeriodElapsed();
     }, this.options.sessionDisconnectGraceMs);
@@ -602,7 +599,7 @@ var SessionNoConnection = class extends IdentifiedSession {
 var import_api = require("@opentelemetry/api");
 
 // package.json
-var version = "0.25.1";
+var version = "0.25.2";
 
 // tracing/index.ts
 function getPropagationContext(ctx) {
@@ -639,16 +636,16 @@ var SessionWaitingForHandshake = class extends CommonSession {
   conn;
   listeners;
   handshakeTimeout;
-  constructor(conn, listeners, ...args) {
-    super(...args);
-    this.conn = conn;
-    this.listeners = listeners;
+  constructor(props) {
+    super(props);
+    this.conn = props.conn;
+    this.listeners = props.listeners;
     this.handshakeTimeout = setTimeout(() => {
-      listeners.onHandshakeTimeout();
+      this.listeners.onHandshakeTimeout();
     }, this.options.handshakeTimeoutMs);
     this.conn.addDataListener(this.onHandshakeData);
-    this.conn.addErrorListener(listeners.onConnectionErrored);
-    this.conn.addCloseListener(listeners.onConnectionClosed);
+    this.conn.addErrorListener(this.listeners.onConnectionErrored);
+    this.conn.addCloseListener(this.listeners.onConnectionClosed);
   }
   onHandshakeData = (msg) => {
     const parsedMsg = this.parseMsg(msg);
@@ -685,16 +682,16 @@ var SessionHandshaking = class extends IdentifiedSession {
   conn;
   listeners;
   handshakeTimeout;
-  constructor(conn, listeners, ...args) {
-    super(...args);
-    this.conn = conn;
-    this.listeners = listeners;
+  constructor(props) {
+    super(props);
+    this.conn = props.conn;
+    this.listeners = props.listeners;
     this.handshakeTimeout = setTimeout(() => {
-      listeners.onHandshakeTimeout();
+      this.listeners.onHandshakeTimeout();
     }, this.options.handshakeTimeoutMs);
     this.conn.addDataListener(this.onHandshakeData);
-    this.conn.addErrorListener(listeners.onConnectionErrored);
-    this.conn.addCloseListener(listeners.onConnectionClosed);
+    this.conn.addErrorListener(this.listeners.onConnectionErrored);
+    this.conn.addCloseListener(this.listeners.onConnectionClosed);
   }
   onHandshakeData = (msg) => {
     const parsedMsg = this.parseMsg(msg);
@@ -726,16 +723,13 @@ var SessionConnected = class extends IdentifiedSession {
   state = "Connected" /* Connected */;
   conn;
   listeners;
-  activeHeartbeatHandle;
-  activeHeartbeatMisses = 0;
-  passiveHeartbeatHandle;
-  get isActivelyHeartbeating() {
-    return this.activeHeartbeatHandle !== void 0;
-  }
+  heartbeatHandle;
+  heartbeatMisses = 0;
+  isActivelyHeartbeating;
   updateBookkeeping(ack, seq) {
     this.sendBuffer = this.sendBuffer.filter((unacked) => unacked.seq >= ack);
     this.ack = seq + 1;
-    this.activeHeartbeatMisses = 0;
+    this.heartbeatMisses = 0;
   }
   send(msg) {
     const constructedMsg = this.constructMsg(msg);
@@ -743,13 +737,13 @@ var SessionConnected = class extends IdentifiedSession {
     this.conn.send(this.options.codec.toBuffer(constructedMsg));
     return constructedMsg.id;
   }
-  constructor(conn, listeners, ...args) {
-    super(...args);
-    this.conn = conn;
-    this.listeners = listeners;
+  constructor(props) {
+    super(props);
+    this.conn = props.conn;
+    this.listeners = props.listeners;
     this.conn.addDataListener(this.onMessageData);
-    this.conn.addCloseListener(listeners.onConnectionClosed);
-    this.conn.addErrorListener(listeners.onConnectionErrored);
+    this.conn.addCloseListener(this.listeners.onConnectionClosed);
+    this.conn.addErrorListener(this.listeners.onConnectionErrored);
     if (this.sendBuffer.length > 0) {
       this.log?.debug(
         `sending ${this.sendBuffer.length} buffered messages`,
@@ -757,12 +751,11 @@ var SessionConnected = class extends IdentifiedSession {
       );
     }
     for (const msg of this.sendBuffer) {
-      conn.send(this.options.codec.toBuffer(msg));
+      this.conn.send(this.options.codec.toBuffer(msg));
     }
-  }
-  startActiveHeartbeat() {
-    this.activeHeartbeatHandle = setInterval(() => {
-      const misses = this.activeHeartbeatMisses;
+    this.isActivelyHeartbeating = false;
+    this.heartbeatHandle = setInterval(() => {
+      const misses = this.heartbeatMisses;
       const missDuration = misses * this.options.heartbeatIntervalMs;
       if (misses >= this.options.heartbeatsUntilDead) {
         this.log?.info(
@@ -771,29 +764,18 @@ var SessionConnected = class extends IdentifiedSession {
         );
         this.telemetry.span.addEvent("closing connection due to inactivity");
         this.conn.close();
-        clearInterval(this.activeHeartbeatHandle);
-        this.activeHeartbeatHandle = void 0;
+        clearInterval(this.heartbeatHandle);
+        this.heartbeatHandle = void 0;
         return;
       }
-      this.sendHeartbeat();
-      this.activeHeartbeatMisses++;
+      if (this.isActivelyHeartbeating) {
+        this.sendHeartbeat();
+      }
+      this.heartbeatMisses++;
     }, this.options.heartbeatIntervalMs);
   }
-  waitForNextHeartbeat() {
-    const duration = this.options.heartbeatsUntilDead * this.options.heartbeatIntervalMs;
-    if (this.passiveHeartbeatHandle) {
-      clearTimeout(this.passiveHeartbeatHandle);
-      this.passiveHeartbeatHandle = void 0;
-    }
-    this.passiveHeartbeatHandle = setTimeout(() => {
-      this.log?.info(
-        `closing connection to ${this.to} due to not receiving a heartbeat in the last ${duration}ms`,
-        this.loggingMetadata
-      );
-      this.telemetry.span.addEvent("closing connection due to inactivity");
-      this.conn.close();
-      this.passiveHeartbeatHandle = void 0;
-    }, duration);
+  startActiveHeartbeat() {
+    this.isActivelyHeartbeating = true;
   }
   sendHeartbeat() {
     this.log?.debug("sending heartbeat", this.loggingMetadata);
@@ -848,7 +830,6 @@ var SessionConnected = class extends IdentifiedSession {
     });
     if (!this.isActivelyHeartbeating) {
       this.sendHeartbeat();
-      this.waitForNextHeartbeat();
     }
   };
   _handleStateExit() {
@@ -856,10 +837,8 @@ var SessionConnected = class extends IdentifiedSession {
     this.conn.removeDataListener(this.onMessageData);
     this.conn.removeCloseListener(this.listeners.onConnectionClosed);
     this.conn.removeErrorListener(this.listeners.onConnectionErrored);
-    clearInterval(this.activeHeartbeatHandle);
-    clearTimeout(this.passiveHeartbeatHandle);
-    this.activeHeartbeatHandle = void 0;
-    this.passiveHeartbeatHandle = void 0;
+    clearInterval(this.heartbeatHandle);
+    this.heartbeatHandle = void 0;
   }
   _handleClose() {
     super._handleClose();
@@ -867,52 +846,76 @@ var SessionConnected = class extends IdentifiedSession {
   }
 };
 
+// transport/sessionStateMachine/SessionBackingOff.ts
+var SessionBackingOff = class extends IdentifiedSession {
+  state = "BackingOff" /* BackingOff */;
+  listeners;
+  backoffTimeout;
+  constructor(props) {
+    super(props);
+    this.listeners = props.listeners;
+    this.backoffTimeout = setTimeout(() => {
+      this.listeners.onBackoffFinished();
+    }, props.backoffMs);
+  }
+  _handleClose() {
+    super._handleClose();
+  }
+  _handleStateExit() {
+    super._handleStateExit();
+    if (this.backoffTimeout) {
+      clearTimeout(this.backoffTimeout);
+      this.backoffTimeout = void 0;
+    }
+  }
+};
+
 // transport/sessionStateMachine/transitions.ts
 function inheritSharedSession(session) {
-  return [
-    session.id,
-    session.from,
-    session.to,
-    session.seq,
-    session.ack,
-    session.sendBuffer,
-    session.telemetry,
-    session.options,
-    session.log
-  ];
+  return {
+    id: session.id,
+    from: session.from,
+    to: session.to,
+    seq: session.seq,
+    ack: session.ack,
+    sendBuffer: session.sendBuffer,
+    telemetry: session.telemetry,
+    options: session.options,
+    log: session.log
+  };
 }
 var SessionStateGraph = {
   entrypoints: {
-    NoConnection(to, from, listeners, options, log) {
+    NoConnection: (to, from, listeners, options, log) => {
       const id = `session-${generateId()}`;
       const telemetry = createSessionTelemetryInfo(id, to, from);
       const sendBuffer = [];
-      const session = new SessionNoConnection(
+      const session = new SessionNoConnection({
         listeners,
         id,
         from,
         to,
-        0,
-        0,
+        seq: 0,
+        ack: 0,
         sendBuffer,
         telemetry,
         options,
         log
-      );
+      });
       session.log?.info(`session ${session.id} created in NoConnection state`, {
         ...session.loggingMetadata,
         tags: ["state-transition"]
       });
       return session;
     },
-    WaitingForHandshake(from, conn, listeners, options, log) {
-      const session = new SessionWaitingForHandshake(
+    WaitingForHandshake: (from, conn, listeners, options, log) => {
+      const session = new SessionWaitingForHandshake({
         conn,
         listeners,
         from,
         options,
         log
-      );
+      });
       session.log?.info(`session created in WaitingForHandshake state`, {
         ...session.loggingMetadata,
         tags: ["state-transition"]
@@ -924,16 +927,33 @@ var SessionStateGraph = {
   // After a session is transitioned, any usage of the old session will throw.
   transition: {
     // happy path transitions
-    NoConnectionToConnecting(oldSession, connPromise, listeners) {
+    NoConnectionToBackingOff: (oldSession, backoffMs, listeners) => {
       const carriedState = inheritSharedSession(oldSession);
       oldSession._handleStateExit();
-      const session = new SessionConnecting(
-        connPromise,
+      const session = new SessionBackingOff({
+        backoffMs,
         listeners,
         ...carriedState
+      });
+      session.log?.info(
+        `session ${session.id} transition from NoConnection to BackingOff`,
+        {
+          ...session.loggingMetadata,
+          tags: ["state-transition"]
+        }
       );
+      return session;
+    },
+    BackingOffToConnecting: (oldSession, connPromise, listeners) => {
+      const carriedState = inheritSharedSession(oldSession);
+      oldSession._handleStateExit();
+      const session = new SessionConnecting({
+        connPromise,
+        listeners,
+        ...carriedState
+      });
       session.log?.info(
-        `session ${session.id} transition from NoConnection to Connecting`,
+        `session ${session.id} transition from BackingOff to Connecting`,
         {
           ...session.loggingMetadata,
           tags: ["state-transition"]
@@ -941,10 +961,14 @@ var SessionStateGraph = {
       );
       return session;
     },
-    ConnectingToHandshaking(oldSession, conn, listeners) {
+    ConnectingToHandshaking: (oldSession, conn, listeners) => {
       const carriedState = inheritSharedSession(oldSession);
       oldSession._handleStateExit();
-      const session = new SessionHandshaking(conn, listeners, ...carriedState);
+      const session = new SessionHandshaking({
+        conn,
+        listeners,
+        ...carriedState
+      });
       session.log?.info(
         `session ${session.id} transition from Connecting to Handshaking`,
         {
@@ -954,11 +978,15 @@ var SessionStateGraph = {
       );
       return session;
     },
-    HandshakingToConnected(oldSession, listeners) {
+    HandshakingToConnected: (oldSession, listeners) => {
       const carriedState = inheritSharedSession(oldSession);
       const conn = oldSession.conn;
       oldSession._handleStateExit();
-      const session = new SessionConnected(conn, listeners, ...carriedState);
+      const session = new SessionConnected({
+        conn,
+        listeners,
+        ...carriedState
+      });
       session.log?.info(
         `session ${session.id} transition from Handshaking to Connected`,
         {
@@ -968,7 +996,7 @@ var SessionStateGraph = {
       );
       return session;
     },
-    WaitingForHandshakeToConnected(pendingSession, oldSession, sessionId, to, propagationCtx, listeners) {
+    WaitingForHandshakeToConnected: (pendingSession, oldSession, sessionId, to, propagationCtx, listeners) => {
       const conn = pendingSession.conn;
       const { from, options } = pendingSession;
       const carriedState = oldSession ? (
@@ -976,21 +1004,30 @@ var SessionStateGraph = {
         inheritSharedSession(oldSession)
       ) : (
         // old session does not exist, create new state
-        [
-          sessionId,
+        {
+          id: sessionId,
           from,
           to,
-          0,
-          0,
-          [],
-          createSessionTelemetryInfo(sessionId, to, from, propagationCtx),
+          seq: 0,
+          ack: 0,
+          sendBuffer: [],
+          telemetry: createSessionTelemetryInfo(
+            sessionId,
+            to,
+            from,
+            propagationCtx
+          ),
           options,
-          pendingSession.log
-        ]
+          log: pendingSession.log
+        }
       );
       pendingSession._handleStateExit();
       oldSession?._handleStateExit();
-      const session = new SessionConnected(conn, listeners, ...carriedState);
+      const session = new SessionConnected({
+        conn,
+        listeners,
+        ...carriedState
+      });
       session.log?.info(
         `session ${session.id} transition from WaitingForHandshake to Connected`,
         {
@@ -1001,11 +1038,24 @@ var SessionStateGraph = {
       return session;
     },
     // disconnect paths
-    ConnectingToNoConnection(oldSession, listeners) {
+    BackingOffToNoConnection: (oldSession, listeners) => {
+      const carriedState = inheritSharedSession(oldSession);
+      oldSession._handleStateExit();
+      const session = new SessionNoConnection({ listeners, ...carriedState });
+      session.log?.info(
+        `session ${session.id} transition from BackingOff to NoConnection`,
+        {
+          ...session.loggingMetadata,
+          tags: ["state-transition"]
+        }
+      );
+      return session;
+    },
+    ConnectingToNoConnection: (oldSession, listeners) => {
       const carriedState = inheritSharedSession(oldSession);
       oldSession.bestEffortClose();
       oldSession._handleStateExit();
-      const session = new SessionNoConnection(listeners, ...carriedState);
+      const session = new SessionNoConnection({ listeners, ...carriedState });
       session.log?.info(
         `session ${session.id} transition from Connecting to NoConnection`,
         {
@@ -1015,11 +1065,11 @@ var SessionStateGraph = {
       );
       return session;
     },
-    HandshakingToNoConnection(oldSession, listeners) {
+    HandshakingToNoConnection: (oldSession, listeners) => {
       const carriedState = inheritSharedSession(oldSession);
       oldSession.conn.close();
       oldSession._handleStateExit();
-      const session = new SessionNoConnection(listeners, ...carriedState);
+      const session = new SessionNoConnection({ listeners, ...carriedState });
       session.log?.info(
         `session ${session.id} transition from Handshaking to NoConnection`,
         {
@@ -1029,11 +1079,11 @@ var SessionStateGraph = {
       );
       return session;
     },
-    ConnectedToNoConnection(oldSession, listeners) {
+    ConnectedToNoConnection: (oldSession, listeners) => {
       const carriedState = inheritSharedSession(oldSession);
       oldSession.conn.close();
       oldSession._handleStateExit();
-      const session = new SessionNoConnection(listeners, ...carriedState);
+      const session = new SessionNoConnection({ listeners, ...carriedState });
       session.log?.info(
         `session ${session.id} transition from Connected to NoConnection`,
         {
@@ -1045,6 +1095,31 @@ var SessionStateGraph = {
     }
   }
 };
+var transitions = SessionStateGraph.transition;
+var ClientSessionStateGraph = {
+  entrypoint: SessionStateGraph.entrypoints.NoConnection,
+  transition: {
+    // happy paths
+    NoConnectionToBackingOff: transitions.NoConnectionToBackingOff,
+    BackingOffToConnecting: transitions.BackingOffToConnecting,
+    ConnectingToHandshaking: transitions.ConnectingToHandshaking,
+    HandshakingToConnected: transitions.HandshakingToConnected,
+    // disconnect paths
+    BackingOffToNoConnection: transitions.BackingOffToNoConnection,
+    ConnectingToNoConnection: transitions.ConnectingToNoConnection,
+    HandshakingToNoConnection: transitions.HandshakingToNoConnection,
+    ConnectedToNoConnection: transitions.ConnectedToNoConnection
+  }
+};
+var ServerSessionStateGraph = {
+  entrypoint: SessionStateGraph.entrypoints.WaitingForHandshake,
+  transition: {
+    // happy paths
+    WaitingForHandshakeToConnected: transitions.WaitingForHandshakeToConnected,
+    // disconnect paths
+    ConnectedToNoConnection: transitions.ConnectedToNoConnection
+  }
+};
 
 // transport/transport.ts
 var Transport = class {
@@ -1224,8 +1299,10 @@ var ClientTransport = class extends Transport {
    * Optional handshake options for this client.
    */
   handshakeExtensions;
+  sessions;
   constructor(clientId, providedOptions) {
     super(clientId, providedOptions);
+    this.sessions = /* @__PURE__ */ new Map();
     this.options = {
       ...defaultClientTransportOptions,
       ...providedOptions
@@ -1257,7 +1334,7 @@ var ClientTransport = class extends Transport {
     return session.send(msg);
   }
   createUnconnectedSession(to) {
-    const session = SessionStateGraph.entrypoints.NoConnection(
+    const session = ClientSessionStateGraph.entrypoint(
       to,
       this.clientId,
       {
@@ -1283,40 +1360,44 @@ var ClientTransport = class extends Transport {
     return noConnectionSession;
   }
   onConnectionEstablished(session, conn) {
-    const handshakingSession = SessionStateGraph.transition.ConnectingToHandshaking(session, conn, {
-      onConnectionErrored: (err) => {
-        const errStr = coerceErrorString(err);
-        this.log?.error(
-          `connection to ${handshakingSession.to} errored during handshake: ${errStr}`,
-          handshakingSession.loggingMetadata
-        );
-      },
-      onConnectionClosed: () => {
-        this.log?.warn(
-          `connection to ${handshakingSession.to} closed during handshake`,
-          handshakingSession.loggingMetadata
-        );
-        this.onConnClosed(handshakingSession);
-      },
-      onHandshake: (msg) => {
-        this.onHandshakeResponse(handshakingSession, msg);
-      },
-      onInvalidHandshake: (reason) => {
-        this.log?.error(
-          `invalid handshake: ${reason}`,
-          handshakingSession.loggingMetadata
-        );
-        this.deleteSession(session);
-        this.protocolError(ProtocolError.HandshakeFailed, reason);
-      },
-      onHandshakeTimeout: () => {
-        this.log?.error(
-          `connection to ${handshakingSession.to} timed out during handshake`,
-          handshakingSession.loggingMetadata
-        );
-        this.onConnClosed(handshakingSession);
+    const handshakingSession = ClientSessionStateGraph.transition.ConnectingToHandshaking(
+      session,
+      conn,
+      {
+        onConnectionErrored: (err) => {
+          const errStr = coerceErrorString(err);
+          this.log?.error(
+            `connection to ${handshakingSession.to} errored during handshake: ${errStr}`,
+            handshakingSession.loggingMetadata
+          );
+        },
+        onConnectionClosed: () => {
+          this.log?.warn(
+            `connection to ${handshakingSession.to} closed during handshake`,
+            handshakingSession.loggingMetadata
+          );
+          this.onConnClosed(handshakingSession);
+        },
+        onHandshake: (msg) => {
+          this.onHandshakeResponse(handshakingSession, msg);
+        },
+        onInvalidHandshake: (reason) => {
+          this.log?.error(
+            `invalid handshake: ${reason}`,
+            handshakingSession.loggingMetadata
+          );
+          this.deleteSession(session);
+          this.protocolError(ProtocolError.HandshakeFailed, reason);
+        },
+        onHandshakeTimeout: () => {
+          this.log?.error(
+            `connection to ${handshakingSession.to} timed out during handshake`,
+            handshakingSession.loggingMetadata
+          );
+          this.onConnClosed(handshakingSession);
+        }
       }
-    });
+    );
     this.updateSession(handshakingSession);
     void this.sendHandshake(handshakingSession);
     return handshakingSession;
@@ -1371,7 +1452,7 @@ var ClientTransport = class extends Transport {
       ...session.loggingMetadata,
       transportMessage: msg
     });
-    const connectedSession = SessionStateGraph.transition.HandshakingToConnected(session, {
+    const connectedSession = ClientSessionStateGraph.transition.HandshakingToConnected(session, {
       onConnectionErrored: (err) => {
         const errStr = coerceErrorString(err);
         this.log?.warn(
@@ -1393,7 +1474,7 @@ var ClientTransport = class extends Transport {
       }
     });
     this.updateSession(connectedSession);
-    this.retryBudget.startRestoringBudget(connectedSession.to);
+    this.retryBudget.startRestoringBudget();
   }
   /**
    * Manually attempts to connect to a client.
@@ -1416,44 +1497,49 @@ var ClientTransport = class extends Transport {
       );
       return;
     }
-    if (!this.retryBudget.hasBudget(to)) {
-      const budgetConsumed = this.retryBudget.getBudgetConsumed(to);
+    if (!this.retryBudget.hasBudget()) {
+      const budgetConsumed = this.retryBudget.getBudgetConsumed();
       const errMsg = `tried to connect to ${to} but retry budget exceeded (more than ${budgetConsumed} attempts in the last ${this.retryBudget.totalBudgetRestoreTime}ms)`;
       this.log?.error(errMsg, session.loggingMetadata);
       this.protocolError(ProtocolError.RetriesExceeded, errMsg);
       return;
     }
-    let sleep = Promise.resolve();
-    const backoffMs = this.retryBudget.getBackoffMs(to);
-    if (backoffMs > 0) {
-      sleep = new Promise((resolve) => setTimeout(resolve, backoffMs));
-    }
+    const backoffMs = this.retryBudget.getBackoffMs();
     this.log?.info(
       `attempting connection to ${to} (${backoffMs}ms backoff)`,
       session.loggingMetadata
     );
-    this.retryBudget.consumeBudget(to);
-    const reconnectPromise = tracing_default.startActiveSpan("connect", async (span) => {
-      try {
-        span.addEvent("backoff", { backoffMs });
-        await sleep;
-        if (this.getStatus() !== "open") {
-          throw new Error("transport state is no longer open");
+    this.retryBudget.consumeBudget();
+    const backingOffSession = ClientSessionStateGraph.transition.NoConnectionToBackingOff(
+      session,
+      backoffMs,
+      {
+        onBackoffFinished: () => {
+          const reconnectPromise = tracing_default.startActiveSpan(
+            "connect",
+            async (span) => {
+              try {
+                return await this.createNewOutgoingConnection(to);
+              } catch (err) {
+                const errStr = coerceErrorString(err);
+                span.recordException(errStr);
+                span.setStatus({ code: import_api3.SpanStatusCode.ERROR });
+                throw err;
+              } finally {
+                span.end();
+              }
+            }
+          );
+          this.onBackoffFinished(backingOffSession, reconnectPromise);
         }
-        span.addEvent("connecting");
-        return await this.createNewOutgoingConnection(to);
-      } catch (err) {
-        const errStr = coerceErrorString(err);
-        span.recordException(errStr);
-        span.setStatus({ code: import_api3.SpanStatusCode.ERROR });
-        throw err;
-      } finally {
-        span.end();
       }
-    });
-    const connectingSession = SessionStateGraph.transition.NoConnectionToConnecting(
+    );
+    this.updateSession(backingOffSession);
+  }
+  onBackoffFinished(session, connPromise) {
+    const connectingSession = ClientSessionStateGraph.transition.BackingOffToConnecting(
       session,
-      reconnectPromise,
+      connPromise,
       {
         onConnectionEstablished: (conn) => {
           this.log?.debug(