@replit/river 0.23.16 → 0.24.0

package/dist/router/index.cjs

@@ -31,7 +31,6 @@ __export(router_exports, {
   createClientHandshakeOptions: () => createClientHandshakeOptions,
   createServer: () => createServer,
   createServerHandshakeOptions: () => createServerHandshakeOptions,
-  diffServerSchema: () => diffServerSchema,
   serializeSchema: () => serializeSchema
 });
 module.exports = __toCommonJS(router_exports);
@@ -174,9 +173,9 @@ var ServiceSchema = class _ServiceSchema {
    * You probably don't need this, usually the River server will handle this
    * for you.
    */
-  instantiate() {
+  instantiate(extendedContext) {
     return Object.freeze({
-      state: this.initializeState(),
+      state: this.initializeState(extendedContext),
       procedures: this.procedures
     });
   }
@@ -238,468 +237,6 @@ var ServiceScaffold = class {
   }
 };
 
-// router/diff.ts
-function diffServerSchema(oldServer, newServer, options) {
-  const allServices = /* @__PURE__ */ new Set([
-    ...Object.keys(oldServer.services),
-    ...Object.keys(newServer.services)
-  ]);
-  const breakages = {};
-  for (const serviceName of allServices) {
-    const oldService = oldServer.services[serviceName];
-    const newService = newServer.services[serviceName];
-    const breakage = diffService(oldService, newService, options);
-    if (breakage) {
-      breakages[serviceName] = breakage;
-    }
-  }
-  if (Object.keys(breakages).length) {
-    return { serviceBreakages: breakages };
-  }
-  return null;
-}
-function diffService(oldService, newService, options) {
-  if (!newService) {
-    return options?.allowServiceRemoval ? null : { reason: "removed" };
-  }
-  if (!oldService) {
-    return null;
-  }
-  const allProcedures = /* @__PURE__ */ new Set([
-    ...Object.keys(oldService.procedures),
-    ...Object.keys(newService.procedures)
-  ]);
-  const breakages = {};
-  for (const procedureName of allProcedures) {
-    const aProcedure = oldService.procedures[procedureName];
-    const bProcedure = newService.procedures[procedureName];
-    const breakage = diffProcedure(aProcedure, bProcedure, options);
-    if (breakage) {
-      breakages[procedureName] = breakage;
-    }
-  }
-  if (Object.keys(breakages).length) {
-    return { reason: "modified", procedureBreakages: breakages };
-  }
-  return null;
-}
-function diffProcedure(oldProcedure, newProcedure, options) {
-  if (!newProcedure) {
-    return options?.allowProcedureRemoval ? null : { reason: "removed" };
-  }
-  if (!oldProcedure) {
-    return null;
-  }
-  if (oldProcedure.type !== newProcedure.type) {
-    return {
-      reason: "type-changed",
-      oldType: oldProcedure.type,
-      newType: newProcedure.type
-    };
-  }
-  const inputBreakage = diffProcedureField(
-    oldProcedure.input,
-    newProcedure.input,
-    "client"
-  );
-  const initBreakage = diffProcedureField(
-    oldProcedure.init,
-    newProcedure.init,
-    "client"
-  );
-  const outputBreakage = diffProcedureField(
-    oldProcedure.output,
-    newProcedure.output,
-    "server"
-  );
-  if (inputBreakage ?? initBreakage ?? outputBreakage) {
-    const result = {
-      reason: "modified"
-    };
-    if (inputBreakage) {
-      result.input = inputBreakage;
-    }
-    if (initBreakage) {
-      result.init = initBreakage;
-    }
-    if (outputBreakage) {
-      result.output = outputBreakage;
-    }
-    return result;
-  }
-  return null;
-}
-function diffProcedureField(oldSchema, newSchema, origin) {
-  if (!oldSchema && !newSchema) {
-    return null;
-  }
-  const diffBreakage = diffRequired(oldSchema, newSchema, origin, false, false);
-  if (diffBreakage) {
-    return diffBreakage;
-  }
-  if (!oldSchema || !newSchema) {
-    throw new Error("Appease typescript, this should never happen");
-  }
-  return diffJSONSchema(oldSchema, newSchema, origin);
-}
-function diffRequired(oldSchema, newSchema, origin, oldRequired, newRequired) {
-  if (!newSchema && !oldSchema) {
-    throw new Error("Both old and new schema are undefined");
-  }
-  if (!newSchema) {
-    if (!oldRequired && origin == "server") {
-      return null;
-    }
-    return { reason: "removed-required" };
-  }
-  if (!oldSchema) {
-    if (newRequired && origin === "client") {
-      return { reason: "new-required" };
-    }
-    return null;
-  }
-  if (origin === "client" && !oldRequired && newRequired) {
-    return { reason: "new-required" };
-  }
-  if (origin === "server" && oldRequired && !newRequired) {
-    return { reason: "removed-required" };
-  }
-  return null;
-}
-function diffJSONSchema(oldSchema, newSchema, origin) {
-  if (oldSchema.type !== newSchema.type) {
-    return {
-      reason: "type-changed",
-      oldType: getReportingType(oldSchema),
-      newType: getReportingType(newSchema)
-    };
-  }
-  if (getReportingType(oldSchema) !== getReportingType(newSchema)) {
-    return {
-      reason: "type-changed",
-      oldType: getReportingType(oldSchema),
-      newType: getReportingType(newSchema)
-    };
-  }
-  if ("const" in oldSchema && "const" in newSchema && oldSchema.const !== newSchema.const) {
-    return {
-      reason: "type-changed",
-      oldType: `${getReportingType(oldSchema)}-const-${oldSchema.const}`,
-      newType: `${getReportingType(newSchema)}-const-${newSchema.const}`
-    };
-  }
-  if ("const" in oldSchema && !("const" in newSchema) && origin === "server") {
-    return {
-      reason: "type-changed",
-      oldType: `${getReportingType(oldSchema)}-const-${oldSchema.const}`,
-      newType: getReportingType(newSchema)
-    };
-  }
-  if ("const" in newSchema && !("const" in oldSchema) && origin === "client") {
-    return {
-      reason: "type-changed",
-      oldType: getReportingType(oldSchema),
-      newType: `${getReportingType(newSchema)}-const-${newSchema.const}`
-    };
-  }
-  const breakages = {};
-  if ("$ref" in newSchema) {
-    if (newSchema.$ref !== oldSchema.$ref) {
-      return {
-        reason: "type-changed",
-        oldType: getReportingType(oldSchema),
-        newType: getReportingType(newSchema)
-      };
-    }
-  } else if ("not" in newSchema) {
-    const notBreakage = diffJSONSchema(
-      oldSchema.not,
-      newSchema.not,
-      origin
-    );
-    if (notBreakage) {
-      breakages.not = notBreakage;
-    }
-  } else if ("anyOf" in newSchema) {
-    const oldAnyOfStringified = oldSchema.anyOf.map((el) => JSON.stringify(el)).sort();
-    const newAnyOfStringified = newSchema.anyOf.map((el) => JSON.stringify(el)).sort();
-    const anyOfBreakages = {};
-    for (let i = 0; i < oldAnyOfStringified.length; i++) {
-      if (newAnyOfStringified.includes(oldAnyOfStringified[i])) {
-        continue;
-      }
-      if (!newAnyOfStringified[i]) {
-        if (origin === "server") {
-          continue;
-        }
-        anyOfBreakages[`old-${i}`] = { reason: "removed-required" };
-      } else {
-        const breakage = diffJSONSchema(
-          JSON.parse(oldAnyOfStringified[i]),
-          JSON.parse(newAnyOfStringified[i]),
-          origin
-        );
-        if (breakage) {
-          anyOfBreakages[`old-${i}`] = breakage;
-        }
-      }
-    }
-    for (let i = 0; i < newAnyOfStringified.length; i++) {
-      if (oldAnyOfStringified.includes(newAnyOfStringified[i])) {
-        continue;
-      }
-      if (!oldAnyOfStringified[i]) {
-        if (origin === "client") {
-          continue;
-        }
-        anyOfBreakages[`new-${i}`] = { reason: "new-required" };
-      } else {
-        const breakage = diffJSONSchema(
-          JSON.parse(oldAnyOfStringified[i]),
-          JSON.parse(newAnyOfStringified[i]),
-          origin
-        );
-        if (breakage) {
-          anyOfBreakages[`new-${i}`] = breakage;
-        }
-      }
-    }
-    if (Object.keys(anyOfBreakages).length > 0) {
-      breakages.anyOf = {
-        reason: "field-breakage",
-        fieldBreakages: anyOfBreakages
-      };
-    }
-  } else if ("oneOf" in newSchema) {
-    throw new Error("oneOf is not supported, typebox does not emit it");
-  } else if ("allOf" in newSchema) {
-    if (newSchema.allOf.length !== oldSchema.allOf.length) {
-      breakages.allOf = {
-        reason: "type-changed",
-        oldType: `${oldSchema.allOf}`,
-        newType: `${newSchema.allOf}`
-      };
-    } else {
-      for (let i = 0; i < newSchema.allOf.length; i++) {
-        const breakage = diffJSONSchema(
-          oldSchema.allOf[i],
-          newSchema.allOf[i],
-          origin
-        );
-        if (breakage) {
-          breakages.allOf = breakage;
-          break;
-        }
-      }
-    }
-  } else if (newSchema.type === "array") {
-    const itemsBreakages = diffJSONSchema(
-      oldSchema.items,
-      newSchema.items,
-      origin
-    );
-    if (itemsBreakages) {
-      breakages.items = itemsBreakages;
-    }
-    if (oldSchema.minItems < newSchema.minItems) {
-      if (origin === "client") {
-        breakages.minItems = {
-          reason: "type-changed",
-          oldType: `${oldSchema.minItems}`,
-          newType: `${newSchema.minItems}`
-        };
-      }
-    } else if (oldSchema.minItems > newSchema.minItems) {
-      if (origin === "server") {
-        breakages.minItems = {
-          reason: "type-changed",
-          oldType: `${oldSchema.minItems}`,
-          newType: `${newSchema.minItems}`
-        };
-      }
-    }
-    if (oldSchema.maxItems < newSchema.maxItems) {
-      if (origin === "server") {
-        breakages.maxItems = {
-          reason: "type-changed",
-          oldType: `${oldSchema.maxItems}`,
-          newType: `${newSchema.maxItems}`
-        };
-      }
-    } else if (oldSchema.maxItems > newSchema.maxItems) {
-      if (origin === "client") {
-        breakages.maxItems = {
-          reason: "type-changed",
-          oldType: `${oldSchema.maxItems}`,
-          newType: `${newSchema.maxItems}`
-        };
-      }
-    }
-    if (!oldSchema.uniqueItems && newSchema.uniqueItems && origin === "client") {
-      breakages.uniqueItems = {
-        reason: "type-changed",
-        oldType: `${!!oldSchema.uniqueItems}`,
-        newType: `${!!newSchema.uniqueItems}`
-      };
-    }
-    if ("contains" in newSchema !== "contains" in oldSchema) {
-      if ("contains" in newSchema && !("contains" in oldSchema) && origin === "client") {
-        breakages.contains = {
-          reason: "type-changed",
-          oldType: "no-contains",
-          newType: "contains"
-        };
-      }
-    } else if ("contains" in newSchema) {
-      const containsBreakage = diffJSONSchema(
-        oldSchema.contains,
-        newSchema.contains,
-        origin
-      );
-      if (containsBreakage) {
-        breakages.contains = containsBreakage;
-      }
-    }
-    if (oldSchema.minContains < newSchema.minContains) {
-      if (origin === "client") {
-        breakages.minContains = {
-          reason: "type-changed",
-          oldType: `${oldSchema.minContains}`,
-          newType: `${newSchema.minContains}`
-        };
-      }
-    } else if (oldSchema.minContains > newSchema.minContains) {
-      if (origin === "server") {
-        breakages.minContains = {
-          reason: "type-changed",
-          oldType: `${oldSchema.minContains}`,
-          newType: `${newSchema.minContains}`
-        };
-      }
-    }
-    if (oldSchema.maxContains < newSchema.maxContains) {
-      if (origin === "server") {
-        breakages.maxContains = {
-          reason: "type-changed",
-          oldType: `${oldSchema.maxContains}`,
-          newType: `${newSchema.maxContains}`
-        };
-      }
-    } else if (oldSchema.maxContains > newSchema.maxContains) {
-      if (origin === "client") {
-        breakages.maxContains = {
-          reason: "type-changed",
-          oldType: `${oldSchema.maxContains}`,
-          newType: `${newSchema.maxContains}`
-        };
-      }
-    }
-  } else if (newSchema.type === "object") {
-    if ("properties" in newSchema !== "properties" in oldSchema) {
-      return {
-        reason: "type-changed",
-        oldType: "properties" in oldSchema ? "probably-object" : "probably-record",
-        newType: "properties" in newSchema ? "probably-object" : "probably-record"
-      };
-    }
-    if ("properties" in newSchema) {
-      const propertiesBreakages = diffObjectProperties(
-        oldSchema.properties,
-        newSchema.properties,
-        origin,
-        oldSchema.required,
-        newSchema.required
-      );
-      if (Object.keys(propertiesBreakages).length) {
-        breakages.properties = {
-          reason: "field-breakage",
-          fieldBreakages: propertiesBreakages
-        };
-      }
-    }
-    if ("patternProperties" in newSchema) {
-      const patternPropertiesBreakages = diffObjectProperties(
-        oldSchema.patternProperties,
-        newSchema.patternProperties,
-        origin,
-        oldSchema.required,
-        newSchema.required
-      );
-      if (Object.keys(patternPropertiesBreakages).length) {
-        breakages.patternProperties = {
-          reason: "field-breakage",
-          fieldBreakages: patternPropertiesBreakages
-        };
-      }
-    }
-    if ("additionalProperties" in newSchema || "additionalProperties" in oldSchema) {
-      throw new Error("additionalProperties is not supported");
-    }
-    if ("minProperties" in newSchema || "minProperties" in oldSchema) {
-      throw new Error("minProperties is not supported");
-    }
-    if ("maxProperties" in newSchema || "maxProperties" in oldSchema) {
-      throw new Error("maxProperties is not supported");
-    }
-  }
-  if (Object.keys(breakages).length) {
-    return {
-      reason: "field-breakage",
-      fieldBreakages: breakages
-    };
-  }
-  return null;
-}
-function diffObjectProperties(oldProperties, newProperties, origin, oldRequiredProperties = [], newRequiredProperties = []) {
-  const allProperties = /* @__PURE__ */ new Set([
-    ...Object.keys(oldProperties),
-    ...Object.keys(newProperties)
-  ]);
-  const breakages = {};
-  for (const propertyName of allProperties) {
-    const requiredBreakage = diffRequired(
-      oldProperties[propertyName],
-      newProperties[propertyName],
-      origin,
-      oldRequiredProperties.includes(propertyName),
-      newRequiredProperties.includes(propertyName)
-    );
-    if (requiredBreakage) {
-      breakages[propertyName] = requiredBreakage;
-    } else if (oldProperties[propertyName] && newProperties[propertyName]) {
-      const propertyBreakage = diffJSONSchema(
-        oldProperties[propertyName],
-        newProperties[propertyName],
-        origin
-      );
-      if (propertyBreakage) {
-        breakages[propertyName] = propertyBreakage;
-      }
-    }
-  }
-  return breakages;
-}
-function getReportingType(schema) {
-  if ("not" in schema) {
-    return "not";
-  }
-  if ("anyOf" in schema) {
-    return "anyOf";
-  }
-  if ("allOf" in schema) {
-    return "allOf";
-  }
-  if ("$ref" in schema) {
-    return "$ref";
-  }
-  if (schema.type && typeof schema.type === "string") {
-    return schema.type;
-  }
-  throw new Error(
-    "Subschema not supported, probably a conditional subschema. Check logs."
-  );
-}
-
 // router/procedures.ts
 var import_typebox2 = require("@sinclair/typebox");
 function rpc({
@@ -1075,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(),
@@ -1109,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([
@@ -1130,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)
     })
   ])
 });
