rivetkit 2.0.6 → 2.0.7

package/dist/tsup/{chunk-4CKHQRXG.js → chunk-NII4KKHD.js}

@@ -3,11 +3,11 @@ import {
   ForeignKeysSchema,
   PatchSchema,
   TablesSchema
-} from "./chunk-ZYLTS2EM.js";
+} from "./chunk-54MAHBLL.js";
 import {
   importWebSocket,
   logger
-} from "./chunk-VCEHU56K.js";
+} from "./chunk-B6N6VM37.js";
 import {
   HTTP_ACTION_REQUEST_VERSIONED,
   HTTP_ACTION_RESPONSE_VERSIONED,
@@ -17,31 +17,43 @@ import {
   TO_SERVER_VERSIONED,
   inputDataToBuffer,
   processMessage
-} from "./chunk-7OUKNSTU.js";
+} from "./chunk-2NL3KGJ7.js";
 import {
   CachedSerializer,
   DeadlineError,
-  HEADER_AUTH_DATA,
   HEADER_CONN_ID,
   HEADER_CONN_PARAMS,
   HEADER_CONN_TOKEN,
   HEADER_ENCODING,
+  HEADER_RIVET_ACTOR,
+  HEADER_RIVET_TARGET,
+  HEADER_RIVET_TOKEN,
   PATH_CONNECT_WEBSOCKET,
   PATH_RAW_WEBSOCKET_PREFIX,
+  WS_PROTOCOL_ACTOR,
+  WS_PROTOCOL_CONN_ID,
+  WS_PROTOCOL_CONN_PARAMS,
+  WS_PROTOCOL_CONN_TOKEN,
+  WS_PROTOCOL_ENCODING,
+  WS_PROTOCOL_STANDARD,
+  WS_PROTOCOL_TARGET,
+  WS_PROTOCOL_TOKEN,
   assertUnreachable as assertUnreachable2,
   contentTypeForEncoding,
   deadline,
   deserializeWithEncoding,
+  encodeDataToString,
   encodingIsBinary,
+  generateRandomString,
   generateSecureToken,
   jsonStringifyCompat,
   serializeWithEncoding
-} from "./chunk-2K3JMDAN.js";
+} from "./chunk-NRELKXIX.js";
 import {
   getBaseLogger,
   getIncludeTarget,
   getLogger
-} from "./chunk-WC2PSJWN.js";
+} from "./chunk-YL4VZMMT.js";
 import {
   SinglePromiseQueue,
   assertUnreachable,
@@ -52,8 +64,9 @@ import {
   httpUserAgent,
   isCborSerializable,
   noopNext,
+  promiseWithResolvers,
   stringifyError
-} from "./chunk-RGQR2J7S.js";
+} from "./chunk-RLBM6D4L.js";
 import {
   ActionNotFound,
   ActionTimedOut,
@@ -69,44 +82,134 @@ import {
   VarsNotEnabled
 } from "./chunk-YPZFLUO6.js";
 
-// src/actor/connection.ts
+// src/actor/conn.ts
 import * as cbor from "cbor-x";
+
+// src/actor/conn-drivers.ts
+var WEBSOCKET_DRIVER = {
+  sendMessage: (actor, _conn, state, message) => {
+    const serialized = message.serialize(state.encoding);
+    actor.rLog.debug({
+      msg: "sending websocket message",
+      encoding: state.encoding,
+      dataType: typeof serialized,
+      isUint8Array: serialized instanceof Uint8Array,
+      isArrayBuffer: serialized instanceof ArrayBuffer,
+      dataLength: serialized.byteLength || serialized.length
+    });
+    if (serialized instanceof Uint8Array) {
+      const buffer = serialized.buffer.slice(
+        serialized.byteOffset,
+        serialized.byteOffset + serialized.byteLength
+      );
+      if (buffer instanceof SharedArrayBuffer) {
+        const arrayBuffer = new ArrayBuffer(buffer.byteLength);
+        new Uint8Array(arrayBuffer).set(new Uint8Array(buffer));
+        actor.rLog.debug({
+          msg: "converted SharedArrayBuffer to ArrayBuffer",
+          byteLength: arrayBuffer.byteLength
+        });
+        state.websocket.send(arrayBuffer);
+      } else {
+        actor.rLog.debug({
+          msg: "sending ArrayBuffer",
+          byteLength: buffer.byteLength
+        });
+        state.websocket.send(buffer);
+      }
+    } else {
+      actor.rLog.debug({
+        msg: "sending string data",
+        length: serialized.length
+      });
+      state.websocket.send(serialized);
+    }
+  },
+  disconnect: async (_actor, _conn, state, reason) => {
+    state.websocket.close(1e3, reason);
+    await state.closePromise.promise;
+  },
+  getConnectionReadyState: (_actor, _conn, state) => {
+    return state.websocket.readyState;
+  }
+};
+var SSE_DRIVER = {
+  sendMessage: (_actor, _conn, state, message) => {
+    state.stream.writeSSE({
+      data: encodeDataToString(message.serialize(state.encoding))
+    });
+  },
+  disconnect: async (_actor, _conn, state, _reason) => {
+    state.stream.close();
+  },
+  getConnectionReadyState: (_actor, _conn, state) => {
+    if (state.stream.aborted || state.stream.closed) {
+      return 3 /* CLOSED */;
+    }
+    return 1 /* OPEN */;
+  }
+};
+var HTTP_DRIVER = {
+  getConnectionReadyState(_actor, _conn) {
+    return 1 /* OPEN */;
+  },
+  disconnect: async () => {
+  }
+};
+var CONN_DRIVERS = {
+  [0 /* WEBSOCKET */]: WEBSOCKET_DRIVER,
+  [1 /* SSE */]: SSE_DRIVER,
+  [2 /* HTTP */]: HTTP_DRIVER
+};
+function getConnDriverKindFromState(state) {
+  if (0 /* WEBSOCKET */ in state) return 0 /* WEBSOCKET */;
+  else if (1 /* SSE */ in state) return 1 /* SSE */;
+  else if (2 /* HTTP */ in state) return 2 /* HTTP */;
+  else assertUnreachable(state);
+}
+
+// src/actor/conn.ts
 function generateConnId() {
   return crypto.randomUUID();
 }
 function generateConnToken() {
   return generateSecureToken(32);
 }
-var CONNECTION_DRIVER_WEBSOCKET = "webSocket";
-var CONNECTION_DRIVER_SSE = "sse";
-var CONNECTION_DRIVER_HTTP = "http";
-var CONNECTION_CHECK_LIVENESS_SYMBOL = Symbol("checkLiveness");
+function generateConnSocketId() {
+  return crypto.randomUUID();
+}
 var Conn = class {
   subscriptions = /* @__PURE__ */ new Set();
-  #stateEnabled;
   // TODO: Remove this cyclical reference
   #actor;
-  #status = "connected";
   /**
    * The proxied state that notifies of changes automatically.
    *
    * Any data that should be stored indefinitely should be held within this object.
    */
   __persist;
+  get __driverState() {
+    var _a;
+    return (_a = this.__socket) == null ? void 0 : _a.driverState;
+  }
   /**
-   * Driver used to manage realtime connection communication.
+   * Socket connected to this connection.
    *
-   * @protected
+   * If undefined, then nothing is connected to this.
    */
-  #driver;
+  __socket;
+  get __status() {
+    if (this.__socket) {
+      return "connected";
+    } else {
+      return "reconnecting";
+    }
+  }
   get params() {
     return this.__persist.params;
   }
-  get driver() {
-    return this.__persist.connDriver;
-  }
-  get _stateEnabled() {
-    return this.#stateEnabled;
+  get __stateEnabled() {
+    return this.#actor.connStateEnabled;
   }
   /**
    * Gets the current state of the connection.
@@ -143,7 +246,7 @@ var Conn = class {
    * Status of the connection.
    */
   get status() {
-    return this.#status;
+    return this.__status;
   }
   /**
    * Timestamp of the last time the connection was seen, i.e. the last time the connection was active and checked for liveness.
@@ -158,14 +261,12 @@ var Conn = class {
    *
    * @protected
    */
-  constructor(actor, persist, driver, stateEnabled) {
+  constructor(actor, persist) {
     this.#actor = actor;
     this.__persist = persist;
-    this.#driver = driver;
-    this.#stateEnabled = stateEnabled;
   }
   #validateStateEnabled() {
-    if (!this.#stateEnabled) {
+    if (!this.__stateEnabled) {
       throw new ConnStateNotEnabled();
     }
   }
@@ -177,14 +278,28 @@ var Conn = class {
    * @protected
    */
   _sendMessage(message) {
-    var _a, _b;
-    (_b = (_a = this.#driver).sendMessage) == null ? void 0 : _b.call(
-      _a,
-      this.#actor,
-      this,
-      this.__persist.connDriverState,
-      message
-    );
+    if (this.__driverState) {
+      const driverKind = getConnDriverKindFromState(this.__driverState);
+      const driver = CONN_DRIVERS[driverKind];
+      if (driver.sendMessage) {
+        driver.sendMessage(
+          this.#actor,
+          this,
+          this.__driverState[driverKind],
+          message
+        );
+      } else {
+        this.#actor.rLog.debug({
+          msg: "conn driver does not support sending messages",
+          conn: this.id
+        });
+      }
+    } else {
+      this.#actor.rLog.warn({
+        msg: "missing connection driver state for send message",
+        conn: this.id
+      });
+    }
   }
   /**
    * Sends an event with arguments to the client.
@@ -221,43 +336,30 @@ var Conn = class {
    * @param reason - The reason for disconnection.
    */
   async disconnect(reason) {
-    this.#status = "reconnecting";
-    await this.#driver.disconnect(
-      this.#actor,
-      this,
-      this.__persist.connDriverState,
-      reason
-    );
-  }
-  /**
-   * This method checks the connection's liveness by querying the driver for its ready state.
-   * If the connection is not closed, it updates the last liveness timestamp and returns `true`.
-   * Otherwise, it returns `false`.
-   * @internal
-   */
-  [CONNECTION_CHECK_LIVENESS_SYMBOL]() {
-    const readyState = this.#driver.getConnectionReadyState(this.#actor, this);
-    const isConnectionClosed = readyState === 3 /* CLOSED */ || readyState === 2 /* CLOSING */ || readyState === void 0;
-    const newLastSeen = Date.now();
-    const newStatus = isConnectionClosed ? "reconnecting" : "connected";
-    this.#actor.rLog.debug({
-      msg: "liveness probe for connection",
-      connId: this.id,
-      actorId: this.#actor.id,
-      readyState,
-      status: this.#status,
-      newStatus,
-      lastSeen: this.__persist.lastSeen,
-      currentTs: newLastSeen
-    });
-    if (!isConnectionClosed) {
-      this.__persist.lastSeen = newLastSeen;
+    if (this.__socket && this.__driverState) {
+      const driverKind = getConnDriverKindFromState(this.__driverState);
+      const driver = CONN_DRIVERS[driverKind];
+      if (driver.disconnect) {
+        driver.disconnect(
+          this.#actor,
+          this,
+          this.__driverState[driverKind],
+          reason
+        );
+      } else {
+        this.#actor.rLog.debug({
+          msg: "no disconnect handler for conn driver",
+          conn: this.id
+        });
+      }
+      this.#actor.__connDisconnected(this, true, this.__socket.socketId);
+    } else {
+      this.#actor.rLog.warn({
+        msg: "missing connection driver state for disconnect",
+        conn: this.id
+      });
     }
