rivetkit 2.0.23 → 2.0.24

package/src/actor/router-endpoints.ts

@@ -1,78 +1,40 @@
 import * as cbor from "cbor-x";
 import type { Context as HonoContext, HonoRequest } from "hono";
-import { type SSEStreamingApi, streamSSE } from "hono/streaming";
-import type { WSContext } from "hono/ws";
-import invariant from "invariant";
-import { ActionContext } from "@/actor/action";
-import type { AnyConn } from "@/actor/conn";
-import {
-	generateConnId,
-	generateConnRequestId,
-	generateConnToken,
-} from "@/actor/conn";
-import { ConnDriverKind } from "@/actor/conn-drivers";
+import type { AnyConn } from "@/actor/conn/mod";
+import { ActionContext } from "@/actor/contexts/action";
 import * as errors from "@/actor/errors";
-import { type AnyActorInstance, PERSIST_SYMBOL } from "@/actor/instance";
-import type { InputData } from "@/actor/protocol/serde";
+import type { AnyActorInstance } from "@/actor/instance/mod";
 import { type Encoding, EncodingSchema } from "@/actor/protocol/serde";
 import {
 	HEADER_ACTOR_QUERY,
-	HEADER_CONN_ID,
 	HEADER_CONN_PARAMS,
-	HEADER_CONN_TOKEN,
 	HEADER_ENCODING,
+	WS_PROTOCOL_CONN_PARAMS,
+	WS_PROTOCOL_ENCODING,
 } from "@/common/actor-router-consts";
-import type { UpgradeWebSocketArgs } from "@/common/inline-websocket-adapter2";
-import { deconstructError, stringifyError } from "@/common/utils";
-import type { UniversalWebSocket } from "@/common/websocket-interface";
-import { HonoWebSocketAdapter } from "@/manager/hono-websocket-adapter";
+import { stringifyError } from "@/common/utils";
 import type { RunnerConfig } from "@/registry/run-config";
 import type * as protocol from "@/schemas/client-protocol/mod";
 import {
 	HTTP_ACTION_REQUEST_VERSIONED,
 	HTTP_ACTION_RESPONSE_VERSIONED,
-	TO_SERVER_VERSIONED,
 } from "@/schemas/client-protocol/versioned";
+import {
+	type HttpActionRequest as HttpActionRequestJson,
+	HttpActionRequestSchema,
+	type HttpActionResponse as HttpActionResponseJson,
+	HttpActionResponseSchema,
+} from "@/schemas/client-protocol-zod/mod";
 import {
 	contentTypeForEncoding,
 	deserializeWithEncoding,
 	serializeWithEncoding,
 } from "@/serde";
-import {
-	arrayBuffersEqual,
-	bufferToArrayBuffer,
-	promiseWithResolvers,
-} from "@/utils";
+import { bufferToArrayBuffer } from "@/utils";
+import { createHttpDriver } from "./conn/drivers/http";
+import { createRawRequestDriver } from "./conn/drivers/raw-request";
 import type { ActorDriver } from "./driver";
 import { loggerWithoutContext } from "./log";
