@rivetkit/cloudflare-workers 0.9.0 → 0.9.2

package/dist/mod.js

@@ -1,18 +1,23 @@
+// src/handler.ts
+import { AsyncLocalStorage } from "async_hooks";
+import { Hono } from "hono";
+import invariant2 from "invariant";
+
 // src/actor-handler-do.ts
 import { DurableObject } from "cloudflare:workers";
+import {
+  createActorRouter,
+  createClientWithDriver,
+  createInlineClientDriver
+} from "@rivetkit/core";
 import { serializeEmptyPersistData } from "@rivetkit/core/driver-helpers";
 
-// src/log.ts
-import { getLogger } from "@rivetkit/core/log";
-var LOGGER_NAME = "driver-cloudflare-workers";
-function logger() {
-  return getLogger(LOGGER_NAME);
-}
-
-// src/actor-handler-do.ts
-import { PartitionTopologyActor } from "@rivetkit/core/topologies/partition";
-
 // src/actor-driver.ts
+import {
+  createGenericConnDrivers,
+  GenericConnGlobalState,
+  lookupInRegistry
+} from "@rivetkit/core";
 import invariant from "invariant";
 var CloudflareDurableObjectGlobalState = class {
   // Single map for all actor state
@@ -26,14 +31,76 @@ var CloudflareDurableObjectGlobalState = class {
     this.#dos.set(actorId, state);
   }
 };
