@replit/river 0.21.1 → 0.23.0

package/dist/{chunk-NCXUFDVL.js → chunk-VOJVLWVX.js}

@@ -1,24 +1,30 @@
 import {
   ControlMessagePayloadSchema,
   coerceErrorString,
+  createHandlerSpan,
+  createProcTelemetryInfo,
+  getPropagationContext,
   isStreamClose,
-  isStreamOpen,
-  tracing_default
-} from "./chunk-5WFL722S.js";
+  isStreamOpen
+} from "./chunk-PCKHBAVP.js";
 import {
   log
 } from "./chunk-OTQNCLFH.js";
 
 // router/services.ts
 import { Type } from "@sinclair/typebox";
-function serializeSchema(services) {
-  return Object.entries(services).reduce(
-    (acc, [name, value]) => {
-      acc[name] = value.serialize();
-      return acc;
-    },
-    {}
-  );
+function serializeSchema(services, handshakeSchema) {
+  const serializedServiceObject = Object.entries(services).reduce((acc, [name, value]) => {
+    acc[name] = value.serialize();
+    return acc;
+  }, {});
+  const schema = {
+    services: serializedServiceObject
+  };
+  if (handshakeSchema) {
+    schema.handshakeSchema = Type.Strict(handshakeSchema);
+  }
+  return schema;
 }
 var ServiceSchema = class _ServiceSchema {
   /**
@@ -298,14 +304,6 @@ var Procedure = {
   stream
 };
 
-// router/client.ts
-import {
-  context,
-  propagation,
-  SpanKind,
-  SpanStatusCode
-} from "@opentelemetry/api";
-
 // node_modules/p-defer/index.js
 function pDefer() {
   const deferred = {};
@@ -642,7 +640,10 @@ var defaultClientOptions = {
   connectOnInvoke: true,
   eagerlyConnect: true
 };
-var createClient = (transport, serverId, providedClientOptions = {}) => {
+function createClient(transport, serverId, providedClientOptions = {}) {
+  if (providedClientOptions.handshakeOptions) {
+    transport.extendHandshake(providedClientOptions.handshakeOptions);
+  }
   const options = { ...defaultClientOptions, ...providedClientOptions };
   if (options.eagerlyConnect) {
     void transport.connect(serverId);
@@ -678,7 +679,7 @@ var createClient = (transport, serverId, providedClientOptions = {}) => {
       throw new Error(`invalid river call, unknown procedure type ${procType}`);
     }
   }, []);
-};
+}
 function createSessionDisconnectHandler(from, cb) {
   return (evt) => {
     if (evt.status === "disconnect" && evt.session.to === from) {
@@ -688,324 +689,256 @@ function createSessionDisconnectHandler(from, cb) {
 }
 function handleRpc(transport, serverId, input, serviceName, procedureName) {
   const streamId = nanoid();
-  return tracing_default.startActiveSpan(
-    `${serviceName}.${procedureName}`,
-    {
-      attributes: {
-        component: "river",
-        "river.method.kind": "rpc",
-        "river.method.service": serviceName,
-        "river.method.name": procedureName,
-        "river.streamId": streamId,
-        "span.kind": "client"
-      },
-      kind: SpanKind.CLIENT
-    },
-    (span) => {
-      const tracing = { traceparent: "", tracestate: "" };
-      propagation.inject(context.active(), tracing);
-      transport.send(serverId, {
-        streamId,
-        serviceName,
-        procedureName,
-        tracing,
-        payload: input,
-        controlFlags: 2 /* StreamOpenBit */ | 4 /* StreamClosedBit */
-      });
-      const responsePromise = new Promise((resolve) => {
-        const onSessionStatus = createSessionDisconnectHandler(serverId, () => {
-          span.setStatus({ code: SpanStatusCode.ERROR });
-          cleanup();
-          resolve(
-            Err({
-              code: UNEXPECTED_DISCONNECT,
-              message: `${serverId} unexpectedly disconnected`
-            })
-          );
-        });
-        function cleanup() {
-          transport.removeEventListener("message", onMessage);
-          transport.removeEventListener("sessionStatus", onSessionStatus);
-          span.end();
-        }
-        function onMessage(msg) {
-          if (msg.streamId !== streamId)
-            return;
-          if (msg.to !== transport.clientId)
-            return;
-          if (msg.payload && typeof msg.payload === "object" && "ok" in msg.payload) {
-            span.setStatus({
-              code: msg.payload.ok ? SpanStatusCode.OK : SpanStatusCode.ERROR
-            });
-          }
-          cleanup();
-          resolve(msg.payload);
-        }
-        transport.addEventListener("message", onMessage);
-        transport.addEventListener("sessionStatus", onSessionStatus);
-      });
-      return responsePromise;
-    }
+  const { span, ctx } = createProcTelemetryInfo(
+    "rpc",
+    serviceName,
+    procedureName,
+    streamId
   );
