rivetkit 2.0.7 → 2.0.9

package/dist/tsup/{chunk-KHRZPP5T.js → chunk-FZP2IBIX.js}

@@ -4,7 +4,7 @@ import {
   inspectorLogger,
   isInspectorEnabled,
   secureInspector
-} from "./chunk-PD6HCAJE.js";
+} from "./chunk-YR2VY4XS.js";
 import {
   ActorDefinition,
   RemoteManagerDriver,
@@ -14,10 +14,10 @@ import {
   generateConnSocketId,
   lookupInRegistry,
   serializeActorKey
-} from "./chunk-NII4KKHD.js";
+} from "./chunk-KHZ2QSQ4.js";
 import {
   CreateActorSchema
-} from "./chunk-54MAHBLL.js";
+} from "./chunk-WBSPHV5V.js";
 import {
   ActionContext,
   HTTP_ACTION_REQUEST_VERSIONED,
@@ -28,7 +28,7 @@ import {
   createVersionedDataHandler,
   parseMessage,
   serializeEmptyPersistData
-} from "./chunk-2NL3KGJ7.js";
+} from "./chunk-VVCL5DXN.js";
 import {
   EncodingSchema,
   HEADER_ACTOR_ID,
@@ -54,12 +54,12 @@ import {
   generateRandomString,
   loggerWithoutContext,
   serializeWithEncoding
-} from "./chunk-NRELKXIX.js";
+} from "./chunk-R7OP5N25.js";
 import {
   configureBaseLogger,
   configureDefaultLogger,
   getLogger
-} from "./chunk-YL4VZMMT.js";
+} from "./chunk-DLPIL3VC.js";
 import {
   bufferToArrayBuffer,
   deconstructError,
@@ -69,7 +69,7 @@ import {
   promiseWithResolvers,
   setLongTimeout,
   stringifyError
-} from "./chunk-RLBM6D4L.js";
+} from "./chunk-E63WU5PL.js";
 import {
   ActorAlreadyExists,
   ConnNotFound,
@@ -571,7 +571,8 @@ async function handleWebSocketConnect(req, runConfig, actorDriver, actorId, enco
       }
       ws.close(1e3, "hack_force_close");
       handlersPromise.then(({ conn, actor: actor3 }) => {
-        actor3.__connDisconnected(conn, event.wasClean, socketId);
+        const wasClean = event.wasClean || event.code === 1e3;
+        actor3.__connDisconnected(conn, wasClean, socketId);
       }).catch((error) => {
         deconstructError(
           error,
@@ -881,7 +882,7 @@ function handleRouteError(error, c) {
 }
 
 // src/actor/router.ts
-function createActorRouter(runConfig, actorDriver) {
+function createActorRouter(runConfig, actorDriver, isTest) {
   const router = new Hono({ strict: false });
   router.use("*", loggerMiddleware(loggerWithoutContext()));
   router.get("/", (c) => {
@@ -892,26 +893,28 @@ function createActorRouter(runConfig, actorDriver) {
   router.get("/health", (c) => {
     return c.text("ok");
   });
-  router.post("/.test/force-disconnect", async (c) => {
-    const connId = c.req.query("conn");
-    if (!connId) {
-      return c.text("Missing conn query parameter", 400);
-    }
-    const actor2 = await actorDriver.loadActor(c.env.actorId);
-    const conn = actor2.__getConnForId(connId);
-    if (!conn) {
-      return c.text(`Connection not found: ${connId}`, 404);
-    }
-    const driverState = conn.__driverState;
-    if (driverState && 0 /* WEBSOCKET */ in driverState) {
-      const ws = driverState[0 /* WEBSOCKET */].websocket;
-      ws.raw.terminate();
-    } else if (driverState && 1 /* SSE */ in driverState) {
-      const stream = driverState[1 /* SSE */].stream;
-      stream.abort();
-    }
-    return c.json({ success: true });
-  });
+  if (isTest) {
+    router.post("/.test/force-disconnect", async (c) => {
+      const connId = c.req.query("conn");
+      if (!connId) {
+        return c.text("Missing conn query parameter", 400);
+      }
+      const actor2 = await actorDriver.loadActor(c.env.actorId);
+      const conn = actor2.__getConnForId(connId);
+      if (!conn) {
+        return c.text(`Connection not found: ${connId}`, 404);
+      }
+      const driverState = conn.__driverState;
+      if (driverState && 0 /* WEBSOCKET */ in driverState) {
+        const ws = driverState[0 /* WEBSOCKET */].websocket;
+        ws.raw.terminate();
+      } else if (driverState && 1 /* SSE */ in driverState) {
+        const stream = driverState[1 /* SSE */].stream;
+        stream.abort();
+      }
+      return c.json({ success: true });
+    });
+  }
   router.get(PATH_CONNECT_WEBSOCKET, async (c) => {
     var _a;
     const upgradeWebSocket = (_a = runConfig.getUpgradeWebSocket) == null ? void 0 : _a.call(runConfig);
@@ -1437,6 +1440,7 @@ var EngineActorDriver = class {
   #actorRouter;
   #version = 1;
   // Version for the runner protocol
+  #alarmTimeout;
   #runnerStarted = Promise.withResolvers();
   #runnerStopped = Promise.withResolvers();
   constructor(registryConfig, runConfig, managerDriver, inlineClient, config2) {
@@ -1445,15 +1449,19 @@ var EngineActorDriver = class {
     this.#managerDriver = managerDriver;
     this.#inlineClient = inlineClient;
     this.#config = config2;
-    this.#actorRouter = createActorRouter(runConfig, this);
+    this.#actorRouter = createActorRouter(
+      runConfig,
+      this,
+      registryConfig.test.enabled
+    );
     let hasDisconnected = false;
     const runnerConfig = {
       version: this.#version,
       endpoint: config2.endpoint,
-      token: config2.token,
+      token: runConfig.token ?? config2.token,
       pegboardEndpoint: config2.pegboardEndpoint,
       namespace: config2.namespace,
-      totalSlots: config2.totalSlots,
+      totalSlots: runConfig.totalSlots ?? config2.totalSlots,
       runnerName: config2.runnerName,
       runnerKey: config2.runnerKey,
       metadata: {
@@ -1541,6 +1549,16 @@ var EngineActorDriver = class {
     await this.#runner.kvPut(actorId, [[KEYS.PERSIST_DATA, data]]);
   }
   async setAlarm(actor2, timestamp) {
+    if (this.#alarmTimeout) {
+      this.#alarmTimeout.abort();
+      this.#alarmTimeout = void 0;
+    }
+    const delay = Math.max(0, timestamp - Date.now());
+    this.#alarmTimeout = setLongTimeout(() => {
+      actor2._onAlarm();
+      this.#alarmTimeout = void 0;
+    }, delay);
+    this.#runner.setAlarm(actor2.id, timestamp);
   }
   async getDatabase(_actorId) {
     return void 0;
@@ -1700,9 +1718,9 @@ var EngineActorDriver = class {
   async serverlessHandleStart(c) {
     await this.#runnerStarted.promise;
     return streamSSE2(c, async (stream) => {
-      const runnerId = this.#runner.runnerId;
-      invariant3(runnerId, "runnerId not set");
-      stream.writeSSE({ data: runnerId });
+      const payload = this.#runner.getServerlessInitPacket();
+      invariant3(payload, "runnerId not set");
+      stream.writeSSE({ data: payload });
       return this.#runnerStopped.promise;
     });
   }
@@ -2555,7 +2573,11 @@ var FileSystemManagerDriver = class {
       this,
       inlineClient
     );
-    this.#actorRouter = createActorRouter(this.#runConfig, this.#actorDriver);
+    this.#actorRouter = createActorRouter(
+      this.#runConfig,
+      this.#actorDriver,
+      registryConfig.test.enabled
+    );
   }
   async sendRequest(actorId, actorRequest) {
     return await this.#actorRouter.fetch(actorRequest, {
@@ -2762,7 +2784,6 @@ import {
   Hono as Hono3
 } from "hono";
 import { cors as corsMiddleware } from "hono/cors";
-import { createMiddleware } from "hono/factory";
 import invariant6 from "invariant";
 import { z as z5 } from "zod";
 
@@ -3105,45 +3126,62 @@ function createManagerRouter(registryConfig, runConfig, managerDriver, serverles
     runConfig.basePath
   );
   router.use("*", loggerMiddleware(logger()));
-  const cors2 = runConfig.cors ? corsMiddleware(runConfig.cors) : createMiddleware((_c, next) => next());
   if (serverlessActorDriverBuilder) {
-    addServerlessRoutes(serverlessActorDriverBuilder, router, cors2);
+    addServerlessRoutes(runConfig, serverlessActorDriverBuilder, router);
   } else {
-    addManagerRoutes(registryConfig, runConfig, managerDriver, router, cors2);
+    addManagerRoutes(registryConfig, runConfig, managerDriver, router);
   }
   router.notFound(handleRouteNotFound);
   router.onError(handleRouteError);
-  return { router, openapi: router, cors: cors2 };
+  return { router, openapi: router };
 }
-function addServerlessRoutes(serverlessActorDriverBuilder, router, cors2) {
-  router.get("/", cors2, (c) => {
+function addServerlessRoutes(runConfig, serverlessActorDriverBuilder, router) {
+  if (runConfig.cors) router.use("*", corsMiddleware(runConfig.cors));
+  router.get("/", (c) => {
     return c.text(
       "This is a RivetKit server.\n\nLearn more at https://rivetkit.org"
     );
   });
-  router.get("/start", cors2, async (c) => {
-    const actorDriver = serverlessActorDriverBuilder();
+  router.get("/start", async (c) => {
+    const token = c.req.header("x-rivet-token");
+    let totalSlots = parseInt(
+      c.req.header("x-rivetkit-total-slots")
+    );
+    if (isNaN(totalSlots)) totalSlots = void 0;
+    const actorDriver = serverlessActorDriverBuilder(token, totalSlots);
     invariant6(
       actorDriver.serverlessHandleStart,
       "missing serverlessHandleStart on ActorDriver"
     );
     return await actorDriver.serverlessHandleStart(c);
   });
-  router.get("/health", cors2, (c) => {
+  router.get("/health", (c) => {
     return c.text("ok");
   });
 }
-function addManagerRoutes(registryConfig, runConfig, managerDriver, router, cors2) {
+function addManagerRoutes(registryConfig, runConfig, managerDriver, router) {
   var _a;
-  router.use("*", cors2, actorGateway.bind(void 0, runConfig, managerDriver));
-  router.get("/", cors2, (c) => {
+  if (isInspectorEnabled(runConfig, "manager")) {
+    if (!managerDriver.inspector) {
+      throw new Unsupported("inspector");
+    }
+    router.route(
+      "/inspect",
+      new Hono3().use(corsMiddleware(runConfig.inspector.cors)).use(secureInspector(runConfig)).use((c, next) => {
+        c.set("inspector", managerDriver.inspector);
+        return next();
+      }).route("/", createManagerInspectorRouter())
+    );
+  }
+  if (runConfig.cors) router.use("*", corsMiddleware(runConfig.cors));
+  router.use("*", actorGateway.bind(void 0, runConfig, managerDriver));
+  router.get("/", (c) => {
     return c.text(
       "This is a RivetKit server.\n\nLearn more at https://rivetkit.org"
     );
   });
   {
     const route = createRoute({
-      middleware: [cors2],
       method: "get",
       path: "/actors",
       request: {
@@ -3203,7 +3241,6 @@ function addManagerRoutes(registryConfig, runConfig, managerDriver, router, cors
   }
   {
     const route = createRoute({
-      middleware: [cors2],
       method: "put",
       path: "/actors",
       request: {
@@ -3248,7 +3285,6 @@ function addManagerRoutes(registryConfig, runConfig, managerDriver, router, cors
   }
   {
     const route = createRoute({
-      middleware: [cors2],
       method: "post",
       path: "/actors",
       request: {
@@ -3435,7 +3471,7 @@ function addManagerRoutes(registryConfig, runConfig, managerDriver, router, cors
       }
     });
   }
-  router.get("/health", cors2, (c) => {
+  router.get("/health", (c) => {
     return c.text("ok");
   });
   (_a = managerDriver.modifyManagerRouter) == null ? void 0 : _a.call(
@@ -3443,18 +3479,6 @@ function addManagerRoutes(registryConfig, runConfig, managerDriver, router, cors
     registryConfig,
     router
   );
-  if (isInspectorEnabled(runConfig, "manager")) {
-    if (!managerDriver.inspector) {
-      throw new Unsupported("inspector");
-    }
-    router.route(
-      "/inspect",
-      new Hono3().use(corsMiddleware(runConfig.inspector.cors)).use(secureInspector(runConfig)).use((c, next) => {
-        c.set("inspector", managerDriver.inspector);
-        return next();
-      }).route("/", createManagerInspectorRouter())
-    );
-  }
 }
 function createApiActor(actor2) {
   return {
@@ -3598,7 +3622,7 @@ var Registry = class {
         console.log(`  - ${k}:${padding}${v}`);
       }
       if (isInspectorEnabled(config2, "manager") && managerDriver.inspector) {
-        console.log(`  - Inspector:     ${getInspectorUrl(config2)}`);
+        console.log(`  - Inspector:    ${getInspectorUrl(config2)}`);
       }
       console.log();
     }
@@ -3677,7 +3701,9 @@ var Registry = class {
       }
       console.log();
     }
-    let serverlessActorDriverBuilder = () => {
+    let serverlessActorDriverBuilder = (token, totalSlots) => {
+      if (token) config2.token = token;
+      if (totalSlots) config2.totalSlots = totalSlots;
       return driver.actor(this.#config, config2, managerDriver, client);
     };
     if (!config2.disableActorDriver) {
@@ -3723,4 +3749,4 @@ export {
   setup
 };
 //! These configs configs hold anything that's not platform-specific about running actors.
-//# sourceMappingURL=chunk-KHRZPP5T.js.map
+//# sourceMappingURL=chunk-FZP2IBIX.js.map