-    this.#status = newStatus;
-    return {
-      status: this.#status,
-      lastSeen: this.__persist.lastSeen
-    };
+    this.__socket = void 0;
   }
 };
 
@@ -335,7 +437,7 @@ function createActorInspectorRouter() {
             id: String(id++)
           });
         });
-        const { promise } = Promise.withResolvers();
+        const { promise } = promiseWithResolvers();
         return promise;
       },
       async () => {
@@ -360,7 +462,7 @@ function createActorInspectorRouter() {
             id: String(id++)
           });
         });
-        const { promise } = Promise.withResolvers();
+        const { promise } = promiseWithResolvers();
         return promise;
       },
       async () => {
@@ -386,7 +488,7 @@ function createActorInspectorRouter() {
             id: String(id++)
           });
         });
-        const { promise } = Promise.withResolvers();
+        const { promise } = promiseWithResolvers();
         return promise;
       },
       async () => {
@@ -703,7 +805,6 @@ var ActorInstance = class {
   #backgroundPromises = [];
   #abortController = new AbortController();
   #config;
-  #connectionDrivers;
   #actorDriver;
   #inlineClient;
   #actorId;
@@ -741,9 +842,9 @@ var ActorInstance = class {
       getConnections: async () => {
         return Array.from(this.#connections.entries()).map(([id, conn]) => ({
           id,
-          stateEnabled: conn._stateEnabled,
+          stateEnabled: conn.__stateEnabled,
           params: conn.params,
-          state: conn._stateEnabled ? conn.state : void 0
+          state: conn.__stateEnabled ? conn.state : void 0
         }));
       },
       setState: async (state) => {
@@ -776,7 +877,7 @@ var ActorInstance = class {
     this.#config = config;
     this.actorContext = new ActorContext(this);
   }
-  async start(connectionDrivers, actorDriver, inlineClient, actorId, name, key, region) {
+  async start(actorDriver, inlineClient, actorId, name, key, region) {
     var _a, _b;
     const logParams = {
       actor: name,
@@ -792,7 +893,6 @@ var ActorInstance = class {
         logParams
       )
     );
-    this.#connectionDrivers = connectionDrivers;
     this.#actorDriver = actorDriver;
     this.#inlineClient = inlineClient;
     this.#actorId = actorId;
@@ -971,7 +1071,7 @@ var ActorInstance = class {
       throw new StateNotEnabled();
     }
   }
-  get #connStateEnabled() {
+  get connStateEnabled() {
     return "createConnState" in this.#config || "connState" in this.#config;
   }
   get #varsEnabled() {
@@ -1110,13 +1210,7 @@ var ActorInstance = class {
       });
       this.#setPersist(persistData);
       for (const connPersist of this.#persist.connections) {
-        const driver = this.__getConnDriver(connPersist.connDriver);
-        const conn = new Conn(
-          this,
-          connPersist,
-          driver,
-          this.#connStateEnabled
-        );
+        const conn = new Conn(this, connPersist);
         this.#connections.set(conn.id, conn);
         for (const sub of connPersist.subscriptions) {
           this.#addSubscription(sub.eventName, conn, true);
@@ -1159,13 +1253,36 @@ var ActorInstance = class {
     return this.#connections.get(id);
   }
   /**
-   * Removes a connection and cleans up its resources.
+   * Connection disconnected.
+   *
+   * If a clean diconnect, will be removed immediately.
+   *
+   * If not a clean disconnect, will keep the connection alive for a given interval to wait for reconnect.
    */
-  __removeConn(conn) {
-    if (!conn) {
-      this.#rLog.warn({ msg: "`conn` does not exist" });
+  __connDisconnected(conn, wasClean, socketId) {
+    if (socketId && conn.__socket && socketId !== conn.__socket.socketId) {
+      this.rLog.debug({
+        msg: "ignoring stale disconnect event",
+        connId: conn.id,
+        eventSocketId: socketId,
+        currentSocketId: conn.__socket.socketId
+      });
       return;
     }
+    if (wasClean) {
+      this.#removeConn(conn);
+    } else {
+      if (!conn.__driverState) {
+        this.rLog.warn("called conn disconnected without driver state");
+      }
+      conn.__persist.lastSeen = Date.now();
+      conn.__socket = void 0;
+    }
+  }
+  /**
+   * Removes a connection and cleans up its resources.
+   */
+  #removeConn(conn) {
     const connIdx = this.#persist.connections.findIndex(
       (c) => c.connId === conn.id
     );
@@ -1203,7 +1320,67 @@ var ActorInstance = class {
     }
     this.#resetSleepTimer();
   }
-  async prepareConn(params, request) {
+  /**
+   * Called to create a new connection or reconnect an existing one.
+   */
+  async createConn(socket, params, request, connectionId, connectionToken) {
+    this.#assertReady();
+    if (connectionId && connectionToken) {
+      this.rLog.debug({
+        msg: "checking for existing connection",
+        connectionId
+      });
+      const existingConn = this.#connections.get(connectionId);
+      if (existingConn && existingConn._token === connectionToken) {
+        this.rLog.debug({
+          msg: "reconnecting existing connection",
+          connectionId
+        });
+        if (existingConn.__driverState) {
+          const driverKind = getConnDriverKindFromState(
+            existingConn.__driverState
+          );
+          const driver = CONN_DRIVERS[driverKind];
+          if (driver.disconnect) {
+            driver.disconnect(
+              this,
+              existingConn,
+              existingConn.__driverState[driverKind],
+              "Reconnecting with new driver state"
+            );
+          }
+        }
+        existingConn.__socket = socket;
+        existingConn.__persist.lastSeen = Date.now();
+        this.#resetSleepTimer();
+        this.inspector.emitter.emit("connectionUpdated");
+        existingConn._sendMessage(
+          new CachedSerializer(
+            {
+              body: {
+                tag: "Init",
+                val: {
+                  actorId: this.id,
+                  connectionId: existingConn.id,
+                  connectionToken: existingConn._token
+                }
+              }
+            },
+            TO_CLIENT_VERSIONED
+          )
+        );
+        return existingConn;
+      }
+      this.rLog.debug({
+        msg: "connection not found or token mismatch, creating new connection",
+        connectionId
+      });
+    }
+    const newConnId = generateConnId();
+    const newConnToken = generateConnToken();
+    if (this.#connections.has(newConnId)) {
+      throw new Error(`Connection already exists: ${newConnId}`);
+    }
     let connState;
     const onBeforeConnectOpts = {
       request
@@ -1215,7 +1392,7 @@ var ActorInstance = class {
         params
       );
     }
-    if (this.#connStateEnabled) {
+    if (this.connStateEnabled) {
       if ("createConnState" in this.#config) {
         const dataOrPromise = this.#config.createConnState(
           this.actorContext,
@@ -1238,39 +1415,16 @@ var ActorInstance = class {
         );
       }
     }
-    return connState;
-  }
-  __getConnDriver(driverId) {
-    const driver = this.#connectionDrivers[driverId];
-    if (!driver) throw new Error(`No connection driver: ${driverId}`);
-    return driver;
-  }
-  /**
-   * Called after establishing a connection handshake.
-   */
-  async createConn(connectionId, connectionToken, params, state, driverId, driverState, authData) {
-    this.#assertReady();
-    if (this.#connections.has(connectionId)) {
-      throw new Error(`Connection already exists: ${connectionId}`);
-    }
-    const driver = this.__getConnDriver(driverId);
     const persist = {
-      connId: connectionId,
-      token: connectionToken,
-      connDriver: driverId,
-      connDriverState: driverState,
+      connId: newConnId,
+      token: newConnToken,
       params,
-      state,
-      authData,
+      state: connState,
       lastSeen: Date.now(),
       subscriptions: []
     };
-    const conn = new Conn(
-      this,
-      persist,
-      driver,
-      this.#connStateEnabled
-    );
+    const conn = new Conn(this, persist);
+    conn.__socket = socket;
     this.#connections.set(conn.id, conn);
     this.#resetSleepTimer();
     this.#persist.connections.push(persist);
@@ -1411,11 +1565,10 @@ var ActorInstance = class {
   #checkConnectionsLiveness() {
     this.#rLog.debug({ msg: "checking connections liveness" });
     for (const conn of this.#connections.values()) {
-      const liveness = conn[CONNECTION_CHECK_LIVENESS_SYMBOL]();
-      if (liveness.status === "connected") {
+      if (conn.__status === "connected") {
         this.#rLog.debug({ msg: "connection is alive", connId: conn.id });
       } else {
-        const lastSeen = liveness.lastSeen;
+        const lastSeen = conn.__persist.lastSeen;
         const sinceLastSeen = Date.now() - lastSeen;
         if (sinceLastSeen < this.#config.options.connectionLivenessTimeout) {
           this.#rLog.debug({
@@ -1431,7 +1584,7 @@ var ActorInstance = class {
           connId: conn.id,
           lastSeen
         });
-        this.__removeConn(conn);
+        this.#removeConn(conn);
       }
     }
   }
@@ -1761,7 +1914,7 @@ var ActorInstance = class {
         await this.#savePersistInner();
       } else {
         if (!this.#onPersistSavedPromise) {
-          this.#onPersistSavedPromise = Promise.withResolvers();
+          this.#onPersistSavedPromise = promiseWithResolvers();
         }
         this.#savePersistThrottled();
         await this.#onPersistSavedPromise.promise;
@@ -1925,10 +2078,6 @@ var ActorInstance = class {
       connections: persist.connections.map((conn) => ({
         id: conn.connId,
         token: conn.token,
-        driver: conn.connDriver,
-        driverState: bufferToArrayBuffer(
-          cbor2.encode(conn.connDriverState || {})
-        ),
         parameters: bufferToArrayBuffer(cbor2.encode(conn.params || {})),
         state: bufferToArrayBuffer(cbor2.encode(conn.state || {})),
         subscriptions: conn.subscriptions.map((sub) => ({
@@ -1957,8 +2106,6 @@ var ActorInstance = class {
       connections: bareData.connections.map((conn) => ({
         connId: conn.id,
         token: conn.token,
-        connDriver: conn.driver,
-        connDriverState: cbor2.decode(new Uint8Array(conn.driverState)),
         params: cbor2.decode(new Uint8Array(conn.parameters)),
         state: cbor2.decode(new Uint8Array(conn.state)),
         subscriptions: conn.subscriptions.map((sub) => ({
@@ -2699,7 +2846,7 @@ function createActorProxy(handle) {
         return Reflect.get(target, prop, receiver);
       }
       if (prop === "constructor" || prop in target) {
-        const value = Reflect.get(target, prop, receiver);
+        const value = Reflect.get(target, prop, target);
         if (typeof value === "function") {
           return value.bind(target);
         }
@@ -2757,7 +2904,7 @@ var ActorConnRaw = class {
   #abortController = new AbortController();
   /** If attempting to connect. Helpful for knowing if in a retry loop when reconnecting. */
   #connecting = false;
-  // These will only be set on SSE driver
+  // Connection info, used for reconnection and HTTP requests
   #actorId;
   #connectionId;
   #connectionToken;
@@ -2811,7 +2958,7 @@ var ActorConnRaw = class {
     logger().debug({ msg: "action", name: opts.name, args: opts.args });
     const actionId = this.#actionIdCounter;
     this.#actionIdCounter += 1;
-    const { promise, resolve, reject } = Promise.withResolvers();
+    const { promise, resolve, reject } = promiseWithResolvers();
     this.#actionsInFlight.set(actionId, { name: opts.name, resolve, reject });
     this.#sendMessage({
       body: {
@@ -2871,7 +3018,7 @@ var ActorConnRaw = class {
     try {
       if (this.#onOpenPromise)
         throw new Error("#onOpenPromise already defined");
-      this.#onOpenPromise = Promise.withResolvers();
+      this.#onOpenPromise = promiseWithResolvers();
       if (this.#client[TRANSPORT_SYMBOL] === "websocket") {
         await this.#connectWebSocket();
       } else if (this.#client[TRANSPORT_SYMBOL] === "sse") {
@@ -2890,15 +3037,34 @@ var ActorConnRaw = class {
       this.#actorQuery,
       this.#driver
     );
+    const isReconnection = this.#connectionId && this.#connectionToken;
+    if (isReconnection) {
+      logger().debug({
+        msg: "attempting websocket reconnection",
+        connectionId: this.#connectionId
+      });
+    }
     const ws = await this.#driver.openWebSocket(
       PATH_CONNECT_WEBSOCKET,
       actorId,
       this.#encoding,
-      this.#params
+      this.#params,
+      // Pass connection ID and token for reconnection if available
+      isReconnection ? this.#connectionId : void 0,
+      isReconnection ? this.#connectionToken : void 0
     );
+    logger().debug({
+      msg: "transport set to new websocket",
+      connectionId: this.#connectionId,
+      readyState: ws.readyState,
+      messageQueueLength: this.#messageQueue.length
+    });
     this.#transport = { websocket: ws };
     ws.addEventListener("open", () => {
-      logger().debug({ msg: "websocket open" });
+      logger().debug({
+        msg: "client websocket open",
+        connectionId: this.#connectionId
+      });
     });
     ws.addEventListener("message", async (ev) => {
       this.#handleOnMessage(ev.data);
@@ -2924,6 +3090,7 @@ var ActorConnRaw = class {
       actorId,
       encoding: this.#encoding
     });
+    const isReconnection = this.#connectionId && this.#connectionToken;
     const eventSource = new EventSource("http://actor/connect/sse", {
       fetch: (input, init) => {
         return this.#driver.sendRequest(
@@ -2934,7 +3101,11 @@ var ActorConnRaw = class {
               ...init == null ? void 0 : init.headers,
               "User-Agent": httpUserAgent(),
               [HEADER_ENCODING]: this.#encoding,
-              ...this.#params !== void 0 ? { [HEADER_CONN_PARAMS]: JSON.stringify(this.#params) } : {}
+              ...this.#params !== void 0 ? { [HEADER_CONN_PARAMS]: JSON.stringify(this.#params) } : {},
+              ...isReconnection ? {
+                [HEADER_CONN_ID]: this.#connectionId,
+                [HEADER_CONN_TOKEN]: this.#connectionToken
+              } : {}
             }
           })
         );
@@ -2942,6 +3113,7 @@ var ActorConnRaw = class {
     });
     this.#transport = { sse: eventSource };
     eventSource.addEventListener("message", (ev) => {
+      if (ev.type === "ping") return;
       this.#handleOnMessage(ev.data);
     });
     eventSource.addEventListener("error", (ev) => {
@@ -2952,7 +3124,8 @@ var ActorConnRaw = class {
   #handleOnOpen() {
     logger().debug({
       msg: "socket open",
-      messageQueueLength: this.#messageQueue.length
+      messageQueueLength: this.#messageQueue.length,
+      connectionId: this.#connectionId
     });
     if (this.#onOpenPromise) {
       this.#onOpenPromise.resolve(void 0);
@@ -2964,6 +3137,10 @@ var ActorConnRaw = class {
     }
     const queue = this.#messageQueue;
     this.#messageQueue = [];
+    logger().debug({
+      msg: "flushing message queue",
+      queueLength: queue.length
+    });
     for (const msg of queue) {
       this.#sendMessage(msg);
     }
@@ -3056,23 +3233,38 @@ var ActorConnRaw = class {
       this.#onOpenPromise.reject(new Error("Closed"));
     }
     const closeEvent = event;
-    if (closeEvent.wasClean) {
-      logger().info({
-        msg: "socket closed",
-        code: closeEvent.code,
-        reason: closeEvent.reason,
-        wasClean: closeEvent.wasClean
-      });
-    } else {
-      logger().warn({
-        msg: "socket closed",
-        code: closeEvent.code,
-        reason: closeEvent.reason,
-        wasClean: closeEvent.wasClean
+    const wasClean = closeEvent.wasClean;
+    logger().info({
+      msg: "socket closed",
+      code: closeEvent.code,
+      reason: closeEvent.reason,
+      wasClean,
+      connectionId: this.#connectionId,
+      messageQueueLength: this.#messageQueue.length,
+      actionsInFlight: this.#actionsInFlight.size
+    });
+    if (this.#actionsInFlight.size > 0) {
+      logger().debug({
+        msg: "rejecting in-flight actions after disconnect",
+        count: this.#actionsInFlight.size,
+        connectionId: this.#connectionId,
+        wasClean
       });
+      const disconnectError = new Error(
+        wasClean ? "Connection closed" : "Connection lost"
+      );
+      for (const actionInfo of this.#actionsInFlight.values()) {
+        actionInfo.reject(disconnectError);
+      }
+      this.#actionsInFlight.clear();
     }
     this.#transport = void 0;
     if (!this.#disposed && !this.#connecting) {
+      logger().debug({
+        msg: "triggering reconnect",
+        connectionId: this.#connectionId,
+        messageQueueLength: this.#messageQueue.length
+      });
       this.#connectWithRetry();
     }
   }
@@ -3176,14 +3368,25 @@ var ActorConnRaw = class {
     };
   }
   #sendMessage(message, opts) {
+    var _a, _b;
     if (this.#disposed) {
       throw new ActorConnDisposed();
     }
     let queueMessage = false;
     if (!this.#transport) {
+      logger().debug({ msg: "no transport, queueing message" });
       queueMessage = true;
     } else if ("websocket" in this.#transport) {
-      if (this.#transport.websocket.readyState === 1) {
+      const readyState = this.#transport.websocket.readyState;
+      logger().debug({
+        msg: "websocket send attempt",
+        readyState,
+        readyStateString: readyState === 0 ? "CONNECTING" : readyState === 1 ? "OPEN" : readyState === 2 ? "CLOSING" : "CLOSED",
+        connectionId: this.#connectionId,
+        messageType: message.body.tag,
+        actionName: (_a = message.body.val) == null ? void 0 : _a.name
+      });
+      if (readyState === 1) {
         try {
           const messageSerialized = serializeWithEncoding(
             this.#encoding,
@@ -3198,11 +3401,16 @@ var ActorConnRaw = class {
         } catch (error) {
           logger().warn({
             msg: "failed to send message, added to queue",
-            error
+            error,
+            connectionId: this.#connectionId
           });
           queueMessage = true;
         }
       } else {
+        logger().debug({
+          msg: "websocket not open, queueing message",
+          readyState
+        });
         queueMessage = true;
       }
     } else if ("sse" in this.#transport) {
@@ -3216,7 +3424,13 @@ var ActorConnRaw = class {
     }
     if (!(opts == null ? void 0 : opts.ephemeral) && queueMessage) {
       this.#messageQueue.push(message);
-      logger().debug({ msg: "queued connection message" });
+      logger().debug({
+        msg: "queued connection message",
+        queueLength: this.#messageQueue.length,
+        connectionId: this.#connectionId,
+        messageType: message.body.tag,
+        actionName: (_b = message.body.val) == null ? void 0 : _b.name
+      });
     }
   }
   async #sendHttpMessage(message, opts) {
@@ -3273,6 +3487,20 @@ var ActorConnRaw = class {
     const buffer = await inputDataToBuffer(data);
     return deserializeWithEncoding(this.#encoding, buffer, TO_CLIENT_VERSIONED);
   }
+  /**
+   * Get the actor ID (for testing purposes).
+   * @internal
+   */
+  get actorId() {
+    return this.#actorId;
+  }
+  /**
+   * Get the connection ID (for testing purposes).
+   * @internal
+   */
+  get connectionId() {
+    return this.#connectionId;
+  }
   /**
    * Disconnects from the actor.
    *
@@ -3294,7 +3522,7 @@ var ActorConnRaw = class {
       if (ws.readyState === 2 || ws.readyState === 3) {
         logger().debug({ msg: "ws already closed or closing" });
       } else {
-        const { promise, resolve } = Promise.withResolvers();
+        const { promise, resolve } = promiseWithResolvers();
         ws.addEventListener("close", () => {
           logger().debug({ msg: "ws closed" });
           resolve(void 0);
@@ -3303,6 +3531,28 @@ var ActorConnRaw = class {
         await promise;
       }
     } else if ("sse" in this.#transport) {
+      if (this.#connectionId && this.#connectionToken) {
+        try {
+          await sendHttpRequest({
+            url: "http://actor/connections/close",
+            method: "POST",
+            headers: {
+              [HEADER_CONN_ID]: this.#connectionId,
+              [HEADER_CONN_TOKEN]: this.#connectionToken
+            },
+            encoding: this.#encoding,
+            skipParseResponse: true,
+            customFetch: this.#driver.sendRequest.bind(
+              this.#driver,
+              this.#actorId
+            ),
+            requestVersionedDataHandler: TO_SERVER_VERSIONED,
+            responseVersionedDataHandler: TO_CLIENT_VERSIONED
+          });
+        } catch (error) {
+          logger().warn({ msg: "failed to send close request", error });
+        }
+      }
       this.#transport.sse.close();
     } else {
       assertUnreachable(this.#transport);
@@ -3352,10 +3602,16 @@ async function apiCall(config, method, path, body) {
     namespace: config.namespace
   });
   logger2().debug({ msg: "making api call", method, url });
+  const headers = {
+    ...config.headers
+  };
+  if (config.token) {
+    headers.Authorization = `Bearer ${config.token}`;
+  }
   return await sendHttpRequest({
     method,
     url,
-    headers: {},
+    headers,
     body,
     encoding: "json",
     skipParseResponse: false,
@@ -3370,18 +3626,21 @@ async function sendHttpRequestToActor(runConfig, actorId, actorRequest) {
   const endpoint = getEndpoint(runConfig);
   const guardUrl = combineUrlPath(endpoint, url.pathname + url.search);
   let bodyToSend = null;
-  const guardHeaders = buildGuardHeadersForHttp(actorRequest, actorId);
-  if (actorRequest.body && actorRequest.method !== "GET" && actorRequest.method !== "HEAD") {
+  const guardHeaders = buildGuardHeadersForHttp(
+    runConfig,
+    actorRequest,
+    actorId
+  );
+  if (actorRequest.method !== "GET" && actorRequest.method !== "HEAD") {
     if (actorRequest.bodyUsed) {
       throw new Error("Request body has already been consumed");
     }
-    const clonedRequest = actorRequest.clone();
-    bodyToSend = await clonedRequest.arrayBuffer();
-    guardHeaders.delete("transfer-encoding");
-    guardHeaders.set(
-      "content-length",
-      String(bodyToSend.byteLength)
-    );
+    const reqBody = await actorRequest.arrayBuffer();
+    if (reqBody.byteLength !== 0) {
+      bodyToSend = reqBody;
+      guardHeaders.delete("transfer-encoding");
+      guardHeaders.set("content-length", String(bodyToSend.byteLength));
+    }
   }
   const guardRequest = new Request(guardUrl, {
     method: actorRequest.method,
@@ -3394,19 +3653,24 @@ async function sendHttpRequestToActor(runConfig, actorId, actorRequest) {
 function mutableResponse(fetchRes) {
   return new Response(fetchRes.body, fetchRes);
 }
-function buildGuardHeadersForHttp(actorRequest, actorId) {
+function buildGuardHeadersForHttp(runConfig, actorRequest, actorId) {
   const headers = new Headers();
   for (const [key, value] of actorRequest.headers.entries()) {
     headers.set(key, value);
   }
-  headers.set("x-rivet-target", "actor");
-  headers.set("x-rivet-actor", actorId);
-  headers.set("x-rivet-port", "main");
+  for (const [key, value] of Object.entries(runConfig.headers)) {
+    headers.set(key, value);
+  }
+  headers.set(HEADER_RIVET_TARGET, "actor");
+  headers.set(HEADER_RIVET_ACTOR, actorId);
+  if (runConfig.token) {
+    headers.set(HEADER_RIVET_TOKEN, runConfig.token);
+  }
   return headers;
 }
 
 // src/remote-manager-driver/actor-websocket-client.ts
-async function openWebSocketToActor(runConfig, path, actorId, encoding, params) {
+async function openWebSocketToActor(runConfig, path, actorId, encoding, params, connId, connToken) {
   const WebSocket2 = await importWebSocket();
   const endpoint = getEndpoint(runConfig);
   const guardUrl = combineUrlPath(endpoint, path);
@@ -3416,49 +3680,65 @@ async function openWebSocketToActor(runConfig, path, actorId, encoding, params)
     path,
     guardUrl
   });
-  const ws = new WebSocket2(guardUrl, {
-    headers: buildGuardHeadersForWebSocket(actorId, encoding, params)
-  });
+  const ws = new WebSocket2(
+    guardUrl,
+    buildWebSocketProtocols(
+      runConfig,
+      actorId,
+      encoding,
+      params,
+      connId,
+      connToken
+    )
+  );
   ws.binaryType = "arraybuffer";
   logger2().debug({ msg: "websocket connection opened", actorId });
   return ws;
 }
-function buildGuardHeadersForWebSocket(actorId, encoding, params, authData) {
-  const headers = {};
-  headers["x-rivet-target"] = "actor";
-  headers["x-rivet-actor"] = actorId;
-  headers["x-rivet-port"] = "main";
-  headers[HEADER_ENCODING] = encoding;
+function buildWebSocketProtocols(runConfig, actorId, encoding, params, connId, connToken) {
+  const protocols = [];
+  protocols.push(WS_PROTOCOL_STANDARD);
+  protocols.push(`${WS_PROTOCOL_TARGET}actor`);
+  protocols.push(`${WS_PROTOCOL_ACTOR}${actorId}`);
+  protocols.push(`${WS_PROTOCOL_ENCODING}${encoding}`);
+  if (runConfig.token) {
+    protocols.push(`${WS_PROTOCOL_TOKEN}${runConfig.token}`);
+  }
   if (params) {
-    headers[HEADER_CONN_PARAMS] = JSON.stringify(params);
+    protocols.push(
+      `${WS_PROTOCOL_CONN_PARAMS}${encodeURIComponent(JSON.stringify(params))}`
+    );
   }
-  if (authData) {
-    headers[HEADER_AUTH_DATA] = JSON.stringify(authData);
+  if (connId) {
+    protocols.push(`${WS_PROTOCOL_CONN_ID}${connId}`);
   }
-  return headers;
+  if (connToken) {
+    protocols.push(`${WS_PROTOCOL_CONN_TOKEN}${connToken}`);
+  }
+  return protocols;
 }
 
 // src/remote-manager-driver/api-endpoints.ts
-async function getActor(config, actorId) {
+async function getActor(config, name, actorId) {
   return apiCall(
     config,
     "GET",
-    `/actors/${encodeURIComponent(actorId)}`
+    `/actors?name=${name}&actor_ids=${encodeURIComponent(actorId)}`
   );
 }
-async function getActorById(config, name, key) {
+async function getActorByKey(config, name, key) {
   const serializedKey = serializeActorKey(key);
   return apiCall(
     config,
     "GET",
-    `/actors/by-id?name=${encodeURIComponent(name)}&key=${encodeURIComponent(serializedKey)}`
+    `/actors?name=${encodeURIComponent(name)}&key=${encodeURIComponent(serializedKey)}`
   );
 }
-async function getOrCreateActorById(config, request) {
+async function getOrCreateActor(config, request) {
   return apiCall(
     config,
     "PUT",
-    `/actors/by-id`,
+    `/actors`,
     request
   );
 }
@@ -3479,13 +3759,8 @@ async function destroyActor(config, actorId) {
 }
 
 // src/remote-manager-driver/ws-proxy.ts
-async function createWebSocketProxy(c, targetUrl, headers) {
+async function createWebSocketProxy(c, targetUrl, protocols) {
   const WebSocket2 = await importWebSocket();
-  for (const [k, v] of c.req.raw.headers.entries()) {
-    if (!k.startsWith("sec-") && k !== "connection" && k !== "upgrade") {
-      headers[k] = v;
-    }
-  }
   const state = {};
   return {
     onOpen: async (event, clientWs) => {
@@ -3498,7 +3773,7 @@ async function createWebSocketProxy(c, targetUrl, headers) {
         });
         return;
       }
-      const targetWs = new WebSocket2(targetUrl, { headers });
+      const targetWs = new WebSocket2(targetUrl, protocols);
       state.targetWs = targetWs;
       state.connectPromise = new Promise((resolve, reject) => {
         targetWs.addEventListener("open", () => {
@@ -3620,31 +3895,26 @@ var RemoteManagerDriver = class {
     name,
     actorId
   }) {
-    try {
-      const response = await getActor(this.#config, actorId);
-      if (response.actor.name !== name) {
-        logger2().debug({
-          msg: "actor name mismatch from api",
-          actorId,
-          apiName: response.actor.name,
-          requestedName: name
-        });
-        return void 0;
-      }
-      const keyRaw = response.actor.key;
-      invariant6(keyRaw, `actor ${actorId} should have key`);
-      const key = deserializeActorKey(keyRaw);
-      return {
+    const response = await getActor(this.#config, name, actorId);
+    const actor = response.actors[0];
+    if (!actor) return void 0;
+    if (actor.name !== name) {
+      logger2().debug({
+        msg: "actor name mismatch from api",
         actorId,
-        name,
-        key
-      };
-    } catch (error) {
-      if (error instanceof EngineApiError && error.group === "actor" && error.code === "not_found") {
-        return void 0;
-      }
-      throw error;
+        apiName: actor.name,
+        requestedName: name
+      });
+      return void 0;
     }
+    const keyRaw = actor.key;
+    invariant6(keyRaw, `actor ${actorId} should have key`);
+    const key = deserializeActorKey(keyRaw);
+    return {
+      actorId,
+      name,
+      key
+    };
   }
   async getWithKey({
     c,
@@ -3653,11 +3923,10 @@ var RemoteManagerDriver = class {
   }) {
     logger2().debug({ msg: "getWithKey: searching for actor", name, key });
     try {
-      const response = await getActorById(this.#config, name, key);
-      if (!response.actor_id) {
-        return void 0;
-      }
-      const actorId = response.actor_id;
+      const response = await getActorByKey(this.#config, name, key);
+      const actor = response.actors[0];
+      if (!actor) return void 0;
+      const actorId = actor.actor_id;
       logger2().debug({
         msg: "getWithKey: found actor via api",
         actorId,
@@ -3683,20 +3952,20 @@ var RemoteManagerDriver = class {
       name,
       key
     });
-    const response = await getOrCreateActorById(this.#config, {
+    const { actor, created } = await getOrCreateActor(this.#config, {
       name,
       key: serializeActorKey(key),
       runner_name_selector: this.#config.runnerName,
-      input: input ? cbor6.encode(actorInput).toString("base64") : void 0,
+      input: actorInput ? cbor6.encode(actorInput).toString("base64") : void 0,
       crash_policy: "sleep"
     });
-    const actorId = response.actor_id;
+    const actorId = actor.actor_id;
     logger2().info({
       msg: "getOrCreateWithKey: actor ready",
       actorId,
       name,
       key,
-      created: response.created
+      created
     });
     return {
       actorId,
@@ -3734,19 +4003,21 @@ var RemoteManagerDriver = class {
   async sendRequest(actorId, actorRequest) {
     return await sendHttpRequestToActor(this.#config, actorId, actorRequest);
   }
-  async openWebSocket(path, actorId, encoding, params) {
+  async openWebSocket(path, actorId, encoding, params, connId, connToken) {
     return await openWebSocketToActor(
       this.#config,
       path,
       actorId,
       encoding,
-      params
+      params,
+      connId,
+      connToken
     );
   }
   async proxyRequest(_c, actorRequest, actorId) {
     return await sendHttpRequestToActor(this.#config, actorId, actorRequest);
   }
-  async proxyWebSocket(c, path, actorId, encoding, params, authData) {
+  async proxyWebSocket(c, path, actorId, encoding, params, connId, connToken) {
     var _a, _b;
     const upgradeWebSocket = (_b = (_a = this.#config).getUpgradeWebSocket) == null ? void 0 : _b.call(_a);
     invariant6(upgradeWebSocket, "missing getUpgradeWebSocket");
@@ -3759,27 +4030,31 @@ var RemoteManagerDriver = class {
       path,
       guardUrl
     });
-    const headers = buildGuardHeadersForWebSocket(
+    const protocols = buildWebSocketProtocols(
+      this.#config,
       actorId,
       encoding,
       params,
-      authData
+      connId,
+      connToken
     );
-    const args = await createWebSocketProxy(c, wsGuardUrl, headers);
+    const args = await createWebSocketProxy(c, wsGuardUrl, protocols);
     return await upgradeWebSocket(() => args)(c, noopNext());
   }
   displayInformation() {
     return { name: "Remote", properties: {} };
   }
+  getOrCreateInspectorAccessToken() {
+    return generateRandomString();
+  }
 };
 
 export {
   generateConnId,
   generateConnToken,
-  CONNECTION_DRIVER_WEBSOCKET,
-  CONNECTION_DRIVER_SSE,
-  CONNECTION_DRIVER_HTTP,
+  generateConnSocketId,
   createActorInspectorRouter,
+  serializeActorKey,
   deserializeActorKey,
   ActorDefinition,
   lookupInRegistry,
@@ -3793,4 +4068,4 @@ export {
   createClientWithDriver,
   RemoteManagerDriver
 };
-//# sourceMappingURL=chunk-4CKHQRXG.js.map
+//# sourceMappingURL=chunk-NII4KKHD.js.map