@replit/river 0.18.5 → 0.19.2

package/dist/transport/impls/ws/server.js

@@ -1,11 +1,11 @@
 import {
   WebSocketConnection
-} from "../../../chunk-K7CUSLWL.js";
+} from "../../../chunk-JH275HID.js";
 import {
   ServerTransport
-} from "../../../chunk-UABIFWM7.js";
-import "../../../chunk-VH3NGOXQ.js";
-import "../../../chunk-WER2DWCP.js";
+} from "../../../chunk-YFPVQTWL.js";
+import "../../../chunk-D5PVGZPQ.js";
+import "../../../chunk-NBE3D667.js";
 import "../../../chunk-GZ7HCLLM.js";
 
 // transport/impls/ws/server.ts

package/dist/transport/index.cjs

@@ -59,7 +59,8 @@ var PROTOCOL_VERSION = "v1.1";
 var ControlMessageHandshakeRequestSchema = import_typebox.Type.Object({
   type: import_typebox.Type.Literal("HANDSHAKE_REQ"),
   protocolVersion: import_typebox.Type.String(),
-  sessionId: import_typebox.Type.String()
+  sessionId: import_typebox.Type.String(),
+  metadata: import_typebox.Type.Optional(import_typebox.Type.Unknown())
 });
 var ControlMessageHandshakeResponseSchema = import_typebox.Type.Object({
   type: import_typebox.Type.Literal("HANDSHAKE_RESP"),
@@ -83,7 +84,7 @@ var ControlMessagePayloadSchema = import_typebox.Type.Union([
 var OpaqueTransportMessageSchema = TransportMessageSchema(
   import_typebox.Type.Unknown()
 );
-function handshakeRequestMessage(from, to, sessionId) {
+function handshakeRequestMessage(from, to, sessionId, metadata) {
   return {
     id: (0, import_nanoid.nanoid)(),
     from,
@@ -95,7 +96,8 @@ function handshakeRequestMessage(from, to, sessionId) {
     payload: {
       type: "HANDSHAKE_REQ",
       protocolVersion: PROTOCOL_VERSION,
-      sessionId
+      sessionId,
+      metadata
     }
   };
 }
@@ -187,6 +189,12 @@ 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)
    */
@@ -289,27 +297,24 @@ var Session = class {
     this.seq = 0;
     this.ack = 0;
   }
-  sendBufferedMessages() {
-    if (!this.connection) {
-      const msg = `tried sending buffered messages without a connection (if you hit this code path something is seriously wrong)`;
-      log?.error(msg, this.loggingMetadata);
-      throw new Error(msg);
-    }
-    log?.info(
-      `resending ${this.sendBuffer.length} buffered messages`,
-      this.loggingMetadata
-    );
+  sendBufferedMessages(conn) {
+    log?.info(`resending ${this.sendBuffer.length} buffered messages`, {
+      ...this.loggingMetadata,
+      connId: conn.debugId
+    });
     for (const msg of this.sendBuffer) {
       log?.debug(`resending msg`, {
         ...this.loggingMetadata,
-        fullTransportMessage: msg
+        fullTransportMessage: msg,
+        connId: conn.debugId
       });
-      const ok = this.connection.send(this.codec.toBuffer(msg));
+      const ok = conn.send(this.codec.toBuffer(msg));
       if (!ok) {
-        const errMsg = `failed to send buffered message to ${this.to} (if you hit this code path something is seriously wrong)`;
+        const errMsg = `failed to send buffered message to ${this.to} (sus, this is a fresh connection)`;
         log?.error(errMsg, {
           ...this.loggingMetadata,
-          fullTransportMessage: msg
+          fullTransportMessage: msg,
+          connId: conn.debugId
         });
         throw new Error(errMsg);
       }
@@ -340,6 +345,7 @@ var Session = class {
     this.closeStaleConnection(newConn);
     this.cancelGrace();
     this.connection = newConn;
+    this.sendBufferedMessages(newConn);
   }
   beginGrace(cb) {
     log?.info(
@@ -538,6 +544,9 @@ var defaultClientTransportOptions = {
   ...defaultTransportOptions,
   ...defaultConnectionRetryOptions
 };
+var defaultServerTransportOptions = {
+  ...defaultTransportOptions
+};
 var Transport = class {
   /**
    * A flag indicating whether the transport has been destroyed.
@@ -591,37 +600,15 @@ var Transport = class {
    * and we know the identity of the connected client.
    * @param conn The connection object.
    */
-  onConnect(conn, connectedTo, advertisedSessionId) {
+  onConnect(conn, connectedTo, session, isReconnect) {
     this.eventDispatcher.dispatchEvent("connectionStatus", {
       status: "connect",
       conn
     });
-    let oldSession = this.sessions.get(connectedTo);
-    if (oldSession?.advertisedSessionId && oldSession.advertisedSessionId !== advertisedSessionId) {
-      log?.warn(
-        `connection from ${connectedTo} is a different session (id: ${advertisedSessionId}, last connected to: ${oldSession.advertisedSessionId}), killing old session and starting a new one`,
-        oldSession.loggingMetadata
-      );
-      this.deleteSession(oldSession);
-      oldSession = void 0;
+    if (isReconnect) {
+      session.replaceWithNewConnection(conn);
+      log?.info(`reconnected to ${connectedTo}`, session.loggingMetadata);
     }
-    if (oldSession === void 0) {
-      const newSession = this.createSession(connectedTo, conn);
-      newSession.advertisedSessionId = advertisedSessionId;
-      log?.info(
-        `new connection for new session to ${connectedTo}`,
-        newSession.loggingMetadata
-      );
-      return newSession;
-    }
-    oldSession.replaceWithNewConnection(conn);
-    oldSession.sendBufferedMessages();
-    oldSession.advertisedSessionId = advertisedSessionId;
-    log?.info(
-      `new connection for existing session to ${connectedTo}`,
-      oldSession.loggingMetadata
-    );
-    return oldSession;
   }
   createSession(to, conn) {
     const session = new Session(
@@ -637,8 +624,18 @@ var Transport = class {
     });
     return session;
   }
-  getOrCreateSession(to, conn) {
+  getOrCreateSession(to, conn, sessionId) {
     let session = this.sessions.get(to);
+    let isReconnect = session !== void 0;
+    if (session?.advertisedSessionId !== void 0 && sessionId !== void 0 && session.advertisedSessionId !== sessionId) {
+      log?.warn(
+        `session for ${to} already exists but has a different session id (expected: ${session.advertisedSessionId}, got: ${sessionId}), creating a new one`,
+        session.loggingMetadata
+      );
+      this.deleteSession(session);
+      isReconnect = false;
+      session = void 0;
+    }
     if (!session) {
       session = this.createSession(to, conn);
       log?.info(
@@ -646,7 +643,10 @@ var Transport = class {
         session.loggingMetadata
       );
     }
-    return session;
+    if (sessionId !== void 0) {
+      session.advertisedSessionId = sessionId;
+    }
+    return { session, isReconnect };
   }
   deleteSession(session) {
     session.close();
@@ -781,7 +781,7 @@ var Transport = class {
       });
       return void 0;
     }
-    return this.getOrCreateSession(to).send(msg);
+    return this.getOrCreateSession(to).session.send(msg);
   }
   // control helpers
   sendCloseStream(to, streamId) {
@@ -940,11 +940,12 @@ var ClientTransport = class extends Transport {
       connectedTo: parsed.from,
       fullTransportMessage: parsed
     });
-    const session = this.onConnect(
-      conn,
+    const { session, isReconnect } = this.getOrCreateSession(
       parsed.from,
+      conn,
       parsed.payload.status.sessionId
     );
+    this.onConnect(conn, parsed.from, session, isReconnect);
     this.retryBudget.startRestoringBudget(parsed.from);
     return session;
   }
@@ -997,8 +998,13 @@ var ClientTransport = class extends Transport {
           conn.close();
           throw new Error("transport state is no longer open");
         }
-        this.sendHandshake(to, conn);
-        return conn;
+        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 {
@@ -1030,14 +1036,35 @@ var ClientTransport = class extends Transport {
     this.inflightConnectionPromises.delete(session.to);
     super.deleteSession(session);
   }
-  sendHandshake(to, conn) {
-    const session = this.getOrCreateSession(to, conn);
-    const requestMsg = handshakeRequestMessage(this.clientId, to, session.id);
+  async sendHandshake(to, conn) {
+    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`, {
+          clientId: this.clientId,
+          connectedTo: to
+        });
+        this.protocolError(
+          ProtocolError.HandshakeFailed,
+          "handshake metadata did not match schema"
+        );
+        return false;
+      }
+    }
+    const { session } = this.getOrCreateSession(to, conn);
+    const requestMsg = handshakeRequestMessage(
+      this.clientId,
+      to,
+      session.id,
+      metadata
+    );
     log?.debug(`sending handshake request to ${to}`, {
       clientId: this.clientId,
       connectedTo: to
     });
     conn.send(this.codec.toBuffer(requestMsg));
+    return true;
   }
   close() {
     this.retryBudget.close();
@@ -1045,8 +1072,16 @@ var ClientTransport = class extends Transport {
   }
 };
 var ServerTransport = class extends Transport {
+  /**
+   * The options for this transport.
+   */
+  options;
   constructor(clientId, providedOptions) {
     super(clientId, providedOptions);
+    this.options = {
+      ...defaultServerTransportOptions,
+      ...providedOptions
+    };
     log?.info(`initiated server transport`, {
       clientId: this.clientId,
       protocolVersion: PROTOCOL_VERSION
@@ -1074,24 +1109,38 @@ var ServerTransport = class extends Transport {
         conn.close();
       }
     }, this.options.sessionDisconnectGraceMs);
+    const buffer = [];
+    let receivedHandshakeMessage = false;
     const handshakeHandler = (data) => {
-      const maybeSession = this.receiveHandshakeRequestMessage(data, conn);
-      if (!maybeSession) {
-        conn.close();
+      if (receivedHandshakeMessage) {
+        buffer.push(data);
         return;
-      } else {
-        session = maybeSession;
-        clearTimeout(handshakeTimeout);
       }
-      conn.removeDataListener(handshakeHandler);
-      conn.addDataListener((data2) => {
-        const parsed = this.parseMsg(data2);
-        if (!parsed) {
-          conn.close();
-          return;
+      receivedHandshakeMessage = true;
+      clearTimeout(handshakeTimeout);
+      void this.receiveHandshakeRequestMessage(data, conn).then(
+        (maybeSession) => {
+          if (!maybeSession) {
+            conn.close();
+            return;
+          }
+          session = maybeSession;
+          const dataHandler = (data2) => {
+            const parsed = this.parseMsg(data2);
+            if (!parsed) {
+              conn.close();
+              return;
+            }
+            this.handleMsg(parsed);
+          };
+          conn.removeDataListener(handshakeHandler);
+          conn.addDataListener(dataHandler);
+          for (const data2 of buffer) {
+            dataHandler(data2);
+          }
+          buffer.length = 0;
         }
-        this.handleMsg(parsed);
-      });
+      );
     };
     conn.addDataListener(handshakeHandler);
     conn.addCloseListener(() => {
@@ -1112,7 +1161,7 @@ var ServerTransport = class extends Transport {
       );
     });
   }
-  receiveHandshakeRequestMessage(data, conn) {
+  async receiveHandshakeRequestMessage(data, conn) {
     const parsed = this.parseMsg(data);
     if (!parsed) {
       this.protocolError(
@@ -1128,7 +1177,8 @@ var ServerTransport = class extends Transport {
         reason
       });
       conn.send(this.codec.toBuffer(responseMsg2));
-      log?.warn(`${reason}: ${JSON.stringify(parsed)}`, {
+      const logData = typeof parsed.payload === "object" ? { ...parsed, payload: { ...parsed.payload, metadata: "redacted" } } : { ...parsed };
+      log?.warn(`${reason}: ${JSON.stringify(logData)}`, {
         clientId: this.clientId,
         connId: conn.debugId
       });
@@ -1153,7 +1203,73 @@ var ServerTransport = class extends Transport {
       this.protocolError(ProtocolError.HandshakeFailed, reason);
       return false;
     }
-    const session = this.getOrCreateSession(parsed.from, conn);
+    const { session, isReconnect } = this.getOrCreateSession(
+      parsed.from,
+      conn,
+      parsed.payload.sessionId
+    );
+    let handshakeMetadata;
+    if (this.options.handshake) {
+      if (!import_value.Value.Check(
+        this.options.handshake.requestSchema,
+        parsed.payload.metadata
+      )) {
+        const reason = "received malformed handshake metadata";
+        const responseMsg2 = handshakeResponseMessage(
+          this.clientId,
+          parsed.from,
+          { ok: false, reason }
+        );
+        conn.send(this.codec.toBuffer(responseMsg2));
+        log?.warn(`received malformed handshake metadata from ${parsed.from}`, {
+          clientId: this.clientId,
+          connId: conn.debugId
+        });
+        this.protocolError(ProtocolError.HandshakeFailed, reason);
+        this.deleteSession(session);
+        return false;
+      }
+      const parsedMetadata = await this.options.handshake.parse(
+        parsed.payload.metadata,
+        session,
+        isReconnect
+      );
+      if (parsedMetadata === false) {
+        const reason = "rejected by server";
+        const responseMsg2 = handshakeResponseMessage(
+          this.clientId,
+          parsed.from,
+          { ok: false, reason }
+        );
+        conn.send(this.codec.toBuffer(responseMsg2));
+        log?.warn(`rejected handshake from ${parsed.from}`, {
+          clientId: this.clientId,
+          connId: conn.debugId
+        });
+        this.protocolError(ProtocolError.HandshakeFailed, reason);
+        this.deleteSession(session);
+        return false;
+      }
+      if (!import_value.Value.Check(this.options.handshake.parsedSchema, parsedMetadata)) {
+        const reason = "failed to parse handshake metadata";
+        const responseMsg2 = handshakeResponseMessage(
+          this.clientId,
+          parsed.from,
+          { ok: false, reason }
+        );
+        conn.send(this.codec.toBuffer(responseMsg2));
+        log?.error(`failed to parse handshake metadata`, {
+          clientId: this.clientId,
+          connId: conn.debugId
+        });
+        this.protocolError(ProtocolError.HandshakeFailed, reason);
+        this.deleteSession(session);
+        return false;
+      }
+      handshakeMetadata = parsedMetadata;
+    }
+    handshakeMetadata ??= {};
+    session.metadata = handshakeMetadata;
     log?.debug(
       `handshake from ${parsed.from} ok, responding with handshake success`,
       { clientId: this.clientId, connId: conn.debugId }
@@ -1163,7 +1279,8 @@ var ServerTransport = class extends Transport {
       sessionId: session.id
     });
     conn.send(this.codec.toBuffer(responseMsg));
-    return this.onConnect(conn, parsed.from, parsed.payload.sessionId);
+    this.onConnect(conn, parsed.from, session, isReconnect);
+    return session;
   }
 };
 // Annotate the CommonJS export names for ESM import in node: