export-runtime 0.0.20 → 0.0.22

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,37 +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;
 
-        // Don't send if connection closed during dispatch
-        if (isClosed) return;
+      // 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;
+          }
 
-        // Extract token from auth responses
-        if (result?.value?.token && msg.type === "auth") {
-          wsSession.token = result.value.token;
-        }
+          if (isClosed) return;
+          const result = await dispatchMessage(dispatcher, msg, env, wsSession);
 
-        if (result) safeSend(stringify({ ...result, id }));
-      } catch (err) {
-        // Don't send errors if connection is closed or error is about closed connection
-        if (isClosed || String(err).includes("Connection closed")) return;
-        if (id !== undefined) safeSend(stringify({ type: "error", id, error: String(err) }));
-      }
+          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", () => {
@@ -372,13 +387,15 @@ export const createHandler = (moduleMap, generatedTypes, minifiedCore, coreId, m
       if (onClose) onClose();
     });
 
-    server.addEventListener("error", () => {
+    server.addEventListener("error", (err) => {
       isClosed = true;
+      if (onClose) onClose();
+      try { server.close(); } catch {}
     });
   };
 
   return {
-    async fetch(request, env) {
+    async fetch(request, env, ctx) {
       const url = new URL(request.url);
       const isShared = url.searchParams.has("shared");
       const origin = request.headers.get("Origin");
@@ -401,6 +418,13 @@ export const createHandler = (moduleMap, generatedTypes, minifiedCore, coreId, m
 
         const pair = new WebSocketPair();
         const [client, server] = Object.values(pair);
+
+        // Create a promise that resolves when the WebSocket closes
+        const wsLifetime = new Promise((resolve) => {
+          server.addEventListener("close", resolve, { once: true });
+          server.addEventListener("error", resolve, { once: true });
+        });
+
         server.accept();
 
         if (isShared && env?.SHARED_EXPORT) {
@@ -413,6 +437,11 @@ export const createHandler = (moduleMap, generatedTypes, minifiedCore, coreId, m
           wireWebSocket(server, dispatcher, env, () => dispatcher.clearAll());
         }
 
+        // Keep worker alive for the WebSocket lifetime
+        if (ctx?.waitUntil) {
+          ctx.waitUntil(wsLifetime);
+        }
+
         return new Response(null, { status: 101, webSocket: client });
       }
 

package/package.json

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

package/rpc.js

@@ -155,25 +155,20 @@ 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 {}
-      }
+      // Synchronously clear all maps to prevent further access
+      const iters = [...iterators.values()];
+      const strms = [...streams.values()];
       instances.clear();
       iterators.clear();
       streams.clear();
+
+      // Cancel iterators and streams without awaiting (fire and forget)
+      for (const iter of iters) {
+        try { iter?.return?.(); } catch {}
+      }
+      for (const entry of strms) {
+        try { (entry.reader || entry.stream)?.cancel?.(); } catch {}
+      }
     },
   };
 }