rivetkit 2.0.5 → 2.0.7-rc.1

package/src/actor/router-endpoints.ts

@@ -2,19 +2,26 @@ 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/connection";
+import type { AnyConn } from "@/actor/conn";
 import {
-	CONNECTION_DRIVER_HTTP,
-	CONNECTION_DRIVER_SSE,
-	CONNECTION_DRIVER_WEBSOCKET,
 	generateConnId,
+	generateConnSocketId,
 	generateConnToken,
-} from "@/actor/connection";
+} from "@/actor/conn";
+import { ConnDriverKind } from "@/actor/conn-drivers";
 import * as errors from "@/actor/errors";
 import type { AnyActorInstance } from "@/actor/instance";
 import type { InputData } from "@/actor/protocol/serde";
 import { type Encoding, EncodingSchema } from "@/actor/protocol/serde";
+import {
+	HEADER_ACTOR_QUERY,
+	HEADER_CONN_ID,
+	HEADER_CONN_PARAMS,
+	HEADER_CONN_TOKEN,
+	HEADER_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";
@@ -31,22 +38,18 @@ import {
 	deserializeWithEncoding,
 	serializeWithEncoding,
 } from "@/serde";
-import { bufferToArrayBuffer } from "@/utils";
+import { bufferToArrayBuffer, promiseWithResolvers } from "@/utils";
 import type { ActorDriver } from "./driver";
-import type {
-	GenericHttpDriverState,
-	GenericSseDriverState,
-	GenericWebSocketDriverState,
-} from "./generic-conn-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;
-	authData: unknown;
 }
 
 export interface ConnectWebSocketOutput {
@@ -60,7 +63,6 @@ export interface ConnectSseOpts {
 	encoding: Encoding;
 	params: unknown;
 	actorId: string;
-	authData: unknown;
 }
 
 export interface ConnectSseOutput {
@@ -74,7 +76,6 @@ export interface ActionOpts {
 	actionName: string;
 	actionArgs: unknown[];
 	actorId: string;
-	authData: unknown;
 }
 
 export interface ActionOutput {
@@ -92,14 +93,12 @@ export interface ConnsMessageOpts {
 export interface FetchOpts {
 	request: Request;
 	actorId: string;
-	authData: unknown;
 }
 
 export interface WebSocketOpts {
 	request: Request;
 	websocket: UniversalWebSocket;
 	actorId: string;
-	authData: unknown;
 }
 
 /**
@@ -112,7 +111,8 @@ export async function handleWebSocketConnect(
 	actorId: string,
 	encoding: Encoding,
 	parameters: unknown,
-	authData: unknown,
+	connId?: string,
+	connToken?: string,
 ): Promise<UpgradeWebSocketArgs> {
 	const exposeInternalError = req ? getRequestExposeInternalError(req) : false;
 
@@ -121,7 +121,7 @@ export async function handleWebSocketConnect(
 		promise: handlersPromise,
 		resolve: handlersResolve,
 		reject: handlersReject,
-	} = Promise.withResolvers<{
+	} = promiseWithResolvers<{
 		conn: AnyConn;
 		actor: AnyActorInstance;
 		connId: string;
@@ -153,40 +153,47 @@ export async function handleWebSocketConnect(
 		};
 	}
 
+	// Promise used to wait for the websocket close in `disconnect`
+	const closePromise = promiseWithResolvers<void>();
+	const socketId = generateConnSocketId();
+
 	return {
 		onOpen: (_evt: any, ws: WSContext) => {
-			actor.rLog.debug("websocket open");
+			actor.rLog.debug("actor websocket open");
 
 			// Run async operations in background
 			(async () => {
 				try {
-					const connId = generateConnId();
-					const connToken = generateConnToken();
-					const connState = await actor.prepareConn(parameters, req);
-
-					// Save socket
-					const connGlobalState =
-						actorDriver.getGenericConnGlobalState(actorId);
-					connGlobalState.websockets.set(connId, ws);
+					let conn: AnyConn;
+
+					// Create or reconnect connection
 					actor.rLog.debug({
-						msg: "registered websocket for conn",
+						msg: connId
+							? "websocket reconnection attempt"
+							: "new websocket connection",
+						connId,
 						actorId,
-						totalCount: connGlobalState.websockets.size,
 					});
 
-					// Create connection
-					const conn = await actor.createConn(
+					conn = await actor.createConn(
+						{
+							socketId,
+							driverState: {
+								[ConnDriverKind.WEBSOCKET]: {
+									encoding,
+									websocket: ws,
+									closePromise,
+								},
+							},
+						},
+						parameters,
+						req,
 						connId,
 						connToken,
-						parameters,
-						connState,
-						CONNECTION_DRIVER_WEBSOCKET,
-						{ encoding } satisfies GenericWebSocketDriverState,
-						authData,
 					);
 
 					// Unblock other handlers
-					handlersResolve({ conn, actor, connId });
+					handlersResolve({ conn, actor, connId: conn.id });
 				} catch (error) {
 					handlersReject(error);
 
@@ -258,6 +265,10 @@ export async function handleWebSocketConnect(
 			},
 			ws: WSContext,
 		) => {
+			handlersReject(`WebSocket closed (${event.code}): ${event.reason}`);
+
+			closePromise.resolve();
+
 			if (event.wasClean) {
 				actor.rLog.info({
 					msg: "websocket closed",
@@ -280,24 +291,8 @@ export async function handleWebSocketConnect(
 
 			// Handle cleanup asynchronously
 			handlersPromise
-				.then(({ conn, actor, connId }) => {
-					const connGlobalState =
-						actorDriver.getGenericConnGlobalState(actorId);
-					const didDelete = connGlobalState.websockets.delete(connId);
-					if (didDelete) {
-						actor.rLog.info({
-							msg: "removing websocket for conn",
-							totalCount: connGlobalState.websockets.size,
-						});
-					} else {
-						actor.rLog.warn({
-							msg: "websocket does not exist for conn",
-							actorId,
-							totalCount: connGlobalState.websockets.size,
-						});
-					}
-
-					actor.__removeConn(conn);
+				.then(({ conn, actor }) => {
+					actor.__connDisconnected(conn, event.wasClean, socketId);
 				})
 				.catch((error) => {
 					deconstructError(
@@ -332,46 +327,50 @@ export async function handleSseConnect(
 	_runConfig: RunConfig,
 	actorDriver: ActorDriver,
 	actorId: string,
-	authData: unknown,
 ) {
+	c.header("Content-Encoding", "Identity");
+
 	const encoding = getRequestEncoding(c.req);
 	const parameters = getRequestConnParams(c.req);
+	const socketId = generateConnSocketId();
+
+	// 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 connId: string | undefined;
-		let connToken: string | undefined;
-		let connState: unknown;
 		let conn: AnyConn | undefined;
 
 		try {
 			// Do all async work inside the handler
 			actor = await actorDriver.loadActor(actorId);
-			connId = generateConnId();
-			connToken = generateConnToken();
-			connState = await actor.prepareConn(parameters, c.req.raw);
 
-			actor.rLog.debug("sse open");
-
-			// Save stream
-			actorDriver
-				.getGenericConnGlobalState(actorId)
-				.sseStreams.set(connId, stream);
+			// Create or reconnect connection
+			actor.rLog.debug({
+				msg: connId ? "sse reconnection attempt" : "sse open",
+				connId,
+			});
 
-			// Create connection
 			conn = await actor.createConn(
+				{
+					socketId,
+					driverState: {
+						[ConnDriverKind.SSE]: {
+							encoding,
+							stream: stream,
+						},
+					},
+				},
+				parameters,
+				c.req.raw,
 				connId,
 				connToken,
-				parameters,
-				connState,
-				CONNECTION_DRIVER_SSE,
-				{ encoding } satisfies GenericSseDriverState,
-				authData,
 			);
 
 			// Wait for close
-			const abortResolver = Promise.withResolvers();
+			const abortResolver = promiseWithResolvers();
 
 			// HACK: This is required so the abort handler below works
 			//
@@ -380,18 +379,14 @@ export async function handleSseConnect(
 
 			// Handle stream abort (when client closes the connection)
 			c.req.raw.signal.addEventListener("abort", async () => {
-				const rLog = actor?.rLog ?? loggerWithoutContext();
+				invariant(actor, "actor should exist");
+				const rLog = actor.rLog ?? loggerWithoutContext();
 				try {
 					rLog.debug("sse stream aborted");
 
 					// Cleanup
-					if (connId) {
-						actorDriver
-							.getGenericConnGlobalState(actorId)
-							.sseStreams.delete(connId);
-					}
-					if (conn && actor) {
-						actor.__removeConn(conn);
+					if (conn) {
+						actor.__connDisconnected(conn, false, socketId);
 					}
 
 					abortResolver.resolve(undefined);
@@ -401,24 +396,33 @@ export async function handleSseConnect(
 				}
 			});
 
-			// HACK: Will throw if not configured
-			try {
-				c.executionCtx.waitUntil(abortResolver.promise);
-			} catch {}
+			// // 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;
+				}
 
-			// Wait until connection aborted
-			await abortResolver.promise;
+				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 (connId !== undefined) {
-				actorDriver
-					.getGenericConnGlobalState(actorId)
-					.sseStreams.delete(connId);
-			}
 			if (conn && actor !== undefined) {
-				actor.__removeConn(conn);
+				actor.__connDisconnected(conn, false, socketId);
 			}
 
 			// Close the stream on error
@@ -436,7 +440,6 @@ export async function handleAction(
 	actorDriver: ActorDriver,
 	actionName: string,
 	actorId: string,
-	authData: unknown,
 ) {
 	const encoding = getRequestEncoding(c.req);
 	const parameters = getRequestConnParams(c.req);
@@ -449,6 +452,7 @@ export async function handleAction(
 		HTTP_ACTION_REQUEST_VERSIONED,
 	);
 	const actionArgs = cbor.decode(new Uint8Array(request.args));
+	const socketId = generateConnSocketId();
 
 	// Invoke the action
 	let actor: AnyActorInstance | undefined;
@@ -460,15 +464,13 @@ export async function handleAction(
 		actor.rLog.debug({ msg: "handling action", actionName, encoding });
 
 		// Create conn
-		const connState = await actor.prepareConn(parameters, c.req.raw);
 		conn = await actor.createConn(
-			generateConnId(),
-			generateConnToken(),
+			{
+				socketId,
+				driverState: { [ConnDriverKind.HTTP]: {} },
+			},
 			parameters,
-			connState,
-			CONNECTION_DRIVER_HTTP,
-			{} satisfies GenericHttpDriverState,
-			authData,
+			c.req.raw,
 		);
 
 		// Call action
@@ -476,7 +478,8 @@ export async function handleAction(
 		output = await actor.executeAction(ctx, actionName, actionArgs);
 	} finally {
 		if (conn) {
-			actor?.__removeConn(conn);
+			// HTTP connections don't have persistent sockets, so no socket ID needed
+			actor?.__connDisconnected(conn, true, socketId);
 		}
 	}
 
@@ -489,7 +492,9 @@ export async function handleAction(
 		responseData,
 		HTTP_ACTION_RESPONSE_VERSIONED,
 	);
-	return c.body(serialized as Uint8Array, 200, {
+
+	// TODO: Remvoe any, Hono is being a dumbass
+	return c.body(serialized as Uint8Array as any, 200, {
 		"Content-Type": contentTypeForEncoding(encoding),
 	});
 }
@@ -534,12 +539,48 @@ export async function handleConnectionMessage(
 	return c.json({});
 }
 
+export async function handleConnectionClose(
+	c: HonoContext,
+	_runConfig: RunConfig,
+	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();
+	}
+
+	// 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",
+		);
+	}
+
+	// Close the SSE connection
+	await conn.disconnect("Connection closed by client request");
+
+	return c.json({});
+}
+
 export async function handleRawWebSocketHandler(
 	req: Request | undefined,
 	path: string,
 	actorDriver: ActorDriver,
 	actorId: string,
-	authData: unknown,
 ): Promise<UpgradeWebSocketArgs> {
 	const actor = await actorDriver.loadActor(actorId);
 
@@ -644,38 +685,6 @@ export function getRequestQuery(c: HonoContext): unknown {
 	}
 }
 
-export const HEADER_ACTOR_QUERY = "X-RivetKit-Query";
-
-export const HEADER_ENCODING = "X-RivetKit-Encoding";
-
-// IMPORTANT: Params must be in headers or in an E2EE part of the request (i.e. NOT the URL or query string) in order to ensure that tokens can be securely passed in params.
-export const HEADER_CONN_PARAMS = "X-RivetKit-Conn-Params";
-
-// Internal header
-export const HEADER_AUTH_DATA = "X-RivetKit-Auth-Data";
-
-export const HEADER_ACTOR_ID = "X-RivetKit-Actor";
-
-export const HEADER_CONN_ID = "X-RivetKit-Conn";
-
-export const HEADER_CONN_TOKEN = "X-RivetKit-Conn-Token";
-
-/**
- * Headers that publics can send from public clients.
- *
- * Used for CORS.
- **/
-export const ALLOWED_PUBLIC_HEADERS = [
-	"Content-Type",
-	"User-Agent",
-	HEADER_ACTOR_QUERY,
-	HEADER_ENCODING,
-	HEADER_CONN_PARAMS,
-	HEADER_ACTOR_ID,
-	HEADER_CONN_ID,
-	HEADER_CONN_TOKEN,
-];
-
 // Helper to get connection parameters for the request
 export function getRequestConnParams(req: HonoRequest): unknown {
 	const paramsParam = req.header(HEADER_CONN_PARAMS);