rivetkit 2.0.6 → 2.0.7

package/src/actor/router-endpoints.ts

@@ -2,22 +2,24 @@ 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";
@@ -36,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 {
@@ -65,7 +63,6 @@ export interface ConnectSseOpts {
 	encoding: Encoding;
 	params: unknown;
 	actorId: string;
-	authData: unknown;
 }
 
 export interface ConnectSseOutput {
@@ -79,7 +76,6 @@ export interface ActionOpts {
 	actionName: string;
 	actionArgs: unknown[];
 	actorId: string;
-	authData: unknown;
 }
 
 export interface ActionOutput {
@@ -97,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;
 }
 
 /**
@@ -117,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;
 
@@ -126,7 +121,7 @@ export async function handleWebSocketConnect(
 		promise: handlersPromise,
 		resolve: handlersResolve,
 		reject: handlersReject,
-	} = Promise.withResolvers<{
+	} = promiseWithResolvers<{
 		conn: AnyConn;
 		actor: AnyActorInstance;
 		connId: string;
@@ -158,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);
 
@@ -263,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",
@@ -285,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(
@@ -337,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
 			//
@@ -385,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);
@@ -406,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
@@ -441,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);
@@ -454,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;
@@ -465,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
@@ -481,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);
 		}
 	}
 
@@ -494,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),
 	});
 }
@@ -539,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);
 

package/src/actor/router.ts

@@ -11,19 +11,24 @@ import {
 	type ConnectWebSocketOutput,
 	type ConnsMessageOpts,
 	handleAction,
+	handleConnectionClose,
 	handleConnectionMessage,
 	handleRawWebSocketHandler,
 	handleSseConnect,
 	handleWebSocketConnect,
 } from "@/actor/router-endpoints";
 import {
-	HEADER_AUTH_DATA,
 	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,
 } from "@/common/actor-router-consts";
 import {
 	handleRouteError,
@@ -35,8 +40,9 @@ import {
 	type ActorInspectorRouterEnv,
 	createActorInspectorRouter,
 } from "@/inspector/actor";
-import { secureInspector } from "@/inspector/utils";
+import { isInspectorEnabled, secureInspector } from "@/inspector/utils";
 import type { RunConfig } from "@/registry/run-config";
+import { ConnDriverKind } from "./conn-drivers";
 import type { ActorDriver } from "./driver";
 import { InternalError } from "./errors";
 import { loggerWithoutContext } from "./log";
@@ -78,19 +84,70 @@ export function createActorRouter(
 		return c.text("ok");
 	});
 
+	// Test endpoint to force disconnect a connection non-cleanly
+	router.post("/.test/force-disconnect", async (c) => {
+		const connId = c.req.query("conn");
+
+		if (!connId) {
+			return c.text("Missing conn query parameter", 400);
+		}
+
+		const actor = await actorDriver.loadActor(c.env.actorId);
+		const conn = actor.__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();
+		}
+
+		return c.json({ success: true });
+	});
+
 	router.get(PATH_CONNECT_WEBSOCKET, async (c) => {
 		const upgradeWebSocket = runConfig.getUpgradeWebSocket?.();
 		if (upgradeWebSocket) {
 			return upgradeWebSocket(async (c) => {
-				const encodingRaw = c.req.header(HEADER_ENCODING);
-				const connParamsRaw = c.req.header(HEADER_CONN_PARAMS);
-				const authDataRaw = c.req.header(HEADER_AUTH_DATA);
+				// 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;
-				const authData = authDataRaw ? JSON.parse(authDataRaw) : undefined;
 
 				return await handleWebSocketConnect(
 					c.req.raw,
@@ -99,7 +156,8 @@ export function createActorRouter(
 					c.env.actorId,
 					encoding,
 					connParams,
-					authData,
+					connIdRaw,
+					connTokenRaw,
 				);
 			})(c, noopNext());
 		} else {
@@ -111,41 +169,38 @@ export function createActorRouter(
 	});
 
 	router.get("/connect/sse", async (c) => {
-		const authDataRaw = c.req.header(HEADER_AUTH_DATA);
-		let authData: unknown;
-		if (authDataRaw) {
-			authData = JSON.parse(authDataRaw);
-		}
-
-		return handleSseConnect(c, runConfig, actorDriver, c.env.actorId, authData);
+		return handleSseConnect(c, runConfig, actorDriver, c.env.actorId);
 	});
 
 	router.post("/action/:action", async (c) => {
 		const actionName = c.req.param("action");
 
-		const authDataRaw = c.req.header(HEADER_AUTH_DATA);
-		let authData: unknown;
-		if (authDataRaw) {
-			authData = JSON.parse(authDataRaw);
-		}
+		return handleAction(c, runConfig, actorDriver, actionName, c.env.actorId);
+	});
 
-		return handleAction(
+	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,
-			actionName,
+			connId,
+			connToken,
 			c.env.actorId,
-			authData,
 		);
 	});
 
-	router.post("/connections/message", async (c) => {
+	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 handleConnectionMessage(
+		return handleConnectionClose(
 			c,
 			runConfig,
 			actorDriver,
@@ -157,12 +212,6 @@ export function createActorRouter(
 
 	// Raw HTTP endpoints - /http/*
 	router.all("/raw/http/*", async (c) => {
-		const authDataRaw = c.req.header(HEADER_AUTH_DATA);
-		let authData: unknown;
-		if (authDataRaw) {
-			authData = JSON.parse(authDataRaw);
-		}
-
 		const actor = await actorDriver.loadActor(c.env.actorId);
 
 		// TODO: This is not a clean way of doing this since `/http/` might exist mid-path
@@ -186,9 +235,7 @@ export function createActorRouter(
 		});
 
 		// Call the actor's onFetch handler - it will throw appropriate errors
-		const response = await actor.handleFetch(correctedRequest, {
-			auth: authData,
-		});
+		const response = await actor.handleFetch(correctedRequest, {});
 
 		// This should never happen now since handleFetch throws errors
 		if (!response) {
@@ -203,16 +250,6 @@ export function createActorRouter(
 		const upgradeWebSocket = runConfig.getUpgradeWebSocket?.();
 		if (upgradeWebSocket) {
 			return upgradeWebSocket(async (c) => {
-				const encodingRaw = c.req.header(HEADER_ENCODING);
-				const connParamsRaw = c.req.header(HEADER_CONN_PARAMS);
-				const authDataRaw = c.req.header(HEADER_AUTH_DATA);
-
-				const encoding = EncodingSchema.parse(encodingRaw);
-				const connParams = connParamsRaw
-					? JSON.parse(connParamsRaw)
-					: undefined;
-				const authData = authDataRaw ? JSON.parse(authDataRaw) : undefined;
-
 				const url = new URL(c.req.url);
 				const pathWithQuery = c.req.path + url.search;
 
@@ -229,7 +266,6 @@ export function createActorRouter(
 					pathWithQuery,
 					actorDriver,
 					c.env.actorId,
-					authData,
 				);
 			})(c, noopNext());
 		} else {
@@ -240,7 +276,7 @@ export function createActorRouter(
 		}
 	});
 
-	if (runConfig.inspector.enabled) {
+	if (isInspectorEnabled(runConfig, "actor")) {
 		router.route(
 			"/inspect",
 			new Hono<ActorInspectorRouterEnv & { Bindings: ActorRouterBindings }>()

package/src/actor/utils.ts

@@ -87,7 +87,11 @@ export class Lock<T> {
 export function generateSecureToken(length = 32) {
 	const array = new Uint8Array(length);
 	crypto.getRandomValues(array);
-	return btoa(String.fromCharCode(...array));
+	// Replace base64 chars that are not URL safe with URL-safe chars and strip padding
+	return btoa(String.fromCharCode(...array))
+		.replace(/\+/g, "-")
+		.replace(/\//g, "_")
+		.replace(/=/g, "");
 }
 
 export function generateRandomString(length = 32) {