rivetkit 2.0.24-rc.1 → 2.0.24

package/src/actor/router.ts

@@ -1,33 +1,16 @@
-import { Hono, type Context as HonoContext } from "hono";
+import { Hono } from "hono";
 import invariant from "invariant";
-import { EncodingSchema } from "@/actor/protocol/serde";
 import {
 	type ActionOpts,
 	type ActionOutput,
-	type ConnectSseOpts,
-	type ConnectSseOutput,
-	type ConnectWebSocketOpts,
-	type ConnectWebSocketOutput,
 	type ConnsMessageOpts,
 	handleAction,
-	handleConnectionClose,
-	handleConnectionMessage,
-	handleRawWebSocketHandler,
-	handleSseConnect,
-	handleWebSocketConnect,
+	handleRawRequest,
 } from "@/actor/router-endpoints";
 import {
-	HEADER_CONN_ID,
-	HEADER_CONN_PARAMS,
-	HEADER_CONN_TOKEN,
-	HEADER_ENCODING,
-	PATH_CONNECT_WEBSOCKET,
-	PATH_RAW_WEBSOCKET_PREFIX,
-	WS_PROTOCOL_CONN_ID,
-	WS_PROTOCOL_CONN_PARAMS,
-	WS_PROTOCOL_CONN_TOKEN,
-	WS_PROTOCOL_ENCODING,
-	WS_PROTOCOL_TOKEN,
+	PATH_CONNECT,
+	PATH_INSPECTOR_CONNECT,
+	PATH_WEBSOCKET_PREFIX,
 } from "@/common/actor-router-consts";
 import {
 	handleRouteError,
@@ -41,21 +24,15 @@ import {
 } from "@/inspector/actor";
 import { isInspectorEnabled, secureInspector } from "@/inspector/utils";
 import type { RunnerConfig } from "@/registry/run-config";
-import { generateConnRequestId } from "./conn";
-import { ConnDriverKind } from "./conn-drivers";
+import { CONN_DRIVER_SYMBOL } from "./conn/mod";
 import type { ActorDriver } from "./driver";
-import { InternalError } from "./errors";
 import { loggerWithoutContext } from "./log";
+import {
+	parseWebSocketProtocols,
+	routeWebSocket,
+} from "./router-websocket-endpoints";
 
-export type {
-	ConnectWebSocketOpts,
-	ConnectWebSocketOutput,
-	ConnectSseOpts,
-	ConnectSseOutput,
-	ActionOpts,
-	ActionOutput,
-	ConnsMessageOpts,
-};
+export type { ActionOpts, ActionOutput, ConnsMessageOpts };
 
 interface ActorRouterBindings {
 	actorId: string;
@@ -80,11 +57,11 @@ export function createActorRouter(
 	// Track all HTTP requests to prevent actor from sleeping during active requests
 	router.use("*", async (c, next) => {
 		const actor = await actorDriver.loadActor(c.env.actorId);
-		actor.__beginHonoHttpRequest();
+		actor.beginHonoHttpRequest();
 		try {
 			await next();
 		} finally {
-			actor.__endHonoHttpRequest();
+			actor.endHonoHttpRequest();
 		}
 	});
 
@@ -108,101 +85,60 @@ export function createActorRouter(
 			}
 
 			const actor = await actorDriver.loadActor(c.env.actorId);
-			const conn = actor.__getConnForId(connId);
+			const conn = actor.connectionManager.getConnForId(connId);
 
 			if (!conn) {
 				return c.text(`Connection not found: ${connId}`, 404);
 			}
 
-			// Force close the websocket/SSE connection without clean shutdown
-			const driverState = conn.__driverState;
-			if (driverState && ConnDriverKind.WEBSOCKET in driverState) {
-				const ws = driverState[ConnDriverKind.WEBSOCKET].websocket;
-
-				// Force close without sending close frame
-				(ws.raw as any).terminate();
-			} else if (driverState && ConnDriverKind.SSE in driverState) {
-				const stream = driverState[ConnDriverKind.SSE].stream;
-
-				// Force close the SSE stream
-				stream.abort();
+			// Force close the connection without clean shutdown
+			if (conn[CONN_DRIVER_SYMBOL]?.terminate) {
+				conn[CONN_DRIVER_SYMBOL].terminate(actor, conn);
 			}
 
 			return c.json({ success: true });
 		});
 	}
 
-	router.get(PATH_CONNECT_WEBSOCKET, async (c) => {
-		const upgradeWebSocket = runConfig.getUpgradeWebSocket?.();
-		if (upgradeWebSocket) {
-			return upgradeWebSocket(async (c) => {
-				// Parse configuration from Sec-WebSocket-Protocol header
-				const protocols = c.req.header("sec-websocket-protocol");
-				let encodingRaw: string | undefined;
-				let connParamsRaw: string | undefined;
-				let connIdRaw: string | undefined;
-				let connTokenRaw: string | undefined;
-
-				if (protocols) {
-					const protocolList = protocols
-						.split(",")
-						.map((p) => p.trim());
-					for (const protocol of protocolList) {
-						if (protocol.startsWith(WS_PROTOCOL_ENCODING)) {
-							encodingRaw = protocol.substring(
-								WS_PROTOCOL_ENCODING.length,
-							);
-						} else if (
-							protocol.startsWith(WS_PROTOCOL_CONN_PARAMS)
-						) {
-							connParamsRaw = decodeURIComponent(
-								protocol.substring(
-									WS_PROTOCOL_CONN_PARAMS.length,
-								),
-							);
-						} else if (protocol.startsWith(WS_PROTOCOL_CONN_ID)) {
-							connIdRaw = protocol.substring(
-								WS_PROTOCOL_CONN_ID.length,
-							);
-						} else if (
-							protocol.startsWith(WS_PROTOCOL_CONN_TOKEN)
-						) {
-							connTokenRaw = protocol.substring(
-								WS_PROTOCOL_CONN_TOKEN.length,
-							);
-						}
-					}
-				}
-
-				const encoding = EncodingSchema.parse(encodingRaw);
-				const connParams = connParamsRaw
-					? JSON.parse(connParamsRaw)
-					: undefined;
-
-				return await handleWebSocketConnect(
-					c.req.raw,
-					runConfig,
-					actorDriver,
-					c.env.actorId,
-					encoding,
-					connParams,
-					generateConnRequestId(),
-					undefined,
-					connIdRaw,
-					connTokenRaw,
+	// Route all WebSocket paths using the same handler
+	//
+	// All WebSockets use a separate underlying router in routeWebSocket since
+	// WebSockets also need to be routed from ManagerDriver.proxyWebSocket and
+	// ManagerDriver.openWebSocket.
+	router.on(
+		"GET",
+		[PATH_CONNECT, `${PATH_WEBSOCKET_PREFIX}*`, PATH_INSPECTOR_CONNECT],
+		async (c) => {
+			const upgradeWebSocket = runConfig.getUpgradeWebSocket?.();
+			if (upgradeWebSocket) {
+				return upgradeWebSocket(async (c) => {
+					const protocols = c.req.header("sec-websocket-protocol");
+					const { encoding, connParams } =
+						parseWebSocketProtocols(protocols);
+
+					return await routeWebSocket(
+						c.req.raw,
+						c.req.path,
+						c.req.header(),
+						runConfig,
+						actorDriver,
+						c.env.actorId,
+						encoding,
+						connParams,
+						undefined,
+						undefined,
+						false,
+						false,
+					);
+				})(c, noopNext());
+			} else {
+				return c.text(
+					"WebSockets are not enabled for this driver.",
+					400,
 				);
-			})(c, noopNext());
-		} else {
-			return c.text(
-				"WebSockets are not enabled for this driver. Use SSE instead.",
-				400,
-			);
-		}
-	});
-
-	router.get("/connect/sse", async (c) => {
-		return handleSseConnect(c, runConfig, actorDriver, c.env.actorId);
-	});
+			}
+		},
+	);
 
 	router.post("/action/:action", async (c) => {
 		const actionName = c.req.param("action");
@@ -216,46 +152,11 @@ export function createActorRouter(
 		);
 	});
 
-	router.post("/connections/message", async (c) => {
-		const connId = c.req.header(HEADER_CONN_ID);
-		const connToken = c.req.header(HEADER_CONN_TOKEN);
-		if (!connId || !connToken) {
-			throw new Error("Missing required parameters");
-		}
-		return handleConnectionMessage(
-			c,
-			runConfig,
-			actorDriver,
-			connId,
-			connToken,
-			c.env.actorId,
-		);
-	});
-
-	router.post("/connections/close", async (c) => {
-		const connId = c.req.header(HEADER_CONN_ID);
-		const connToken = c.req.header(HEADER_CONN_TOKEN);
-		if (!connId || !connToken) {
-			throw new Error("Missing required parameters");
-		}
-		return handleConnectionClose(
-			c,
-			runConfig,
-			actorDriver,
-			connId,
-			connToken,
-			c.env.actorId,
-		);
-	});
-
-	// Raw HTTP endpoints - /http/*
-	router.all("/raw/http/*", async (c) => {
-		const actor = await actorDriver.loadActor(c.env.actorId);
-
+	router.all("/request/*", async (c) => {
 		// TODO: This is not a clean way of doing this since `/http/` might exist mid-path
 		// Strip the /http prefix from the URL to get the original path
 		const url = new URL(c.req.url);
-		const originalPath = url.pathname.replace(/^\/raw\/http/, "") || "/";
+		const originalPath = url.pathname.replace(/^\/request/, "") || "/";
 
 		// Create a new request with the corrected URL
 		const correctedUrl = new URL(originalPath + url.search, url.origin);
@@ -272,47 +173,12 @@ export function createActorRouter(
 			to: correctedRequest.url,
 		});
 
-		// Call the actor's onFetch handler - it will throw appropriate errors
-		const response = await actor.handleFetch(correctedRequest, {});
-
-		// This should never happen now since handleFetch throws errors
-		if (!response) {
-			throw new InternalError("handleFetch returned void unexpectedly");
-		}
-
-		return response;
-	});
-
-	// Raw WebSocket endpoint - /websocket/*
-	router.get(`${PATH_RAW_WEBSOCKET_PREFIX}*`, async (c) => {
-		const upgradeWebSocket = runConfig.getUpgradeWebSocket?.();
-		if (upgradeWebSocket) {
-			return upgradeWebSocket(async (c) => {
-				const url = new URL(c.req.url);
-				const pathWithQuery = c.req.path + url.search;
-
-				loggerWithoutContext().debug({
-					msg: "actor router raw websocket",
-					path: c.req.path,
-					url: c.req.url,
-					search: url.search,
-					pathWithQuery,
-				});
-
-				return await handleRawWebSocketHandler(
-					c.req.raw,
-					pathWithQuery,
-					actorDriver,
-					c.env.actorId,
-					undefined,
-				);
-			})(c, noopNext());
-		} else {
-			return c.text(
-				"WebSockets are not enabled for this driver. Use SSE instead.",
-				400,
-			);
-		}
+		return await handleRawRequest(
+			c,
+			correctedRequest,
+			actorDriver,
+			c.env.actorId,
+		);
 	});
 
 	if (isInspectorEnabled(runConfig, "actor")) {

package/src/actor/schedule.ts

@@ -1,4 +1,4 @@
-import type { AnyActorInstance } from "./instance";
+import type { AnyActorInstance } from "./instance/mod";
 
 export class Schedule {
 	#actor: AnyActorInstance;