@replit/river 0.23.18 → 0.24.0

package/dist/router/index.cjs

@@ -612,7 +612,15 @@ function _pushable(getNext, options) {
 
 // transport/message.ts
 var import_typebox3 = require("@sinclair/typebox");
+
+// transport/id.ts
 var import_nanoid = require("nanoid");
+var alphabet = (0, import_nanoid.customAlphabet)(
+  "1234567890abcdefghijklmnopqrstuvxyzABCDEFGHIJKLMNOPQRSTUVXYZ"
+);
+var generateId = () => alphabet(12);
+
+// transport/message.ts
 var TransportMessageSchema = (t) => import_typebox3.Type.Object({
   id: import_typebox3.Type.String(),
   from: import_typebox3.Type.String(),
@@ -646,18 +654,29 @@ var ControlMessageHandshakeRequestSchema = import_typebox3.Type.Object({
    * used by the server to know whether this is a new or a reestablished connection, and whether it
    * is compatible with what it already has.
    */
-  expectedSessionState: import_typebox3.Type.Optional(
-    import_typebox3.Type.Object({
-      /**
-       * reconnect is set to true if the client explicitly wants to reestablish an existing
-       * connection.
-       */
-      reconnect: import_typebox3.Type.Boolean(),
-      nextExpectedSeq: import_typebox3.Type.Integer()
-    })
-  ),
+  expectedSessionState: import_typebox3.Type.Object({
+    // what the client expects the server to send next
+    nextExpectedSeq: import_typebox3.Type.Integer(),
+    // TODO: remove optional once we know all servers
+    // are nextSentSeq here
+    // what the server expects the client to send next
+    nextSentSeq: import_typebox3.Type.Optional(import_typebox3.Type.Integer())
+  }),
   metadata: import_typebox3.Type.Optional(import_typebox3.Type.Unknown())
 });
+var HandshakeErrorRetriableResponseCodes = import_typebox3.Type.Union([
+  import_typebox3.Type.Literal("SESSION_STATE_MISMATCH")
+]);
+var HandshakeErrorFatalResponseCodes = import_typebox3.Type.Union([
+  import_typebox3.Type.Literal("MALFORMED_HANDSHAKE_META"),
+  import_typebox3.Type.Literal("MALFORMED_HANDSHAKE"),
+  import_typebox3.Type.Literal("PROTOCOL_VERSION_MISMATCH"),
+  import_typebox3.Type.Literal("REJECTED_BY_CUSTOM_HANDLER")
+]);
+var HandshakeErrorResponseCodes = import_typebox3.Type.Union([
+  HandshakeErrorRetriableResponseCodes,
+  HandshakeErrorFatalResponseCodes
+]);
 var ControlMessageHandshakeResponseSchema = import_typebox3.Type.Object({
   type: import_typebox3.Type.Literal("HANDSHAKE_RESP"),
   status: import_typebox3.Type.Union([
@@ -667,7 +686,10 @@ var ControlMessageHandshakeResponseSchema = import_typebox3.Type.Object({
     }),
     import_typebox3.Type.Object({
       ok: import_typebox3.Type.Literal(false),
-      reason: import_typebox3.Type.String()
+      reason: import_typebox3.Type.String(),
+      // TODO: remove optional once we know all servers
+      // are sending code here
+      code: import_typebox3.Type.Optional(HandshakeErrorResponseCodes)
     })
   ])
 });
@@ -680,6 +702,15 @@ var ControlMessagePayloadSchema = import_typebox3.Type.Union([
 var OpaqueTransportMessageSchema = TransportMessageSchema(
   import_typebox3.Type.Unknown()
 );
+function closeStreamMessage(streamId) {
+  return {
+    streamId,
+    controlFlags: 4 /* StreamClosedBit */,
+    payload: {
+      type: "CLOSE"
+    }
+  };
+}
 function isStreamOpen(controlFlag) {
   return (
     /* eslint-disable-next-line @typescript-eslint/no-unsafe-enum-comparison */
@@ -693,9 +724,6 @@ function isStreamClose(controlFlag) {
   );
 }
 
-// router/client.ts
-var import_nanoid2 = require("nanoid");
-
 // router/result.ts
 var import_typebox4 = require("@sinclair/typebox");
 var UNCAUGHT_ERROR = "UNCAUGHT_ERROR";
@@ -724,7 +752,7 @@ function Err(error) {
 var import_api = require("@opentelemetry/api");
 
 // package.json
-var version = "0.23.18";
+var version = "0.24.0";
 
 // tracing/index.ts
 function getPropagationContext(ctx) {
@@ -818,7 +846,7 @@ function createClient(transport, serverId, providedClientOptions = {}) {
   }
   const options = { ...defaultClientOptions, ...providedClientOptions };
   if (options.eagerlyConnect) {
-    void transport.connect(serverId);
+    transport.connect(serverId);
   }
   return _createRecursiveProxy(async (opts) => {
     const [serviceName, procName, procType] = [...opts.path];
@@ -828,8 +856,8 @@ function createClient(transport, serverId, providedClientOptions = {}) {
       );
     }
     const [input] = opts.args;
-    if (options.connectOnInvoke && !transport.connections.has(serverId)) {
-      void transport.connect(serverId);
+    if (options.connectOnInvoke && !transport.sessions.has(serverId)) {
+      transport.connect(serverId);
     }
     if (procType === "rpc") {
       return handleRpc(transport, serverId, input, serviceName, procName);
@@ -852,7 +880,7 @@ function createSessionDisconnectHandler(from, cb) {
   };
 }
 function handleRpc(transport, serverId, input, serviceName, procedureName) {
-  const streamId = (0, import_nanoid2.nanoid)();
+  const streamId = generateId();
   const { span, ctx } = createProcTelemetryInfo(
     transport,
     "rpc",
@@ -897,7 +925,7 @@ function handleRpc(transport, serverId, input, serviceName, procedureName) {
   return responsePromise;
 }
 function handleStream(transport, serverId, init, serviceName, procedureName) {
-  const streamId = (0, import_nanoid2.nanoid)();
+  const streamId = generateId();
   const { span, ctx } = createProcTelemetryInfo(
     transport,
     "stream",
@@ -938,7 +966,7 @@ function handleStream(transport, serverId, init, serviceName, procedureName) {
     }
     if (!healthyClose)
       return;
-    transport.sendCloseStream(serverId, streamId);
+    transport.send(serverId, closeStreamMessage(streamId));
   };
   void pipeInputToTransport();
   function onMessage(msg) {
@@ -974,7 +1002,7 @@ function handleStream(transport, serverId, init, serviceName, procedureName) {
   return [inputStream, outputStream, cleanup];
 }
 function handleSubscribe(transport, serverId, input, serviceName, procedureName) {
-  const streamId = (0, import_nanoid2.nanoid)();
+  const streamId = generateId();
   const { span, ctx } = createProcTelemetryInfo(
     transport,
     "subscription",
@@ -1013,7 +1041,7 @@ function handleSubscribe(transport, serverId, input, serviceName, procedureName)
     cleanup();
     if (!healthyClose)
       return;
-    transport.sendCloseStream(serverId, streamId);
+    transport.send(serverId, closeStreamMessage(streamId));
   };
   const onSessionStatus = createSessionDisconnectHandler(serverId, () => {
     outputStream.push(
@@ -1030,7 +1058,7 @@ function handleSubscribe(transport, serverId, input, serviceName, procedureName)
   return [outputStream, closeHandler];
 }
 function handleUpload(transport, serverId, init, serviceName, procedureName) {
-  const streamId = (0, import_nanoid2.nanoid)();
+  const streamId = generateId();
   const { span, ctx } = createProcTelemetryInfo(
     transport,
     "upload",
@@ -1070,7 +1098,7 @@ function handleUpload(transport, serverId, init, serviceName, procedureName) {
     }
     if (!healthyClose)
       return;
-    transport.sendCloseStream(serverId, streamId);
+    transport.send(serverId, closeStreamMessage(streamId));
   };
   void pipeInputToTransport();
   const responsePromise = new Promise((resolve) => {
@@ -1121,11 +1149,11 @@ var RiverServer = class {
   transport;
   services;
   contextMap;
-  // map of streamId to ProcStream
   streamMap;
-  // map of client to their open streams by streamId
-  clientStreams;
-  disconnectedSessions;
+  sessionToStreamId;
+  // streams that are in the process of being cleaned up
+  // this is to prevent output handlers from sending after the stream is cleaned up
+  sessionsBeingCleanedUp = /* @__PURE__ */ new Set();
   log;
   constructor(transport, services, handshakeOptions, extendedContext) {
     const instances = {};
@@ -1143,9 +1171,8 @@ var RiverServer = class {
       transport.extendHandshake(handshakeOptions);
     }
     this.transport = transport;
-    this.disconnectedSessions = /* @__PURE__ */ new Set();
     this.streamMap = /* @__PURE__ */ new Map();
-    this.clientStreams = /* @__PURE__ */ new Map();
+    this.sessionToStreamId = /* @__PURE__ */ new Map();
     this.transport.addEventListener("message", this.onMessage);
     this.transport.addEventListener("sessionStatus", this.onSessionStatus);
     this.log = transport.log;
@@ -1180,24 +1207,35 @@ var RiverServer = class {
     }
     await this.pushToStream(procStream, message, isInitMessage);
   };
-  // cleanup streams on session close
   onSessionStatus = async (evt) => {
-    if (evt.status !== "disconnect")
+    const streamsFromThisClient = this.sessionToStreamId.get(evt.session.id);
+    const cleanupStreams = async (ids) => {
+      this.sessionsBeingCleanedUp.add(evt.session.id);
+      await Promise.all(Array.from(ids).map(this.cleanupStream));
+      this.sessionToStreamId.delete(evt.session.id);
+      this.sessionsBeingCleanedUp.delete(evt.session.id);
+    };
+    if (evt.status === "connect") {
+      if (streamsFromThisClient) {
+        this.log?.error(
+          `got session connect from ${evt.session.to} but there are still streams open from this session`,
+          {
+            clientId: this.transport.clientId,
+            tags: ["invariant-violation"]
+          }
+        );
+        await cleanupStreams(streamsFromThisClient);
+      }
+      this.sessionToStreamId.set(evt.session.id, /* @__PURE__ */ new Set());
       return;
-    const disconnectedClientId = evt.session.to;
+    }
     this.log?.info(
-      `got session disconnect from ${disconnectedClientId}, cleaning up streams`,
+      `got session disconnect from ${evt.session.to}, cleaning up streams for session`,
       evt.session.loggingMetadata
     );
-    const streamsFromThisClient = this.clientStreams.get(disconnectedClientId);
     if (!streamsFromThisClient)
       return;
-    this.disconnectedSessions.add(disconnectedClientId);
-    await Promise.all(
-      Array.from(streamsFromThisClient).map(this.cleanupStream)
-    );
-    this.disconnectedSessions.delete(disconnectedClientId);
-    this.clientStreams.delete(disconnectedClientId);
+    await cleanupStreams(streamsFromThisClient);
   };
   createNewProcStream(message) {
     if (!isStreamOpen(message.controlFlags)) {
@@ -1248,26 +1286,34 @@ var RiverServer = class {
       });
       return;
     }
+    const {
+      to,
+      id: sessionId,
+      loggingMetadata: sessionLoggingMetadata
+    } = session;
     const procedure = service.procedures[message.procedureName];
     const incoming = pushable({ objectMode: true });
     const outgoing = pushable({ objectMode: true });
     const needsClose = procedure.type === "subscription" || procedure.type === "stream";
     const disposables = [];
+    const wrappedSend = (payload) => {
+      if (!this.sessionsBeingCleanedUp.has(sessionId)) {
+        this.transport.send(to, payload);
+      }
+    };
     const outputHandler = (
       // sending outgoing messages back to client
       needsClose ? (
         // subscription and stream case, we need to send a close bit after the response stream
         (async () => {
           for await (const response of outgoing) {
-            this.transport.send(session.to, {
+            wrappedSend({
               streamId: message.streamId,
               controlFlags: 0,
               payload: response
             });
           }
-          if (!this.disconnectedSessions.has(message.from)) {
-            this.transport.sendCloseStream(session.to, message.streamId);
-          }
+          wrappedSend(closeStreamMessage(message.streamId));
           disposables.forEach((d) => d());
         })()
       ) : (
@@ -1275,7 +1321,7 @@ var RiverServer = class {
         (async () => {
           const response = await outgoing.next().then((res) => res.value);
           if (response) {
-            this.transport.send(session.to, {
+            wrappedSend({
               streamId: message.streamId,
               controlFlags: 4 /* StreamClosedBit */,
               payload: response
@@ -1289,7 +1335,7 @@ var RiverServer = class {
       const errorMsg = coerceErrorString(err);
       this.log?.error(
         `procedure ${message.serviceName}.${message.procedureName} threw an uncaught error: ${errorMsg}`,
-        session.loggingMetadata
+        sessionLoggingMetadata
       );
       span.recordException(err instanceof Error ? err : new Error(errorMsg));
       span.setStatus({ code: import_api2.SpanStatusCode.ERROR });
@@ -1300,10 +1346,10 @@ var RiverServer = class {
         })
       );
     };
-    const sessionMeta = this.transport.sessionHandshakeMetadata.get(session);
+    const sessionMeta = this.transport.sessionHandshakeMetadata.get(to);
     if (!sessionMeta) {
       this.log?.error(`session doesn't have handshake metadata`, {
-        ...session.loggingMetadata,
+        ...sessionLoggingMetadata,
         tags: ["invariant-violation"]
       });
       return;
@@ -1315,7 +1361,6 @@ var RiverServer = class {
       to: message.to,
       from: message.from,
       streamId: message.streamId,
-      session,
       metadata: sessionMeta
     };
     switch (procedure.type) {
@@ -1434,9 +1479,7 @@ var RiverServer = class {
                   initMessage.value,
                   incoming
                 );
-                if (!this.disconnectedSessions.has(message.from)) {
-                  outgoing.push(outputMessage);
-                }
+                outgoing.push(outputMessage);
               } catch (err) {
                 errorHandler(err, span);
               } finally {
@@ -1454,9 +1497,7 @@ var RiverServer = class {
                   serviceContextWithTransportInfo,
                   incoming
                 );
-                if (!this.disconnectedSessions.has(message.from)) {
-                  outgoing.push(outputMessage);
-                }
+                outgoing.push(outputMessage);
               } catch (err) {
                 errorHandler(err, span);
               } finally {
@@ -1469,7 +1510,7 @@ var RiverServer = class {
       default:
         this.log?.warn(
           `got request for invalid procedure type ${procedure.type} at ${message.serviceName}.${message.procedureName}`,
-          { ...session.loggingMetadata, transportMessage: message }
+          { ...sessionLoggingMetadata, transportMessage: message }
         );
         return;
     }
@@ -1482,9 +1523,9 @@ var RiverServer = class {
       promises: { inputHandler, outputHandler }
     };
     this.streamMap.set(message.streamId, procStream);
-    const streamsFromThisClient = this.clientStreams.get(message.from) ?? /* @__PURE__ */ new Set();
-    streamsFromThisClient.add(message.streamId);
-    this.clientStreams.set(message.from, streamsFromThisClient);
+    const streamsForThisSession = this.sessionToStreamId.get(sessionId) ?? /* @__PURE__ */ new Set();
+    streamsForThisSession.add(message.streamId);
+    this.sessionToStreamId.set(sessionId, streamsForThisSession);
     return procStream;
   }
   async pushToStream(procStream, message, isInit) {
@@ -1520,11 +1561,11 @@ var RiverServer = class {
     }
     if (isStreamClose(message.controlFlags)) {
       await this.cleanupStream(message.streamId);
-      const streamsFromThisClient = this.clientStreams.get(message.from);
+      const streamsFromThisClient = this.sessionToStreamId.get(message.from);
       if (streamsFromThisClient) {
         streamsFromThisClient.delete(message.streamId);
         if (streamsFromThisClient.size === 0) {
-          this.clientStreams.delete(message.from);
+          this.sessionToStreamId.delete(message.from);
         }
       }
     }