@@ -1143,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 */
@@ -1156,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";
@@ -1187,7 +752,7 @@ function Err(error) {
 var import_api = require("@opentelemetry/api");
 
 // package.json
-var version = "0.23.16";
+var version = "0.24.0";
 
 // tracing/index.ts
 function getPropagationContext(ctx) {
@@ -1281,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];
@@ -1291,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);
@@ -1315,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",
@@ -1360,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",
@@ -1401,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) {
@@ -1437,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",
@@ -1476,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(
@@ -1493,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",
@@ -1533,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) => {
@@ -1584,18 +1149,18 @@ 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 = {};
     this.services = instances;
     this.contextMap = /* @__PURE__ */ new Map();
     for (const [name, service] of Object.entries(services)) {
-      const instance = service.instantiate();
+      const instance = service.instantiate(extendedContext ?? {});
       instances[name] = instance;
       this.contextMap.set(instance, {
         ...extendedContext,
@@ -1606,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;
@@ -1643,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)) {
@@ -1711,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());
         })()
       ) : (
@@ -1738,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
@@ -1752,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 });
@@ -1763,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;
@@ -1778,7 +1361,6 @@ var RiverServer = class {
       to: message.to,
       from: message.from,
       streamId: message.streamId,
-      session,
       metadata: sessionMeta
     };
     switch (procedure.type) {
@@ -1897,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 {
@@ -1917,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 {
@@ -1932,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;
     }
@@ -1945,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) {
@@ -1983,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);
         }
       }
     }
@@ -2045,7 +1623,6 @@ function createServerHandshakeOptions(schema, validate) {
   createClientHandshakeOptions,
   createServer,
   createServerHandshakeOptions,
-  diffServerSchema,
   serializeSchema
 });
 //# sourceMappingURL=index.cjs.map