-import { parseMessage } from "./protocol/old";
-
-export const SSE_PING_INTERVAL = 1000;
-
-export interface ConnectWebSocketOpts {
-	req?: HonoRequest;
-	encoding: Encoding;
-	actorId: string;
-	params: unknown;
-}
-
-export interface ConnectWebSocketOutput {
-	onOpen: (ws: WSContext) => void;
-	onMessage: (message: protocol.ToServer) => void;
-	onClose: () => void;
-}
-
-export interface ConnectSseOpts {
-	req?: HonoRequest;
-	encoding: Encoding;
-	params: unknown;
-	actorId: string;
-}
-
-export interface ConnectSseOutput {
-	onOpen: (stream: SSEStreamingApi) => void;
-	onClose: () => Promise<void>;
-}
 
 export interface ActionOpts {
 	req?: HonoRequest;
@@ -89,7 +51,6 @@ export interface ActionOutput {
 export interface ConnsMessageOpts {
 	req?: HonoRequest;
 	connId: string;
-	connToken: string;
 	message: protocol.ToServer;
 	actorId: string;
 }
@@ -99,362 +60,6 @@ export interface FetchOpts {
 	actorId: string;
 }
 
-export interface WebSocketOpts {
-	request: Request;
-	websocket: UniversalWebSocket;
-	actorId: string;
-}
-
-/**
- * Creates a WebSocket connection handler
- */
-export async function handleWebSocketConnect(
-	req: Request | undefined,
-	runConfig: RunnerConfig,
-	actorDriver: ActorDriver,
-	actorId: string,
-	encoding: Encoding,
-	parameters: unknown,
-	requestId: string,
-	requestIdBuf: ArrayBuffer | undefined,
-	connId: string | undefined,
-	connToken: string | undefined,
-): Promise<UpgradeWebSocketArgs> {
-	const exposeInternalError = req
-		? getRequestExposeInternalError(req)
-		: false;
-
-	// Setup promise for the init handlers since all other behavior depends on this
-	const {
-		promise: handlersPromise,
-		resolve: handlersResolve,
-		reject: handlersReject,
-	} = promiseWithResolvers<{
-		conn: AnyConn;
-		actor: AnyActorInstance;
-		connId: string;
-	}>();
-
-	// Pre-load the actor to catch errors early
-	let actor: AnyActorInstance;
-	try {
-		actor = await actorDriver.loadActor(actorId);
-	} catch (error) {
-		// Return handler that immediately closes with error
-		return {
-			onOpen: (_evt: any, ws: WSContext) => {
-				const { code } = deconstructError(
-					error,
-					actor.rLog,
-					{
-						wsEvent: "open",
-					},
-					exposeInternalError,
-				);
-				ws.close(1011, code);
-			},
-			onMessage: (_evt: { data: any }, ws: WSContext) => {
-				ws.close(1011, "Actor not loaded");
-			},
-			onClose: (_event: any, _ws: WSContext) => {},
-			onError: (_error: unknown) => {},
-		};
-	}
-
-	// Promise used to wait for the websocket close in `disconnect`
-	const closePromise = promiseWithResolvers<void>();
-
-	// Track connection outside of scope for cleanup
-	let createdConn: AnyConn | undefined;
-
-	return {
-		onOpen: (_evt: any, ws: WSContext) => {
-			actor.rLog.debug("actor websocket open");
-
-			// Run async operations in background
-			(async () => {
-				try {
-					let conn: AnyConn;
-
-					// Create or reconnect connection
-					actor.rLog.debug({
-						msg: connId
-							? "websocket reconnection attempt"
-							: "new websocket connection",
-						connId,
-						actorId,
-					});
-
-					// Check if this is a hibernatable websocket
-					const isHibernatable =
-						!!requestIdBuf &&
-						actor[PERSIST_SYMBOL].hibernatableWebSocket.findIndex(
-							(ws) =>
-								arrayBuffersEqual(ws.requestId, requestIdBuf),
-						) !== -1;
-
-					conn = await actor.createConn(
-						{
-							requestId: requestId,
-							requestIdBuf: requestIdBuf,
-							hibernatable: isHibernatable,
-							driverState: {
-								[ConnDriverKind.WEBSOCKET]: {
-									encoding,
-									websocket: ws,
-									closePromise,
-								},
-							},
-						},
-						parameters,
-						req,
-						connId,
-						connToken,
-					);
-
-					// Store connection so we can clean on close
-					createdConn = conn;
-
-					// Unblock other handlers
-					handlersResolve({ conn, actor, connId: conn.id });
-				} catch (error) {
-					handlersReject(error);
-
-					const { code } = deconstructError(
-						error,
-						actor.rLog,
-						{
-							wsEvent: "open",
-						},
-						exposeInternalError,
-					);
-					ws.close(1011, code);
-				}
-			})();
-		},
-		onMessage: (evt: { data: any }, ws: WSContext) => {
-			// Handle message asynchronously
-			handlersPromise
-				.then(({ conn, actor }) => {
-					actor.rLog.debug({ msg: "received message" });
-
-					const value = evt.data.valueOf() as InputData;
-					parseMessage(value, {
-						encoding: encoding,
-						maxIncomingMessageSize:
-							runConfig.maxIncomingMessageSize,
-					})
-						.then((message) => {
-							actor
-								.processMessage(message, conn)
-								.catch((error) => {
-									const { code } = deconstructError(
-										error,
-										actor.rLog,
-										{
-											wsEvent: "message",
-										},
-										exposeInternalError,
-									);
-									ws.close(1011, code);
-								});
-						})
-						.catch((error) => {
-							const { code } = deconstructError(
-								error,
-								actor.rLog,
-								{
-									wsEvent: "message",
-								},
-								exposeInternalError,
-							);
-							ws.close(1011, code);
-						});
-				})
-				.catch((error) => {
-					const { code } = deconstructError(
-						error,
-						actor.rLog,
-						{
-							wsEvent: "message",
-						},
-						exposeInternalError,
-					);
-					ws.close(1011, code);
-				});
-		},
-		onClose: (
-			event: {
-				wasClean: boolean;
-				code: number;
-				reason: string;
-			},
-			ws: WSContext,
-		) => {
-			handlersReject(`WebSocket closed (${event.code}): ${event.reason}`);
-
-			closePromise.resolve();
-
-			if (event.wasClean) {
-				actor.rLog.info({
-					msg: "websocket closed",
-					code: event.code,
-					reason: event.reason,
-					wasClean: event.wasClean,
-				});
-			} else {
-				actor.rLog.warn({
-					msg: "websocket closed",
-					code: event.code,
-					reason: event.reason,
-					wasClean: event.wasClean,
-				});
-			}
-
-			// HACK: Close socket in order to fix bug with Cloudflare leaving WS in closing state
-			// https://github.com/cloudflare/workerd/issues/2569
-			ws.close(1000, "hack_force_close");
-
-			// Wait for actor.createConn to finish before removing the connection
-			handlersPromise.finally(() => {
-				if (createdConn) {
-					const wasClean = event.wasClean || event.code === 1000;
-					actor.__connDisconnected(createdConn, wasClean, requestId);
-				}
-			});
-		},
-		onError: (_error: unknown) => {
-			try {
-				// Actors don't need to know about this, since it's abstracted away
-				actor.rLog.warn({ msg: "websocket error" });
-			} catch (error) {
-				deconstructError(
-					error,
-					actor.rLog,
-					{ wsEvent: "error" },
-					exposeInternalError,
-				);
-			}
-		},
-	};
-}
-
-/**
- * Creates an SSE connection handler
- */
-export async function handleSseConnect(
-	c: HonoContext,
-	_runConfig: RunnerConfig,
-	actorDriver: ActorDriver,
-	actorId: string,
-) {
-	c.header("Content-Encoding", "Identity");
-
-	const encoding = getRequestEncoding(c.req);
-	const parameters = getRequestConnParams(c.req);
-	const requestId = generateConnRequestId();
-
-	// Check for reconnection parameters
-	const connId = c.req.header(HEADER_CONN_ID);
-	const connToken = c.req.header(HEADER_CONN_TOKEN);
-
-	// Return the main handler with all async work inside
-	return streamSSE(c, async (stream) => {
-		let actor: AnyActorInstance | undefined;
-		let conn: AnyConn | undefined;
-
-		try {
-			// Do all async work inside the handler
-			actor = await actorDriver.loadActor(actorId);
-
-			// Create or reconnect connection
-			actor.rLog.debug({
-				msg: connId ? "sse reconnection attempt" : "sse open",
-				connId,
-			});
-
-			conn = await actor.createConn(
-				{
-					requestId: requestId,
-					hibernatable: false,
-					driverState: {
-						[ConnDriverKind.SSE]: {
-							encoding,
-							stream: stream,
-						},
-					},
-				},
-				parameters,
-				c.req.raw,
-				connId,
-				connToken,
-			);
-
-			// Wait for close
-			const abortResolver = promiseWithResolvers();
-
-			// HACK: This is required so the abort handler below works
-			//
-			// See https://github.com/honojs/hono/issues/1770#issuecomment-2461966225
-			stream.onAbort(() => {});
-
-			// Handle stream abort (when client closes the connection)
-			c.req.raw.signal.addEventListener("abort", async () => {
-				invariant(actor, "actor should exist");
-				const rLog = actor.rLog ?? loggerWithoutContext();
-				try {
-					rLog.debug("sse stream aborted");
-
-					// Cleanup
-					if (conn) {
-						actor.__connDisconnected(conn, false, requestId);
-					}
-
-					abortResolver.resolve(undefined);
-				} catch (error) {
-					rLog.error({ msg: "error closing sse connection", error });
-					abortResolver.resolve(undefined);
-				}
-			});
-
-			// // HACK: Will throw if not configured
-			// try {
-			// 	c.executionCtx.waitUntil(abortResolver.promise);
-			// } catch {}
-
-			// Send ping every second to keep the connection alive
-			//
-			// NOTE: This is required on Cloudflare Workers in order to detect when the connection is closed
-			while (true) {
-				if (stream.closed || stream.aborted) {
-					actor?.rLog.debug({
-						msg: "sse stream closed",
-						closed: stream.closed,
-						aborted: stream.aborted,
-					});
-					break;
-				}
-
-				await stream.writeSSE({ event: "ping", data: "" });
-				await stream.sleep(SSE_PING_INTERVAL);
-			}
-		} catch (error) {
-			loggerWithoutContext().error({
-				msg: "error in sse connection",
-				error,
-			});
-
-			// Cleanup on error
-			if (conn && actor !== undefined) {
-				actor.__connDisconnected(conn, false, requestId);
-			}
-
-			// Close the stream on error
-			stream.close();
-		}
-	});
-}
-
 /**
  * Creates an action handler
  */
@@ -474,9 +79,14 @@ export async function handleAction(
 		encoding,
 		new Uint8Array(arrayBuffer),
 		HTTP_ACTION_REQUEST_VERSIONED,
+		HttpActionRequestSchema,
+		// JSON: args is already the decoded value (raw object/array)
+		(json: HttpActionRequestJson) => json.args,
+		// BARE/CBOR: args is ArrayBuffer that needs CBOR-decoding
+		(bare: protocol.HttpActionRequest) =>
+			cbor.decode(new Uint8Array(bare.args)),
 	);
-	const actionArgs = cbor.decode(new Uint8Array(request.args));
-	const requestId = generateConnRequestId();
+	const actionArgs = request;
 
 	// Invoke the action
 	let actor: AnyActorInstance | undefined;
@@ -488,198 +98,82 @@ export async function handleAction(
 		actor.rLog.debug({ msg: "handling action", actionName, encoding });
 
 		// Create conn
-		conn = await actor.createConn(
-			{
-				requestId: requestId,
-				hibernatable: false,
-				driverState: { [ConnDriverKind.HTTP]: {} },
-			},
+		conn = await actor.connectionManager.prepareAndConnectConn(
+			createHttpDriver(),
 			parameters,
 			c.req.raw,
+			c.req.path,
+			c.req.header(),
 		);
 
 		// Call action
-		const ctx = new ActionContext(actor.actorContext!, conn!);
+		const ctx = new ActionContext(actor, conn!);
 		output = await actor.executeAction(ctx, actionName, actionArgs);
 	} finally {
 		if (conn) {
-			// HTTP connections don't have persistent sockets, so no socket ID needed
-			actor?.__connDisconnected(conn, true, requestId);
+			conn.disconnect();
 		}
 	}
 
 	// Send response
-	const responseData: protocol.HttpActionResponse = {
-		output: bufferToArrayBuffer(cbor.encode(output)),
-	};
 	const serialized = serializeWithEncoding(
 		encoding,
-		responseData,
+		output,
 		HTTP_ACTION_RESPONSE_VERSIONED,
+		HttpActionResponseSchema,
+		// JSON: output is the raw value (will be serialized by jsonStringifyCompat)
+		(value): HttpActionResponseJson => ({ output: value }),
+		// BARE/CBOR: output needs to be CBOR-encoded to ArrayBuffer
+		(value): protocol.HttpActionResponse => ({
+			output: bufferToArrayBuffer(cbor.encode(value)),
+		}),
 	);
 
-	// TODO: Remvoe any, Hono is being a dumbass
+	// TODO: Remove any, Hono is being a dumbass
 	return c.body(serialized as Uint8Array as any, 200, {
 		"Content-Type": contentTypeForEncoding(encoding),
 	});
 }
 
-/**
- * Create a connection message handler
- */
-export async function handleConnectionMessage(
+export async function handleRawRequest(
 	c: HonoContext,
-	_runConfig: RunnerConfig,
+	req: Request,
 	actorDriver: ActorDriver,
-	connId: string,
-	connToken: string,
 	actorId: string,
-) {
-	const encoding = getRequestEncoding(c.req);
-
-	// Validate incoming request
-	const arrayBuffer = await c.req.arrayBuffer();
-	const message = deserializeWithEncoding(
-		encoding,
-		new Uint8Array(arrayBuffer),
-		TO_SERVER_VERSIONED,
-	);
-
+): Promise<Response> {
 	const actor = await actorDriver.loadActor(actorId);
+	const parameters = getRequestConnParams(c.req);
 
-	// Find connection
-	const conn = actor.conns.get(connId);
-	if (!conn) {
-		throw new errors.ConnNotFound(connId);
-	}
-
-	// Authenticate connection
-	if (conn._token !== connToken) {
-		throw new errors.IncorrectConnToken();
-	}
-
-	// Process message
-	await actor.processMessage(message, conn);
-
-	return c.json({});
-}
-
-export async function handleConnectionClose(
-	c: HonoContext,
-	_runConfig: RunnerConfig,
-	actorDriver: ActorDriver,
-	connId: string,
-	connToken: string,
-	actorId: string,
-) {
-	const actor = await actorDriver.loadActor(actorId);
-
-	// Find connection
-	const conn = actor.conns.get(connId);
-	if (!conn) {
-		throw new errors.ConnNotFound(connId);
-	}
-
-	// Authenticate connection
-	if (conn._token !== connToken) {
-		throw new errors.IncorrectConnToken();
-	}
+	// Track connection outside of scope for cleanup
+	let createdConn: AnyConn | undefined;
 
-	// Check if this is an SSE connection
-	if (
-		!conn.__socket?.driverState ||
-		!(ConnDriverKind.SSE in conn.__socket.driverState)
-	) {
-		throw new errors.UserError(
-			"Connection close is only supported for SSE connections",
+	try {
+		const conn = await actor.connectionManager.prepareAndConnectConn(
+			createRawRequestDriver(),
+			parameters,
+			req,
+			c.req.path,
+			c.req.header(),
 		);
-	}
-
-	// Close the SSE connection
-	await conn.disconnect("Connection closed by client request");
 
-	return c.json({});
-}
+		createdConn = conn;
 
-export async function handleRawWebSocketHandler(
-	req: Request | undefined,
-	path: string,
-	actorDriver: ActorDriver,
-	actorId: string,
-	requestIdBuf: ArrayBuffer | undefined,
-): Promise<UpgradeWebSocketArgs> {
-	const actor = await actorDriver.loadActor(actorId);
-
-	// Return WebSocket event handlers
-	return {
-		onOpen: (evt: any, ws: any) => {
-			// Extract rivetRequestId provided by engine runner
-			const rivetRequestId = evt?.rivetRequestId;
-			const isHibernatable =
-				actor[PERSIST_SYMBOL].hibernatableWebSocket.findIndex((ws) =>
-					arrayBuffersEqual(ws.requestId, rivetRequestId),
-				) !== -1;
-
-			// Wrap the Hono WebSocket in our adapter
-			const adapter = new HonoWebSocketAdapter(
-				ws,
-				rivetRequestId,
-				isHibernatable,
-			);
-
-			// Store adapter reference on the WebSocket for event handlers
-			(ws as any).__adapter = adapter;
-
-			const newPath = truncateRawWebSocketPathPrefix(path);
-			let newRequest: Request;
-			if (req) {
-				newRequest = new Request(`http://actor${newPath}`, req);
-			} else {
-				newRequest = new Request(`http://actor${newPath}`, {
-					method: "GET",
-				});
-			}
-
-			actor.rLog.debug({
-				msg: "rewriting websocket url",
-				fromPath: path,
-				toUrl: newRequest.url,
-			});
-
-			// Call the actor's onWebSocket handler with the adapted WebSocket
-			actor.handleWebSocket(adapter, {
-				request: newRequest,
-			});
-		},
-		onMessage: (event: any, ws: any) => {
-			// Find the adapter for this WebSocket
-			const adapter = (ws as any).__adapter;
-			if (adapter) {
-				adapter._handleMessage(event);
-			}
-		},
-		onClose: (evt: any, ws: any) => {
-			// Find the adapter for this WebSocket
-			const adapter = (ws as any).__adapter;
-			if (adapter) {
-				adapter._handleClose(evt?.code || 1006, evt?.reason || "");
-			}
-		},
-		onError: (error: any, ws: any) => {
-			// Find the adapter for this WebSocket
-			const adapter = (ws as any).__adapter;
-			if (adapter) {
-				adapter._handleError(error);
-			}
-		},
-	};
+		return await actor.handleRawRequest(conn, req);
+	} finally {
+		// Clean up the connection after the request completes
+		if (createdConn) {
+			createdConn.disconnect();
+		}
+	}
 }
 
 // Helper to get the connection encoding from a request
+//
+// Defaults to JSON if not provided so we can support vanilla curl requests easily.
 export function getRequestEncoding(req: HonoRequest): Encoding {
 	const encodingParam = req.header(HEADER_ENCODING);
 	if (!encodingParam) {
-		throw new errors.InvalidEncoding("undefined");
+		return "json";
 	}
 
 	const result = EncodingSchema.safeParse(encodingParam);
@@ -728,22 +222,3 @@ export function getRequestConnParams(req: HonoRequest): unknown {
 		);
 	}
 }
-
-/**
- * Truncase the PATH_RAW_WEBSOCKET_PREFIX path prefix in order to pass a clean
- * path to the onWebSocket handler.
- *
- * Example:
- * - `/raw/websocket/foo` -> `/foo`
- * - `/raw/websocket` -> `/`
- */
-export function truncateRawWebSocketPathPrefix(path: string): string {
-	// Extract the path after prefix and preserve query parameters
-	// Use URL API for cleaner parsing
-	const url = new URL(path, "http://actor");
-	const pathname = url.pathname.replace(/^\/raw\/websocket\/?/, "") || "/";
-	const normalizedPath =
-		(pathname.startsWith("/") ? pathname : "/" + pathname) + url.search;
-
-	return normalizedPath;
-}