+var ActorHandler = class {
+  actor;
+  actorPromise = Promise.withResolvers();
+  genericConnGlobalState = new GenericConnGlobalState();
+};
 var CloudflareActorsActorDriver = class {
+  #registryConfig;
+  #runConfig;
+  #managerDriver;
+  #inlineClient;
   #globalState;
-  constructor(globalState) {
+  #actors = /* @__PURE__ */ new Map();
+  constructor(registryConfig, runConfig, managerDriver, inlineClient, globalState) {
+    this.#registryConfig = registryConfig;
+    this.#runConfig = runConfig;
+    this.#managerDriver = managerDriver;
+    this.#inlineClient = inlineClient;
     this.#globalState = globalState;
   }
   #getDOCtx(actorId) {
     return this.#globalState.getDOState(actorId).ctx;
   }
+  async loadActor(actorId) {
+    var _a;
+    let handler = this.#actors.get(actorId);
+    if (handler) {
+      if (handler.actorPromise) await handler.actorPromise.promise;
+      if (!handler.actor) throw new Error("Actor should be loaded");
+      return handler.actor;
+    }
+    handler = new ActorHandler();
+    this.#actors.set(actorId, handler);
+    const doState = this.#globalState.getDOState(actorId);
+    const storage = doState.ctx.storage;
+    const [name, key] = await Promise.all([
+      storage.get(KEYS.NAME),
+      storage.get(KEYS.KEY)
+    ]);
+    if (!name) {
+      throw new Error(`Actor ${actorId} is not initialized - missing name`);
+    }
+    if (!key) {
+      throw new Error(`Actor ${actorId} is not initialized - missing key`);
+    }
+    const definition = lookupInRegistry(this.#registryConfig, name);
+    handler.actor = definition.instantiate();
+    const connDrivers = createGenericConnDrivers(
+      handler.genericConnGlobalState
+    );
+    await handler.actor.start(
+      connDrivers,
+      this,
+      this.#inlineClient,
+      actorId,
+      name,
+      key,
+      "unknown"
+      // TODO: Support regions in Cloudflare
+    );
+    (_a = handler.actorPromise) == null ? void 0 : _a.resolve();
+    handler.actorPromise = void 0;
+    return handler.actor;
+  }
+  getGenericConnGlobalState(actorId) {
+    const handler = this.#actors.get(actorId);
+    if (!handler) {
+      throw new Error(`Actor ${actorId} not loaded`);
+    }
+    return handler.genericConnGlobalState;
+  }
   getContext(actorId) {
     const state = this.#globalState.getDOState(actorId);
     return { ctx: state.ctx, env: state.env };
@@ -51,6 +118,24 @@ var CloudflareActorsActorDriver = class {
     return this.#getDOCtx(actorId).storage.sql;
   }
 };
+function createCloudflareActorsActorDriverBuilder(globalState) {
+  return (registryConfig, runConfig, managerDriver, inlineClient) => {
+    return new CloudflareActorsActorDriver(
+      registryConfig,
+      runConfig,
+      managerDriver,
+      inlineClient,
+      globalState
+    );
+  };
+}
+
+// src/log.ts
+import { getLogger } from "@rivetkit/core/log";
+var LOGGER_NAME = "driver-cloudflare-workers";
+function logger() {
+  return getLogger(LOGGER_NAME);
+}
 
 // src/actor-handler-do.ts
 var KEYS = {
@@ -92,23 +177,27 @@ function createActorDurableObject(registry, runConfig) {
         return this.#actor;
       }
       if (!this.#initialized) throw new Error("Not initialized");
-      runConfig.driver.actor = new CloudflareActorsActorDriver(globalState);
-      const actorTopology = new PartitionTopologyActor(
+      const actorId = this.ctx.id.toString();
+      globalState.setDOState(actorId, { ctx: this.ctx, env: this.env });
+      runConfig.driver.actor = createCloudflareActorsActorDriverBuilder(globalState);
+      const managerDriver = runConfig.driver.manager(
         registry.config,
         runConfig
       );
-      const actorId = this.ctx.id.toString();
-      globalState.setDOState(actorId, { ctx: this.ctx, env: this.env });
+      const inlineClient = createClientWithDriver(
+        createInlineClientDriver(managerDriver)
+      );
+      const actorDriver = runConfig.driver.actor(
+        registry.config,
+        runConfig,
+        managerDriver,
+        inlineClient
+      );
+      const actorRouter = createActorRouter(runConfig, actorDriver);
       this.#actor = {
-        actorTopology
+        actorRouter
       };
-      await actorTopology.start(
-        actorId,
-        this.#initialized.name,
-        this.#initialized.key,
-        // TODO:
-        "unknown"
-      );
+      await actorDriver.loadActor(actorId);
       return this.#actor;
     }
     /** RPC called by the service that creates the DO to initialize it. */
@@ -129,27 +218,32 @@ function createActorDurableObject(registry, runConfig) {
     }
     async fetch(request) {
       return await CF_AMBIENT_ENV.run(this.env, async () => {
-        const { actorTopology } = await this.#loadActor();
-        const ctx = this.ctx;
-        return await actorTopology.router.fetch(
-          request,
-          this.env,
-          // Implement execution context so we can wait on requests
-          {
-            waitUntil(promise) {
-              ctx.waitUntil(promise);
-            },
-            passThroughOnException() {
-            },
-            props: {}
-          }
-        );
+        const { actorRouter } = await this.#loadActor();
+        const actorId = this.ctx.id.toString();
+        return await actorRouter.fetch(request, {
+          actorId
+        });
       });
     }
     async alarm() {
       return await CF_AMBIENT_ENV.run(this.env, async () => {
-        const { actorTopology } = await this.#loadActor();
-        await actorTopology.actor.onAlarm();
+        await this.#loadActor();
+        const actorId = this.ctx.id.toString();
+        const managerDriver = runConfig.driver.manager(
+          registry.config,
+          runConfig
+        );
+        const inlineClient = createClientWithDriver(
+          createInlineClientDriver(managerDriver)
+        );
+        const actorDriver = runConfig.driver.actor(
+          registry.config,
+          runConfig,
+          managerDriver,
+          inlineClient
+        );
+        const actor = await actorDriver.loadActor(actorId);
+        await actor.onAlarm();
       });
     }
   };
@@ -162,16 +256,12 @@ var ConfigSchema = RunConfigSchema.removeDefault().omit({ driver: true, getUpgra
   app: z.custom().optional()
 }).default({});
 
-// src/handler.ts
-import { Hono } from "hono";
-import { PartitionTopologyManager } from "@rivetkit/core/topologies/partition";
-
 // src/manager-driver.ts
 import {
-  HEADER_EXPOSE_INTERNAL_ERROR,
-  HEADER_ENCODING,
+  HEADER_AUTH_DATA,
   HEADER_CONN_PARAMS,
-  HEADER_AUTH_DATA
+  HEADER_ENCODING,
+  HEADER_EXPOSE_INTERNAL_ERROR
 } from "@rivetkit/core/driver-helpers";
 import { ActorAlreadyExists, InternalError } from "@rivetkit/core/errors";
 
@@ -221,105 +311,103 @@ var STANDARD_WEBSOCKET_HEADERS = [
   "sec-websocket-extensions"
 ];
 var CloudflareActorsManagerDriver = class {
-  connRoutingHandler;
-  constructor() {
-    this.connRoutingHandler = {
-      custom: {
-        sendRequest: async (actorId, actorRequest) => {
-          const env = getCloudflareAmbientEnv();
-          logger().debug("sending request to durable object", {
-            actorId,
-            method: actorRequest.method,
-            url: actorRequest.url
-          });
-          const id = env.ACTOR_DO.idFromString(actorId);
-          const stub = env.ACTOR_DO.get(id);
-          return await stub.fetch(actorRequest);
-        },
-        openWebSocket: async (actorId, encodingKind, params) => {
-          const env = getCloudflareAmbientEnv();
-          logger().debug("opening websocket to durable object", { actorId });
-          const id = env.ACTOR_DO.idFromString(actorId);
-          const stub = env.ACTOR_DO.get(id);
-          const headers = {
-            Upgrade: "websocket",
-            Connection: "Upgrade",
-            [HEADER_EXPOSE_INTERNAL_ERROR]: "true",
-            [HEADER_ENCODING]: encodingKind
-          };
-          if (params) {
-            headers[HEADER_CONN_PARAMS] = JSON.stringify(params);
-          }
-          headers["sec-websocket-protocol"] = "rivetkit";
-          const response = await stub.fetch("http://actor/connect/websocket", {
-            headers
-          });
-          const webSocket = response.webSocket;
-          if (!webSocket) {
-            throw new InternalError(
-              "missing websocket connection in response from DO"
-            );
-          }
-          logger().debug("durable object websocket connection open", {
-            actorId
-          });
-          webSocket.accept();
-          setTimeout(() => {
-            var _a;
-            (_a = webSocket.onopen) == null ? void 0 : _a.call(webSocket, new Event("open"));
-          }, 0);
-          return webSocket;
-        },
-        proxyRequest: async (c, actorRequest, actorId) => {
-          logger().debug("forwarding request to durable object", {
-            actorId,
-            method: actorRequest.method,
-            url: actorRequest.url
-          });
-          const id = c.env.ACTOR_DO.idFromString(actorId);
-          const stub = c.env.ACTOR_DO.get(id);
-          return await stub.fetch(actorRequest);
-        },
-        proxyWebSocket: async (c, path, actorId, encoding, params, authData) => {
-          logger().debug("forwarding websocket to durable object", {
-            actorId,
-            path
-          });
-          const upgradeHeader = c.req.header("Upgrade");
-          if (!upgradeHeader || upgradeHeader !== "websocket") {
-            return new Response("Expected Upgrade: websocket", {
-              status: 426
-            });
-          }
-          const newUrl = new URL(`http://actor${path}`);
-          const actorRequest = new Request(newUrl, c.req.raw);
-          const headerKeys = [];
-          actorRequest.headers.forEach((v, k) => headerKeys.push(k));
-          for (const k of headerKeys) {
-            if (!STANDARD_WEBSOCKET_HEADERS.includes(k)) {
-              actorRequest.headers.delete(k);
-            }
-          }
-          actorRequest.headers.set(HEADER_EXPOSE_INTERNAL_ERROR, "true");
-          actorRequest.headers.set(HEADER_ENCODING, encoding);
-          if (params) {
-            actorRequest.headers.set(
-              HEADER_CONN_PARAMS,
-              JSON.stringify(params)
-            );
-          }
-          if (authData) {
-            actorRequest.headers.set(
-              HEADER_AUTH_DATA,
-              JSON.stringify(authData)
-            );
-          }
-          const id = c.env.ACTOR_DO.idFromString(actorId);
-          const stub = c.env.ACTOR_DO.get(id);
-          return await stub.fetch(actorRequest);
-        }
-      }
+  async sendRequest(actorId, actorRequest) {
+    const env = getCloudflareAmbientEnv();
+    logger().debug("sending request to durable object", {
+      actorId,
+      method: actorRequest.method,
+      url: actorRequest.url
+    });
+    const id = env.ACTOR_DO.idFromString(actorId);
+    const stub = env.ACTOR_DO.get(id);
+    return await stub.fetch(actorRequest);
+  }
+  async openWebSocket(path, actorId, encoding, params) {
+    const env = getCloudflareAmbientEnv();
+    logger().debug("opening websocket to durable object", { actorId, path });
+    const id = env.ACTOR_DO.idFromString(actorId);
+    const stub = env.ACTOR_DO.get(id);
+    const headers = {
+      Upgrade: "websocket",
+      Connection: "Upgrade",
+      [HEADER_EXPOSE_INTERNAL_ERROR]: "true",
+      [HEADER_ENCODING]: encoding
     };
+    if (params) {
+      headers[HEADER_CONN_PARAMS] = JSON.stringify(params);
+    }
+    headers["sec-websocket-protocol"] = "rivetkit";
+    const url = `http://actor${path}`;
+    logger().debug("rewriting websocket url", {
+      from: path,
+      to: url
+    });
+    const response = await stub.fetch(url, {
+      headers
+    });
+    const webSocket = response.webSocket;
+    if (!webSocket) {
+      throw new InternalError(
+        "missing websocket connection in response from DO"
+      );
+    }
+    logger().debug("durable object websocket connection open", {
+      actorId
+    });
+    webSocket.accept();
+    setTimeout(() => {
+      var _a;
+      const event = new Event("open");
+      (_a = webSocket.onopen) == null ? void 0 : _a.call(webSocket, event);
+      webSocket.dispatchEvent(event);
+    }, 0);
+    return webSocket;
+  }
+  async proxyRequest(c, actorRequest, actorId) {
+    logger().debug("forwarding request to durable object", {
+      actorId,
+      method: actorRequest.method,
+      url: actorRequest.url
+    });
+    const id = c.env.ACTOR_DO.idFromString(actorId);
+    const stub = c.env.ACTOR_DO.get(id);
+    return await stub.fetch(actorRequest);
+  }
+  async proxyWebSocket(c, path, actorId, encoding, params, authData) {
+    logger().debug("forwarding websocket to durable object", {
+      actorId,
+      path
+    });
+    const upgradeHeader = c.req.header("Upgrade");
+    if (!upgradeHeader || upgradeHeader !== "websocket") {
+      return new Response("Expected Upgrade: websocket", {
+        status: 426
+      });
+    }
+    const newUrl = new URL(`http://actor${path}`);
+    const actorRequest = new Request(newUrl, c.req.raw);
+    logger().debug("rewriting websocket url", {
+      from: c.req.url,
+      to: actorRequest.url
+    });
+    const headerKeys = [];
+    actorRequest.headers.forEach((v, k) => headerKeys.push(k));
+    for (const k of headerKeys) {
+      if (!STANDARD_WEBSOCKET_HEADERS.includes(k)) {
+        actorRequest.headers.delete(k);
+      }
+    }
+    actorRequest.headers.set(HEADER_EXPOSE_INTERNAL_ERROR, "true");
+    actorRequest.headers.set(HEADER_ENCODING, encoding);
+    if (params) {
+      actorRequest.headers.set(HEADER_CONN_PARAMS, JSON.stringify(params));
+    }
+    if (authData) {
+      actorRequest.headers.set(HEADER_AUTH_DATA, JSON.stringify(authData));
+    }
+    const id = c.env.ACTOR_DO.idFromString(actorId);
+    const stub = c.env.ACTOR_DO.get(id);
+    return await stub.fetch(actorRequest);
   }
   async getForId({
     c,
@@ -423,7 +511,7 @@ var CloudflareActorsManagerDriver = class {
 };
 
 // src/websocket.ts
-import { WSContext, defineWebSocketHelper } from "hono/ws";
+import { defineWebSocketHelper, WSContext } from "hono/ws";
 var upgradeWebSocket = defineWebSocketHelper(async (c, events) => {
   var _a, _b;
   const upgradeHeader = c.req.header("Upgrade");
@@ -487,8 +575,6 @@ var upgradeWebSocket = defineWebSocketHelper(async (c, events) => {
 });
 
 // src/handler.ts
-import invariant2 from "invariant";
-import { AsyncLocalStorage } from "async_hooks";
 var CF_AMBIENT_ENV = new AsyncLocalStorage();
 function getCloudflareAmbientEnv() {
   const env = CF_AMBIENT_ENV.getStore();
@@ -503,30 +589,28 @@ function createServer(registry, inputConfig) {
   const config = ConfigSchema.parse(inputConfig);
   const runConfig = {
     driver: {
-      topology: "partition",
-      manager: new CloudflareActorsManagerDriver(),
-      // HACK: We can't build the actor driver until we're inside the Druable Object
+      manager: () => new CloudflareActorsManagerDriver(),
+      // HACK: We can't build the actor driver until we're inside the Durable Object
       actor: void 0
     },
     getUpgradeWebSocket: () => upgradeWebSocket,
     ...config
   };
-  const ActorHandler = createActorDurableObject(registry, runConfig);
-  const managerTopology = new PartitionTopologyManager(
-    registry.config,
-    runConfig
-  );
+  const ActorHandler2 = createActorDurableObject(registry, runConfig);
+  const serverOutput = registry.createServer(runConfig);
   return {
-    client: managerTopology.inlineClient,
+    client: serverOutput.client,
     createHandler: (hono) => {
       const app = hono ?? new Hono();
-      app.route("/registry", managerTopology.router);
+      if (!hono) {
+        app.route("/registry", serverOutput.hono);
+      }
       const handler = {
         fetch: (request, env, ctx) => {
           return CF_AMBIENT_ENV.run(env, () => app.fetch(request, env, ctx));
         }
       };
-      return { handler, ActorHandler };
+      return { handler, ActorHandler: ActorHandler2 };
     }
   };
 }