+  transport.send(serverId, {
+    streamId,
+    serviceName,
+    procedureName,
+    payload: input,
+    tracing: getPropagationContext(ctx),
+    controlFlags: 2 /* StreamOpenBit */ | 4 /* StreamClosedBit */
+  });
+  const responsePromise = new Promise((resolve) => {
+    const onSessionStatus = createSessionDisconnectHandler(serverId, () => {
+      cleanup();
+      resolve(
+        Err({
+          code: UNEXPECTED_DISCONNECT,
+          message: `${serverId} unexpectedly disconnected`
+        })
+      );
+    });
+    function cleanup() {
+      transport.removeEventListener("message", onMessage);
+      transport.removeEventListener("sessionStatus", onSessionStatus);
+      span.end();
+    }
+    function onMessage(msg) {
+      if (msg.streamId !== streamId)
+        return;
+      if (msg.to !== transport.clientId)
+        return;
+      cleanup();
+      resolve(msg.payload);
+    }
+    transport.addEventListener("message", onMessage);
+    transport.addEventListener("sessionStatus", onSessionStatus);
+  });
+  return responsePromise;
 }
 function handleStream(transport, serverId, init, serviceName, procedureName) {
   const streamId = nanoid();
-  return tracing_default.startActiveSpan(
-    `${serviceName}.${procedureName}`,
-    {
-      attributes: {
-        component: "river",
-        "river.method.kind": init ? "stream-with-init" : "stream",
-        "river.method.service": serviceName,
-        "river.method.name": procedureName,
-        "river.streamId": streamId,
-        "span.kind": "client"
-      },
-      kind: SpanKind.CLIENT
-    },
-    (span) => {
-      const tracing = { traceparent: "", tracestate: "" };
-      propagation.inject(context.active(), tracing);
-      const inputStream = pushable({ objectMode: true });
-      const outputStream = pushable({ objectMode: true });
-      let firstMessage = true;
-      let healthyClose = true;
-      if (init) {
-        transport.send(serverId, {
-          streamId,
-          serviceName,
-          procedureName,
-          tracing,
-          payload: init,
-          controlFlags: 2 /* StreamOpenBit */
-        });
-        firstMessage = false;
-      }
-      const pipeInputToTransport = async () => {
-        for await (const rawIn of inputStream) {
-          const m = {
-            streamId,
-            payload: rawIn,
-            controlFlags: 0
-          };
-          if (firstMessage) {
-            m.serviceName = serviceName;
-            m.procedureName = procedureName;
-            m.tracing = tracing;
-            m.controlFlags |= 2 /* StreamOpenBit */;
-            firstMessage = false;
-          }
-          transport.send(serverId, m);
-        }
-        if (!healthyClose)
-          return;
-        transport.sendCloseStream(serverId, streamId);
-        span.setStatus({ code: SpanStatusCode.OK });
+  const { span, ctx } = createProcTelemetryInfo(
+    "stream",
+    serviceName,
+    procedureName,
+    streamId
+  );
+  const inputStream = pushable({ objectMode: true });
+  const outputStream = pushable({ objectMode: true });
+  let firstMessage = true;
+  let healthyClose = true;
+  if (init) {
+    transport.send(serverId, {
+      streamId,
+      serviceName,
+      procedureName,
+      payload: init,
+      tracing: getPropagationContext(ctx),
+      controlFlags: 2 /* StreamOpenBit */
+    });
+    firstMessage = false;
+  }
+  const pipeInputToTransport = async () => {
+    for await (const rawIn of inputStream) {
+      const m = {
+        streamId,
+        payload: rawIn,
+        controlFlags: 0
       };
-      void pipeInputToTransport();
-      function onMessage(msg) {
-        if (msg.streamId !== streamId)
-          return;
-        if (msg.to !== transport.clientId)
-          return;
-        if (isStreamClose(msg.controlFlags)) {
-          cleanup();
-        } else {
-          outputStream.push(msg.payload);
-        }
-      }
-      function cleanup() {
-        inputStream.end();
-        outputStream.end();
-        transport.removeEventListener("message", onMessage);
-        transport.removeEventListener("sessionStatus", onSessionStatus);
-        span.end();
+      if (firstMessage) {
+        m.serviceName = serviceName;
+        m.procedureName = procedureName;
+        m.tracing = getPropagationContext(ctx);
+        m.controlFlags |= 2 /* StreamOpenBit */;
+        firstMessage = false;
       }
-      const onSessionStatus = createSessionDisconnectHandler(serverId, () => {
-        outputStream.push(
-          Err({
-            code: UNEXPECTED_DISCONNECT,
-            message: `${serverId} unexpectedly disconnected`
-          })
-        );
-        healthyClose = false;
-        span.setStatus({ code: SpanStatusCode.ERROR });
-        cleanup();
-      });
-      transport.addEventListener("message", onMessage);
-      transport.addEventListener("sessionStatus", onSessionStatus);
-      return [inputStream, outputStream, cleanup];
+      transport.send(serverId, m);
     }
-  );
+    if (!healthyClose)
+      return;
+    transport.sendCloseStream(serverId, streamId);
+  };
+  void pipeInputToTransport();
+  function onMessage(msg) {
+    if (msg.streamId !== streamId)
+      return;
+    if (msg.to !== transport.clientId)
+      return;
+    if (isStreamClose(msg.controlFlags)) {
+      cleanup();
+    } else {
+      outputStream.push(msg.payload);
+    }
+  }
+  function cleanup() {
+    inputStream.end();
+    outputStream.end();
+    transport.removeEventListener("message", onMessage);
+    transport.removeEventListener("sessionStatus", onSessionStatus);
+    span.end();
+  }
+  const onSessionStatus = createSessionDisconnectHandler(serverId, () => {
+    outputStream.push(
+      Err({
+        code: UNEXPECTED_DISCONNECT,
+        message: `${serverId} unexpectedly disconnected`
+      })
+    );
+    healthyClose = false;
+    cleanup();
+  });
+  transport.addEventListener("message", onMessage);
+  transport.addEventListener("sessionStatus", onSessionStatus);
+  return [inputStream, outputStream, cleanup];
 }
 function handleSubscribe(transport, serverId, input, serviceName, procedureName) {
   const streamId = nanoid();
-  return tracing_default.startActiveSpan(
-    `${serviceName}.${procedureName}`,
-    {
-      attributes: {
-        component: "river",
-        "river.method.kind": "subscribe",
-        "river.method.service": serviceName,
-        "river.method.name": procedureName,
-        "river.streamId": streamId,
-        "span.kind": "client"
-      },
-      kind: SpanKind.CLIENT
-    },
-    (span) => {
-      const tracing = { traceparent: "", tracestate: "" };
-      propagation.inject(context.active(), tracing);
-      transport.send(serverId, {
-        streamId,
-        serviceName,
-        procedureName,
-        tracing,
-        payload: input,
-        controlFlags: 2 /* StreamOpenBit */
-      });
-      let healthyClose = true;
-      const outputStream = pushable({ objectMode: true });
-      function onMessage(msg) {
-        if (msg.streamId !== streamId)
-          return;
-        if (msg.to !== transport.clientId)
-          return;
-        if (isStreamClose(msg.controlFlags)) {
-          cleanup();
-        } else {
-          outputStream.push(msg.payload);
-        }
-      }
-      function cleanup() {
-        outputStream.end();
-        transport.removeEventListener("message", onMessage);
-        transport.removeEventListener("sessionStatus", onSessionStatus);
-        span.end();
-      }
-      const closeHandler = () => {
-        cleanup();
-        if (!healthyClose)
-          return;
-        transport.sendCloseStream(serverId, streamId);
-      };
-      const onSessionStatus = createSessionDisconnectHandler(serverId, () => {
-        outputStream.push(
-          Err({
-            code: UNEXPECTED_DISCONNECT,
-            message: `${serverId} unexpectedly disconnected`
-          })
-        );
-        healthyClose = false;
-        span.setStatus({ code: SpanStatusCode.ERROR });
-        cleanup();
-      });
-      transport.addEventListener("message", onMessage);
-      transport.addEventListener("sessionStatus", onSessionStatus);
-      return [outputStream, closeHandler];
-    }
+  const { span, ctx } = createProcTelemetryInfo(
+    "subscription",
+    serviceName,
+    procedureName,
+    streamId
   );
+  transport.send(serverId, {
+    streamId,
+    serviceName,
+    procedureName,
+    payload: input,
+    tracing: getPropagationContext(ctx),
+    controlFlags: 2 /* StreamOpenBit */
+  });
+  let healthyClose = true;
+  const outputStream = pushable({ objectMode: true });
+  function onMessage(msg) {
+    if (msg.streamId !== streamId)
+      return;
+    if (msg.to !== transport.clientId)
+      return;
+    if (isStreamClose(msg.controlFlags)) {
+      cleanup();
+    } else {
+      outputStream.push(msg.payload);
+    }
+  }
+  function cleanup() {
+    outputStream.end();
+    transport.removeEventListener("message", onMessage);
+    transport.removeEventListener("sessionStatus", onSessionStatus);
+    span.end();
+  }
+  const closeHandler = () => {
+    cleanup();
+    if (!healthyClose)
+      return;
+    transport.sendCloseStream(serverId, streamId);
+  };
+  const onSessionStatus = createSessionDisconnectHandler(serverId, () => {
+    outputStream.push(
+      Err({
+        code: UNEXPECTED_DISCONNECT,
+        message: `${serverId} unexpectedly disconnected`
+      })
+    );
+    healthyClose = false;
+    cleanup();
+  });
+  transport.addEventListener("message", onMessage);
+  transport.addEventListener("sessionStatus", onSessionStatus);
+  return [outputStream, closeHandler];
 }
 function handleUpload(transport, serverId, init, serviceName, procedureName) {
   const streamId = nanoid();
-  return tracing_default.startActiveSpan(
-    `${serviceName}.${procedureName}`,
-    {
-      attributes: {
-        component: "river",
-        "river.method.kind": init ? "upload-with-init" : "upload",
-        "river.method.service": serviceName,
-        "river.method.name": procedureName,
-        "river.streamId": streamId,
-        "span.kind": "client"
-      },
-      kind: SpanKind.CLIENT
-    },
-    (span) => {
-      const tracing = { traceparent: "", tracestate: "" };
-      propagation.inject(context.active(), tracing);
-      const inputStream = pushable({ objectMode: true });
-      let firstMessage = true;
-      let healthyClose = true;
-      if (init) {
-        transport.send(serverId, {
-          streamId,
-          serviceName,
-          procedureName,
-          tracing,
-          payload: init,
-          controlFlags: 2 /* StreamOpenBit */
-        });
+  const { span, ctx } = createProcTelemetryInfo(
+    "upload",
+    serviceName,
+    procedureName,
+    streamId
+  );
+  const inputStream = pushable({ objectMode: true });
+  let firstMessage = true;
+  let healthyClose = true;
+  if (init) {
+    transport.send(serverId, {
+      streamId,
+      serviceName,
+      procedureName,
+      payload: init,
+      tracing: getPropagationContext(ctx),
+      controlFlags: 2 /* StreamOpenBit */
+    });
+    firstMessage = false;
+  }
+  const pipeInputToTransport = async () => {
+    for await (const rawIn of inputStream) {
+      const m = {
+        streamId,
+        payload: rawIn,
+        controlFlags: 0
+      };
+      if (firstMessage) {
+        m.serviceName = serviceName;
+        m.procedureName = procedureName;
+        m.tracing = getPropagationContext(ctx);
+        m.controlFlags |= 2 /* StreamOpenBit */;
         firstMessage = false;
       }
-      const pipeInputToTransport = async () => {
-        for await (const rawIn of inputStream) {
-          const m = {
-            streamId,
-            payload: rawIn,
-            controlFlags: 0
-          };
-          if (firstMessage) {
-            m.serviceName = serviceName;
-            m.procedureName = procedureName;
-            m.tracing = tracing;
-            m.controlFlags |= 2 /* StreamOpenBit */;
-            firstMessage = false;
-          }
-          transport.send(serverId, m);
-        }
-        if (!healthyClose)
-          return;
-        transport.sendCloseStream(serverId, streamId);
-      };
-      void pipeInputToTransport();
-      const responsePromise = new Promise((resolve) => {
-        const onSessionStatus = createSessionDisconnectHandler(serverId, () => {
-          healthyClose = false;
-          span.setStatus({ code: SpanStatusCode.ERROR });
-          cleanup();
-          resolve(
-            Err({
-              code: UNEXPECTED_DISCONNECT,
-              message: `${serverId} unexpectedly disconnected`
-            })
-          );
-        });
-        function cleanup() {
-          inputStream.end();
-          transport.removeEventListener("message", onMessage);
-          transport.removeEventListener("sessionStatus", onSessionStatus);
-          span.end();
-        }
-        function onMessage(msg) {
-          if (msg.streamId !== streamId)
-            return;
-          if (msg.to !== transport.clientId)
-            return;
-          if (msg.payload && typeof msg.payload === "object" && "ok" in msg.payload) {
-            span.setStatus({
-              code: msg.payload.ok ? SpanStatusCode.OK : SpanStatusCode.ERROR
-            });
-          }
-          cleanup();
-          resolve(msg.payload);
-        }
-        transport.addEventListener("message", onMessage);
-        transport.addEventListener("sessionStatus", onSessionStatus);
-      });
-      return [inputStream, responsePromise];
+      transport.send(serverId, m);
     }
-  );
+    if (!healthyClose)
+      return;
+    transport.sendCloseStream(serverId, streamId);
+  };
+  void pipeInputToTransport();
+  const responsePromise = new Promise((resolve) => {
+    const onSessionStatus = createSessionDisconnectHandler(serverId, () => {
+      healthyClose = false;
+      cleanup();
+      resolve(
+        Err({
+          code: UNEXPECTED_DISCONNECT,
+          message: `${serverId} unexpectedly disconnected`
+        })
+      );
+    });
+    function cleanup() {
+      inputStream.end();
+      transport.removeEventListener("message", onMessage);
+      transport.removeEventListener("sessionStatus", onSessionStatus);
+      span.end();
+    }
+    function onMessage(msg) {
+      if (msg.streamId !== streamId)
+        return;
+      if (msg.to !== transport.clientId)
+        return;
+      cleanup();
+      resolve(msg.payload);
+    }
+    transport.addEventListener("message", onMessage);
+    transport.addEventListener("sessionStatus", onSessionStatus);
+  });
+  return [inputStream, responsePromise];
 }
 
 // router/server.ts
-import {
-  context as context2,
-  propagation as propagation2,
-  SpanKind as SpanKind2,
-  SpanStatusCode as SpanStatusCode2
-} from "@opentelemetry/api";
 import { Value } from "@sinclair/typebox/value";
+import { SpanStatusCode } from "@opentelemetry/api";
 var RiverServer = class {
   transport;
   services;
@@ -1015,7 +948,7 @@ var RiverServer = class {
   // map of client to their open streams by streamId
   clientStreams;
   disconnectedSessions;
-  constructor(transport, services, extendedContext) {
+  constructor(transport, services, handshakeOptions, extendedContext) {
     const instances = {};
     this.services = instances;
     this.contextMap = /* @__PURE__ */ new Map();
@@ -1027,6 +960,9 @@ var RiverServer = class {
         state: instance.state
       });
     }
+    if (handshakeOptions) {
+      transport.extendHandshake(handshakeOptions);
+    }
     this.transport = transport;
     this.disconnectedSessions = /* @__PURE__ */ new Set();
     this.streamMap = /* @__PURE__ */ new Map();
@@ -1076,9 +1012,9 @@ var RiverServer = class {
     this.transport.removeEventListener("message", this.onMessage);
     this.transport.removeEventListener("sessionStatus", this.onSessionStatus);
     await Promise.all([...this.streamMap.keys()].map(this.cleanupStream));
-    for (const context3 of this.contextMap.values()) {
-      if (Symbol.dispose in context3.state) {
-        const dispose = context3.state[Symbol.dispose];
+    for (const context of this.contextMap.values()) {
+      if (Symbol.dispose in context.state) {
+        const dispose = context.state[Symbol.dispose];
         if (typeof dispose === "function") {
           dispose();
         }
@@ -1111,10 +1047,6 @@ var RiverServer = class {
       });
       return;
     }
-    let activeContext = context2.active();
-    if (message.tracing) {
-      activeContext = propagation2.extract(activeContext, message.tracing);
-    }
     const service = this.services[message.serviceName];
     const serviceContext = this.getContext(service, message.serviceName);
     if (!(message.procedureName in service.procedures)) {
@@ -1178,12 +1110,8 @@ var RiverServer = class {
         `procedure ${message.serviceName}.${message.procedureName} threw an uncaught error: ${errorMsg}`,
         session.loggingMetadata
       );
-      if (err instanceof Error) {
-        span.recordException(err);
-      } else {
-        span.recordException(errorMsg);
-      }
-      span.setStatus({ code: SpanStatusCode2.ERROR });
+      span.recordException(err instanceof Error ? err : new Error(errorMsg));
+      span.setStatus({ code: SpanStatusCode.ERROR });
       outgoing.push(
         Err({
           code: UNCAUGHT_ERROR,
@@ -1191,7 +1119,8 @@ var RiverServer = class {
         })
       );
     };
-    if (session.metadata === void 0) {
+    const sessionMeta = this.transport.sessionHandshakeMetadata.get(session);
+    if (!sessionMeta) {
       log?.error(`session doesn't have handshake metadata`, {
         ...session.loggingMetadata,
         tags: ["invariant-violation"]
@@ -1205,37 +1134,53 @@ var RiverServer = class {
       to: message.to,
       from: message.from,
       streamId: message.streamId,
-      // we've already validated that the session has handshake metadata
-      session
+      session,
+      metadata: sessionMeta
     };
     switch (procedure.type) {
       case "rpc":
-        inputHandler = (async () => {
-          const inputMessage = await incoming.next();
-          if (inputMessage.done) {
-            return;
+        inputHandler = createHandlerSpan(
+          procedure.type,
+          message,
+          async (span) => {
+            const inputMessage = await incoming.next();
+            if (inputMessage.done) {
+              return;
+            }
+            try {
+              const outputMessage = await procedure.handler(
+                serviceContextWithTransportInfo,
+                inputMessage.value
+              );
+              outgoing.push(outputMessage);
+            } catch (err) {
+              errorHandler(err, span);
+            } finally {
+              span.end();
+            }
           }
-          await tracing_default.startActiveSpan(
-            `${message.serviceName}.${message.procedureName}`,
-            {
-              attributes: {
-                component: "river",
-                "river.method.kind": "rpc",
-                "river.method.service": message.serviceName,
-                "river.method.name": message.procedureName,
-                "river.streamId": message.streamId,
-                "span.kind": "server"
-              },
-              kind: SpanKind2.SERVER
-            },
-            activeContext,
+        );
+        break;
+      case "stream":
+        if (procHasInitMessage) {
+          inputHandler = createHandlerSpan(
+            procedure.type,
+            message,
             async (span) => {
+              const initMessage = await incoming.next();
+              if (initMessage.done) {
+                return;
+              }
               try {
-                const outputMessage = await procedure.handler(
+                const dispose = await procedure.handler(
                   serviceContextWithTransportInfo,
-                  inputMessage.value
+                  initMessage.value,
+                  incoming,
+                  outgoing
                 );
-                outgoing.push(outputMessage);
+                if (dispose) {
+                  disposables.push(dispose);
+                }
               } catch (err) {
                 errorHandler(err, span);
               } finally {
@@ -1243,67 +1188,10 @@ var RiverServer = class {
               }
             }
           );
-        })();
-        break;
-      case "stream":
-        if (procHasInitMessage) {
-          inputHandler = (async () => {
-            const initMessage = await incoming.next();
-            if (initMessage.done) {
-              return;
-            }
-            await tracing_default.startActiveSpan(
-              `${message.serviceName}.${message.procedureName}`,
-              {
-                attributes: {
-                  component: "river",
-                  "river.method.kind": "stream-with-init",
-                  "river.method.service": message.serviceName,
-                  "river.method.name": message.procedureName,
-                  "river.streamId": message.streamId,
-                  "span.kind": "server"
-                },
-                kind: SpanKind2.SERVER
-              },
-              activeContext,
-              async (span) => {
-                try {
-                  const dispose = await procedure.handler(
-                    serviceContextWithTransportInfo,
-                    initMessage.value,
-                    incoming,
-                    outgoing
-                  );
-                  if (dispose) {
-                    disposables.push(() => {
-                      dispose();
-                      span.end();
-                    });
-                  } else {
-                    span.end();
-                  }
-                } catch (err) {
-                  errorHandler(err, span);
-                  span.end();
-                }
-              }
-            );
-          })();
         } else {
-          inputHandler = tracing_default.startActiveSpan(
-            `${message.serviceName}.${message.procedureName}`,
-            {
-              attributes: {
-                component: "river",
-                "river.method.kind": "stream",
-                "river.method.service": message.serviceName,
-                "river.method.name": message.procedureName,
-                "river.streamId": message.streamId,
-                "span.kind": "server"
-              },
-              kind: SpanKind2.SERVER
-            },
-            activeContext,
+          inputHandler = createHandlerSpan(
+            procedure.type,
+            message,
             async (span) => {
               try {
                 const dispose = await procedure.handler(
@@ -1312,15 +1200,11 @@ var RiverServer = class {
                   outgoing
                 );
                 if (dispose) {
-                  disposables.push(() => {
-                    dispose();
-                    span.end();
-                  });
-                } else {
-                  span.end();
+                  disposables.push(dispose);
                 }
               } catch (err) {
                 errorHandler(err, span);
+              } finally {
                 span.end();
               }
             }
@@ -1328,102 +1212,61 @@ var RiverServer = class {
         }
         break;
       case "subscription":
-        inputHandler = (async () => {
-          const inputMessage = await incoming.next();
-          if (inputMessage.done) {
-            return;
+        inputHandler = createHandlerSpan(
+          procedure.type,
+          message,
+          async (span) => {
+            const inputMessage = await incoming.next();
+            if (inputMessage.done) {
+              return;
+            }
+            try {
+              const dispose = await procedure.handler(
+                serviceContextWithTransportInfo,
+                inputMessage.value,
+                outgoing
+              );
+              if (dispose) {
+                disposables.push(dispose);
+              }
+            } catch (err) {
+              errorHandler(err, span);
+            } finally {
+              span.end();
+            }
           }
-          await tracing_default.startActiveSpan(
-            `${message.serviceName}.${message.procedureName}`,
-            {
-              attributes: {
-                component: "river",
-                "river.method.kind": "subscription",
-                "river.method.service": message.serviceName,
-                "river.method.name": message.procedureName,
-                "river.streamId": message.streamId,
-                "span.kind": "server"
-              },
-              kind: SpanKind2.SERVER
-            },
-            activeContext,
+        );
+        break;
+      case "upload":
+        if (procHasInitMessage) {
+          inputHandler = createHandlerSpan(
+            procedure.type,
+            message,
             async (span) => {
+              const initMessage = await incoming.next();
+              if (initMessage.done) {
+                return;
+              }
               try {
-                const dispose = await procedure.handler(
+                const outputMessage = await procedure.handler(
                   serviceContextWithTransportInfo,
-                  inputMessage.value,
-                  outgoing
+                  initMessage.value,
+                  incoming
                 );
-                if (dispose) {
-                  disposables.push(() => {
-                    dispose();
-                    span.end();
-                  });
-                } else {
-                  span.end();
+                if (!this.disconnectedSessions.has(message.from)) {
+                  outgoing.push(outputMessage);
                 }
               } catch (err) {
                 errorHandler(err, span);
+              } finally {
                 span.end();
               }
             }
           );
-        })();
-        break;
-      case "upload":
-        if (procHasInitMessage) {
-          inputHandler = (async () => {
-            const initMessage = await incoming.next();
-            if (initMessage.done) {
-              return;
-            }
-            await tracing_default.startActiveSpan(
-              `${message.serviceName}.${message.procedureName}`,
-              {
-                attributes: {
-                  component: "river",
-                  "river.method.kind": "upload-with-init",
-                  "river.method.service": message.serviceName,
-                  "river.method.name": message.procedureName,
-                  "river.streamId": message.streamId,
-                  "span.kind": "server"
-                },
-                kind: SpanKind2.SERVER
-              },
-              activeContext,
-              async (span) => {
-                try {
-                  const outputMessage = await procedure.handler(
-                    serviceContextWithTransportInfo,
-                    initMessage.value,
-                    incoming
-                  );
-                  if (!this.disconnectedSessions.has(message.from)) {
-                    outgoing.push(outputMessage);
-                  }
-                } catch (err) {
-                  errorHandler(err, span);
-                } finally {
-                  span.end();
-                }
-              }
-            );
-          })();
         } else {
-          inputHandler = tracing_default.startActiveSpan(
-            `${message.serviceName}.${message.procedureName}`,
-            {
-              attributes: {
-                component: "river",
-                "river.method.kind": "upload",
-                "river.method.service": message.serviceName,
-                "river.method.name": message.procedureName,
-                "river.streamId": message.streamId,
-                "span.kind": "server"
-              },
-              kind: SpanKind2.SERVER
-            },
-            activeContext,
+          inputHandler = createHandlerSpan(
+            procedure.type,
+            message,
             async (span) => {
               try {
                 const outputMessage = await procedure.handler(
@@ -1489,8 +1332,8 @@ var RiverServer = class {
     }
   }
   getContext(service, serviceName) {
-    const context3 = this.contextMap.get(service);
-    if (!context3) {
+    const context = this.contextMap.get(service);
+    if (!context) {
       const err = `no context found for ${serviceName}`;
       log?.error(err, {
         clientId: this.transport.clientId,
@@ -1498,7 +1341,7 @@ var RiverServer = class {
       });
       throw new Error(err);
     }
-    return context3;
+    return context;
   }
   cleanupStream = async (id) => {
     const stream2 = this.streamMap.get(id);
@@ -1512,13 +1355,15 @@ var RiverServer = class {
     this.streamMap.delete(id);
   };
 };
-function createServer(transport, services, extendedContext) {
-  return new RiverServer(transport, services, extendedContext);
+function createServer(transport, services, providedServerOptions) {
+  return new RiverServer(
+    transport,
+    services,
+    providedServerOptions?.handshakeOptions,
+    providedServerOptions?.extendedContext
+  );
 }
 
-// package.json
-var version = "0.21.1";
-
 export {
   serializeSchema,
   ServiceSchema,
@@ -1529,7 +1374,6 @@ export {
   Ok,
   Err,
   createClient,
-  createServer,
-  version
+  createServer
 };
-//# sourceMappingURL=chunk-NCXUFDVL.js.map
+//# sourceMappingURL=chunk-VOJVLWVX.js.map