diff --git a/rivetkit-typescript/packages/rivetkit/src/actor/router-endpoints.ts b/rivetkit-typescript/packages/rivetkit/src/actor/router-endpoints.ts index 37aa8c0801..6ea33604b8 100644 --- a/rivetkit-typescript/packages/rivetkit/src/actor/router-endpoints.ts +++ b/rivetkit-typescript/packages/rivetkit/src/actor/router-endpoints.ts @@ -163,6 +163,9 @@ export async function handleWebSocketConnect( // Promise used to wait for the websocket close in `disconnect` const closePromise = promiseWithResolvers(); + // Track connection outside of scope for cleanup + let createdConn: AnyConn | undefined; + return { onOpen: (_evt: any, ws: WSContext) => { actor.rLog.debug("actor websocket open"); @@ -198,6 +201,9 @@ export async function handleWebSocketConnect( connToken, ); + // Store connection so we can clean on close + createdConn = conn; + // Unblock other handlers handlersResolve({ conn, actor, connId: conn.id }); } catch (error) { @@ -298,20 +304,13 @@ export async function handleWebSocketConnect( // https://github.com/cloudflare/workerd/issues/2569 ws.close(1000, "hack_force_close"); - // Handle cleanup asynchronously - handlersPromise - .then(({ conn, actor }) => { + // Wait for actor.createConn to finish before removing the connection + handlersPromise.finally(() => { + if (createdConn) { const wasClean = event.wasClean || event.code === 1000; - actor.__connDisconnected(conn, wasClean, requestId); - }) - .catch((error) => { - deconstructError( - error, - actor.rLog, - { wsEvent: "close" }, - exposeInternalError, - ); - }); + actor.__connDisconnected(createdConn, wasClean, requestId); + } + }); }, onError: (_error: unknown) => { try {