diff --git a/packages/next-js/src/mod.ts b/packages/next-js/src/mod.ts index 1cc75ee43..c418075b8 100644 --- a/packages/next-js/src/mod.ts +++ b/packages/next-js/src/mod.ts @@ -5,10 +5,10 @@ export const toNextHandler = ( inputConfig: RunConfigInput = {}, ) => { // Don't run server locally since we're using the fetch handler directly - inputConfig.disableServer = true; + inputConfig.disableDefaultServer = true; inputConfig.disableActorDriver = true; - const { fetch } = registry.startServerless(inputConfig); + const { fetch } = registry.start(inputConfig); const fetchWrapper = async ( request: Request, diff --git a/packages/rivetkit/src/client/config.ts b/packages/rivetkit/src/client/config.ts index 73a5fa915..8369152c8 100644 --- a/packages/rivetkit/src/client/config.ts +++ b/packages/rivetkit/src/client/config.ts @@ -19,8 +19,6 @@ export const ClientConfigSchema = z.object({ .optional() .transform((x) => x ?? getEnvUniversal("RIVET_TOKEN")), - totalSlots: z.number().optional(), - headers: z.record(z.string()).optional().default({}), /** Endpoint to connect to the Rivet engine. Can be configured via RIVET_ENGINE env var. */ diff --git a/packages/rivetkit/src/drivers/default.ts b/packages/rivetkit/src/drivers/default.ts index 846d2471d..23b1691ff 100644 --- a/packages/rivetkit/src/drivers/default.ts +++ b/packages/rivetkit/src/drivers/default.ts @@ -10,7 +10,13 @@ import type { DriverConfig, RunConfig } from "@/registry/run-config"; export function chooseDefaultDriver(runConfig: RunConfig): DriverConfig { if (runConfig.endpoint && runConfig.driver) { throw new UserError( - "Cannot specify both 'engine' and 'driver' in configuration", + "Cannot specify both 'endpoint' and 'driver' in configuration", + ); + } + + if (runConfig.runnerKind === "serverless" && !runConfig.endpoint) { + throw new UserError( + "Cannot use 'serverless' runnerKind without the 'endpoint' config set. This condition should not happen, since the endpoint should be configured by default by Rivet when using serverless runners.", ); } diff --git a/packages/rivetkit/src/drivers/engine/actor-driver.ts b/packages/rivetkit/src/drivers/engine/actor-driver.ts index c6db28bce..3d95b453e 100644 --- a/packages/rivetkit/src/drivers/engine/actor-driver.ts +++ b/packages/rivetkit/src/drivers/engine/actor-driver.ts @@ -92,9 +92,9 @@ export class EngineActorDriver implements ActorDriver { endpoint: config.endpoint, token: runConfig.token ?? config.token, pegboardEndpoint: config.pegboardEndpoint, - namespace: config.namespace, + namespace: runConfig.namespace ?? config.namespace, totalSlots: runConfig.totalSlots ?? config.totalSlots, - runnerName: config.runnerName, + runnerName: runConfig.runnerName ?? config.runnerName, runnerKey: config.runnerKey, metadata: { inspectorToken: this.#runConfig.inspector.token(), diff --git a/packages/rivetkit/src/inspector/utils.ts b/packages/rivetkit/src/inspector/utils.ts index 96e6bde8b..e125640eb 100644 --- a/packages/rivetkit/src/inspector/utils.ts +++ b/packages/rivetkit/src/inspector/utils.ts @@ -65,8 +65,10 @@ export function getInspectorUrl(runConfig: RunConfigInput | undefined) { url.searchParams.set("t", accessToken); - if (runConfig?.inspector?.defaultEndpoint) { - url.searchParams.set("u", runConfig.inspector.defaultEndpoint); + const overrideDefaultEndpoint = + runConfig?.inspector?.defaultEndpoint ?? runConfig.overrideServerAddress; + if (overrideDefaultEndpoint) { + url.searchParams.set("u", overrideDefaultEndpoint); } return url.href; diff --git a/packages/rivetkit/src/manager/router.ts b/packages/rivetkit/src/manager/router.ts index a1a81a153..8fbeb94af 100644 --- a/packages/rivetkit/src/manager/router.ts +++ b/packages/rivetkit/src/manager/router.ts @@ -50,6 +50,7 @@ import { type Actor as ApiActor, } from "@/manager-api/actors"; import { RivetIdSchema } from "@/manager-api/common"; +import type { ServerlessActorDriverBuilder } from "@/mod"; import type { RegistryConfig } from "@/registry/config"; import type { RunConfig } from "@/registry/run-config"; import type { ActorOutput, ManagerDriver } from "./driver"; @@ -79,7 +80,7 @@ export function createManagerRouter( registryConfig: RegistryConfig, runConfig: RunConfig, managerDriver: ManagerDriver, - serverlessActorDriverBuilder: (() => ActorDriver) | undefined, + serverlessActorDriverBuilder: ServerlessActorDriverBuilder | undefined, ): { router: Hono; openapi: OpenAPIHono } { const router = new OpenAPIHono({ strict: false }).basePath( runConfig.basePath, @@ -87,6 +88,25 @@ export function createManagerRouter( router.use("*", loggerMiddleware(logger())); + // HACK: Add Sec-WebSocket-Protocol header to fix KIT-339 + // + // Some Deno WebSocket providers do not auto-set the protocol, which + // will cause some WebSocket clients to fail + router.use( + "*", + createMiddleware(async (c, next) => { + const upgrade = c.req.header("upgrade"); + const isWebSocket = upgrade?.toLowerCase() === "websocket"; + const isGet = c.req.method === "GET"; + + if (isGet && isWebSocket) { + c.header("Sec-WebSocket-Protocol", "rivet"); + } + + await next(); + }), + ); + if (serverlessActorDriverBuilder) { addServerlessRoutes(runConfig, serverlessActorDriverBuilder, router); } else { @@ -102,10 +122,7 @@ export function createManagerRouter( function addServerlessRoutes( runConfig: RunConfig, - serverlessActorDriverBuilder: ( - token: string | undefined, - totalSlots: number | undefined, - ) => ActorDriver, + serverlessActorDriverBuilder: ServerlessActorDriverBuilder, router: OpenAPIHono, ) { // Apply CORS @@ -122,11 +139,18 @@ function addServerlessRoutes( router.get("/start", async (c) => { const token = c.req.header("x-rivet-token"); let totalSlots: number | undefined = parseInt( - c.req.header("x-rivetkit-total-slots") as any, + c.req.header("x-rivet-total-slots") as any, + ); + if (!Number.isFinite(totalSlots)) totalSlots = undefined; + const runnerName = c.req.header("x-rivet-runner-name"); + const namespace = c.req.header("x-rivet-namespace-id"); + + const actorDriver = serverlessActorDriverBuilder( + token, + totalSlots, + runnerName, + namespace, ); - if (isNaN(totalSlots)) totalSlots = undefined; - - const actorDriver = serverlessActorDriverBuilder(token, totalSlots); invariant( actorDriver.serverlessHandleStart, "missing serverlessHandleStart on ActorDriver", diff --git a/packages/rivetkit/src/registry/mod.ts b/packages/rivetkit/src/registry/mod.ts index 4f95fda7b..60ce1723b 100644 --- a/packages/rivetkit/src/registry/mod.ts +++ b/packages/rivetkit/src/registry/mod.ts @@ -24,6 +24,13 @@ import { } from "./run-config"; import { crossPlatformServe } from "./serve"; +export type ServerlessActorDriverBuilder = ( + token?: string, + totalSlots?: number, + runnerName?: string, + namespace?: string, +) => ActorDriver; + interface ServerOutput> { /** Client to communicate with the actors. */ client: Client; @@ -64,14 +71,17 @@ export class Registry { // TODO: Find cleaner way of disabling by default if (driver.name === "engine") { config.inspector.enabled = { manager: false, actor: true }; - config.disableServer = true; + config.disableDefaultServer = true; } if (driver.name === "cloudflare-workers") { config.inspector.enabled = { manager: false, actor: true }; - config.disableServer = true; + config.disableDefaultServer = true; config.disableActorDriver = true; config.noWelcome = true; } + if (config.runnerKind === "serverless") { + config.disableActorDriver = true; + } // Configure getUpgradeWebSocket lazily so we can assign it in crossPlatformServe let upgradeWebSocket: any; @@ -102,10 +112,10 @@ export class Registry { const displayInfo = managerDriver.displayInformation(); console.log(); console.log(` RivetKit ${pkg.version} (${displayInfo.name})`); - if (config.disableServer) { - console.log(` - Endpoint: (default server disabled)`); - } else { + if (!config.disableDefaultServer) { console.log(` - Endpoint: ${config.endpoint}`); + } else if (config.overrideServerAddress) { + console.log(` - Endpoint: ${config.overrideServerAddress}`); } for (const [k, v] of Object.entries(displayInfo.properties)) { const padding = " ".repeat(Math.max(0, 13 - k.length)); @@ -117,6 +127,7 @@ export class Registry { console.log(); } + let serverlessActorDriverBuilder: undefined | ServerlessActorDriverBuilder; // HACK: We need to find a better way to let the driver itself decide when to start the actor driver // Create runner // @@ -128,117 +139,22 @@ export class Registry { managerDriver, client, ); - } - - const { router: hono } = createManagerRouter( - this.#config, - config, - managerDriver, - undefined, - ); - - // Start server - if (!config.disableServer) { - (async () => { - const out = await crossPlatformServe(hono, undefined); - upgradeWebSocket = out.upgradeWebSocket; - })(); - } - - return { - client, - fetch: hono.fetch.bind(hono), - }; - } - - public startServerless(inputConfig?: RunConfigInput): ServerOutput { - const config = RunConfigSchema.parse(inputConfig); - - // Configure logger - if (config.logging?.baseLogger) { - // Use provided base logger - configureBaseLogger(config.logging.baseLogger); } else { - // Configure default logger with log level from config - // getPinoLevel will handle env variable priority - configureDefaultLogger(config.logging?.level); - } - - // Choose the driver based on configuration - const driver = chooseDefaultDriver(config); - - // TODO: Find cleaner way of disabling by default - if (driver.name === "engine") { - config.inspector.enabled = false; - config.disableServer = true; - config.disableActorDriver = true; - } - if (driver.name === "cloudflare-workers") { - config.inspector.enabled = false; - config.disableServer = true; - config.disableActorDriver = true; - config.noWelcome = true; - } - - // Configure getUpgradeWebSocket lazily so we can assign it in crossPlatformServe - let upgradeWebSocket: any; - if (!config.getUpgradeWebSocket) { - config.getUpgradeWebSocket = () => upgradeWebSocket!; - } - - // Create router - const managerDriver = driver.manager(this.#config, config); - - // Create client - const client = createClientWithDriver(managerDriver, config); - - const driverLog = managerDriver.extraStartupLog?.() ?? {}; - logger().info({ - msg: "rivetkit ready", - driver: driver.name, - definitions: Object.keys(this.#config.use).length, - ...driverLog, - }); - if (config.inspector?.enabled && managerDriver.inspector) { - logger().info({ msg: "inspector ready", url: getInspectorUrl(config) }); - } - - // Print welcome information - if (!config.noWelcome) { - const displayInfo = managerDriver.displayInformation(); - console.log(); - console.log(` RivetKit ${pkg.version} (${displayInfo.name})`); - console.log(` - Endpoint: http://127.0.0.1:6420`); - for (const [k, v] of Object.entries(displayInfo.properties)) { - const padding = " ".repeat(Math.max(0, 13 - k.length)); - console.log(` - ${k}:${padding}${v}`); - } - if (config.inspector?.enabled && managerDriver.inspector) { - console.log(` - Inspector: ${getInspectorUrl(config)}`); - } - console.log(); - } - - let serverlessActorDriverBuilder: - | ((token?: string, totalSlots?: number) => ActorDriver) - | undefined = ( - token: string | undefined, - totalSlots: number | undefined, - ) => { - // Override config - if (token) config.token = token; - if (totalSlots) config.totalSlots = totalSlots; - - return driver.actor(this.#config, config, managerDriver, client); - }; - - // HACK: We need to find a better way to let the driver itself decide when to start the actor driver - // Create runner - // - // Even though we do not use the return value, this is required to start the code that will handle incoming actors - if (!config.disableActorDriver) { - const _actorDriver = serverlessActorDriverBuilder(); - serverlessActorDriverBuilder = undefined; + serverlessActorDriverBuilder = ( + token, + totalSlots, + runnerName, + namespace, + ) => { + // Override config + if (token) config.token = token; + if (totalSlots) config.totalSlots = totalSlots; + if (runnerName) config.runnerName = runnerName; + if (namespace) config.namespace = namespace; + + // Create new actor driver with updated config + return driver.actor(this.#config, config, managerDriver, client); + }; } const { router: hono } = createManagerRouter( @@ -249,7 +165,7 @@ export class Registry { ); // Start server - if (!config.disableServer) { + if (!config.disableDefaultServer) { (async () => { const out = await crossPlatformServe(hono, undefined); upgradeWebSocket = out.upgradeWebSocket; diff --git a/packages/rivetkit/src/registry/run-config.ts b/packages/rivetkit/src/registry/run-config.ts index 8bf13efe9..d370f8d66 100644 --- a/packages/rivetkit/src/registry/run-config.ts +++ b/packages/rivetkit/src/registry/run-config.ts @@ -6,6 +6,7 @@ import { ClientConfigSchema } from "@/client/config"; import { LogLevelSchema } from "@/common/log"; import { InspectorConfigSchema } from "@/inspector/config"; import type { ManagerDriverBuilder } from "@/manager/driver"; +import { getEnvUniversal } from "@/utils"; type CorsOptions = NonNullable[0]>; @@ -30,11 +31,30 @@ export const RunConfigSchema = ClientConfigSchema.extend({ inspector: InspectorConfigSchema, /** @experimental */ - disableServer: z.boolean().optional().default(false), + disableDefaultServer: z.boolean().optional().default(false), + + /** @experimental */ + overrideServerAddress: z.string().optional(), /** @experimental */ disableActorDriver: z.boolean().optional().default(false), + /** + * @experimental + * + * Whether to run runners normally or have them managed + * serverlessly (by the Rivet Engine for example). + */ + runnerKind: z + .enum(["serverless", "normal"]) + .optional() + .default(() => + getEnvUniversal("RIVET_RUNNER_KIND") === "serverless" + ? "serverless" + : "normal", + ), + totalSlots: z.number().optional(), + /** * @experimental *