export-runtime 0.0.19 → 0.0.21

package/handler.js

@@ -323,6 +323,7 @@ export const createHandler = (moduleMap, generatedTypes, minifiedCore, coreId, m
     // Track session state and connection state for this WebSocket connection
     const wsSession = { token: null };
     let isClosed = false;
+    const pendingOps = new Set();
 
     // Safe send that ignores errors when connection is closed
     const safeSend = (data) => {
@@ -334,32 +335,51 @@ export const createHandler = (moduleMap, generatedTypes, minifiedCore, coreId, m
       }
     };
 
-    server.addEventListener("message", async (event) => {
-      if (isClosed) return;
-      let id;
+    // Wrap async operations to track them
+    const trackOp = async (op) => {
+      const promise = op();
+      pendingOps.add(promise);
       try {
-        const msg = parse(event.data);
-        id = msg.id;
-
-        // Handle auth token updates (on reconnect or explicit setToken)
-        if (msg.type === "auth" && msg.token && !msg.method) {
-          // Direct token send on reconnect - just update session
-          wsSession.token = msg.token;
-          safeSend(stringify({ type: "auth-result", id, success: true }));
-          return;
-        }
+        return await promise;
+      } finally {
+        pendingOps.delete(promise);
+      }
+    };
 
-        const result = await dispatchMessage(dispatcher, msg, env, wsSession);
+    server.addEventListener("message", (event) => {
+      if (isClosed) return;
 
-        // Extract token from auth responses
-        if (result?.value?.token && msg.type === "auth") {
-          wsSession.token = result.value.token;
-        }
+      // Handle message asynchronously but don't await it
+      trackOp(async () => {
+        let id;
+        try {
+          if (isClosed) return;
+          const msg = parse(event.data);
+          id = msg.id;
+
+          // Handle auth token updates (on reconnect or explicit setToken)
+          if (msg.type === "auth" && msg.token && !msg.method) {
+            wsSession.token = msg.token;
+            safeSend(stringify({ type: "auth-result", id, success: true }));
+            return;
+          }
 
-        if (result) safeSend(stringify({ ...result, id }));
-      } catch (err) {
-        if (id !== undefined) safeSend(stringify({ type: "error", id, error: String(err) }));
-      }
+          if (isClosed) return;
+          const result = await dispatchMessage(dispatcher, msg, env, wsSession);
+
+          if (isClosed) return;
+
+          // Extract token from auth responses
+          if (result?.value?.token && msg.type === "auth") {
+            wsSession.token = result.value.token;
+          }
+
+          if (result) safeSend(stringify({ ...result, id }));
+        } catch (err) {
+          if (isClosed || String(err).includes("Connection closed")) return;
+          if (id !== undefined) safeSend(stringify({ type: "error", id, error: String(err) }));
+        }
+      });
     });
 
     server.addEventListener("close", () => {
@@ -369,6 +389,7 @@ export const createHandler = (moduleMap, generatedTypes, minifiedCore, coreId, m
 
     server.addEventListener("error", () => {
       isClosed = true;
+      if (onClose) onClose();
     });
   };
 

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "export-runtime",
-  "version": "0.0.19",
+  "version": "0.0.21",
   "description": "Cloudflare Workers ESM Export Framework Runtime",
   "keywords": [
     "cloudflare",

package/rpc.js

@@ -38,6 +38,7 @@ export function createRpcDispatcher(moduleMap) {
   const iterators = new Map();
   const streams = new Map();
   let nextId = 1;
+  let closed = false;
 
   const requireInstance = (id) => {
     const inst = instances.get(id);
@@ -115,9 +116,11 @@ export function createRpcDispatcher(moduleMap) {
     },
 
     async rpcIterateNext(iteratorId) {
+      if (closed) throw new Error("Connection closed");
       const iter = iterators.get(iteratorId);
       if (!iter) throw new Error("Iterator not found");
       const { value, done } = await iter.next();
+      if (closed) throw new Error("Connection closed");
       if (done) iterators.delete(iteratorId);
       return { type: "iterate-result", value, done: !!done };
     },
@@ -130,10 +133,12 @@ export function createRpcDispatcher(moduleMap) {
     },
 
     async rpcStreamRead(streamId) {
+      if (closed) throw new Error("Connection closed");
       const entry = streams.get(streamId);
       if (!entry) throw new Error("Stream not found");
       if (!entry.reader) entry.reader = entry.stream.getReader();
       const { value, done } = await entry.reader.read();
+      if (closed) throw new Error("Connection closed");
       if (done) streams.delete(streamId);
       const v = value instanceof Uint8Array ? Array.from(value) : value;
       return { type: "stream-result", value: v, done: !!done };
@@ -149,6 +154,23 @@ export function createRpcDispatcher(moduleMap) {
     },
 
     clearAll() {
+      closed = true;
+      // Cancel all active iterators
+      for (const iter of iterators.values()) {
+        try {
+          if (iter?.return) iter.return(undefined);
+        } catch {}
+      }
+      // Cancel all active streams
+      for (const entry of streams.values()) {
+        try {
+          if (entry.reader) {
+            entry.reader.cancel();
+          } else if (entry.stream) {
+            entry.stream.cancel();
+          }
+        } catch {}
+      }
       instances.clear();
       iterators.clear();
       streams.clear();