chore(rivetkit): specialize context types for better compiler errors

rivetkit-typescript/packages/rivetkit/src/actor/config.ts

@@ -1,24 +1,25 @@
 import { z } from "zod";
 import type { UniversalWebSocket } from "@/common/websocket-interface";
 import type { Conn } from "./conn/mod";
-import type { ActionContext } from "./contexts/action";
-import type { ActorContext } from "./contexts/actor";
-import type { CreateConnStateContext } from "./contexts/create-conn-state";
-import type { OnBeforeConnectContext } from "./contexts/on-before-connect";
-import type { OnConnectContext } from "./contexts/on-connect";
-import type { RequestContext } from "./contexts/request";
-import type { WebSocketContext } from "./contexts/websocket";
+import type {
+	ActionContext,
+	ActorContext,
+	BeforeActionResponseContext,
+	BeforeConnectContext,
+	ConnectContext,
+	CreateConnStateContext,
+	CreateContext,
+	CreateVarsContext,
+	DestroyContext,
+	DisconnectContext,
+	RequestContext,
+	SleepContext,
+	StateChangeContext,
+	WakeContext,
+	WebSocketContext,
+} from "./contexts";
 import type { AnyDatabaseProvider } from "./database";
 
-export type InitContext = ActorContext<
-	undefined,
-	undefined,
-	undefined,
-	undefined,
-	undefined,
-	undefined
->;
-
 export interface ActorTypes<
 	TState,
 	TConnParams,
@@ -134,7 +135,7 @@ type CreateState<TState, TConnParams, TConnState, TVars, TInput, TDatabase> =
 	| { state: TState }
 	| {
 			createState: (
-				c: InitContext,
+				c: CreateContext<TState, TInput, TDatabase>,
 				input: TInput,
 			) => TState | Promise<TState>;
 	  }
@@ -168,7 +169,14 @@ type CreateConnState<
 /**
  * @experimental
  */
-type CreateVars<TState, TConnParams, TConnState, TVars, TInput, TDatabase> =
+type CreateVars<
+	TState,
+	TConnParams,
+	TConnState,
+	TVars,
+	TInput,
+	TDatabase extends AnyDatabaseProvider,
+> =
 	| {
 			/**
 			 * @experimental
@@ -180,7 +188,7 @@ type CreateVars<TState, TConnParams, TConnState, TVars, TInput, TDatabase> =
 			 * @experimental
 			 */
 			createVars: (
-				c: InitContext,
+				c: CreateVarsContext<TState, TInput, TDatabase>,
 				driverCtx: any,
 			) => TVars | Promise<TVars>;
 	  }
@@ -240,22 +248,15 @@ interface BaseActorConfig<
 	 * This is called before any other lifecycle hooks.
 	 */
 	onCreate?: (
-		c: ActorContext<
-			TState,
-			TConnParams,
-			TConnState,
-			TVars,
-			TInput,
-			TDatabase
-		>,
+		c: CreateContext<TState, TInput, TDatabase>,
 		input: TInput,
 	) => void | Promise<void>;
 
 	/**
 	 * Called when the actor is destroyed.
 	 */
 	onDestroy?: (
-		c: ActorContext<
+		c: DestroyContext<
 			TState,
 			TConnParams,
 			TConnState,
@@ -274,7 +275,7 @@ interface BaseActorConfig<
 	 * @returns Void or a Promise that resolves when startup is complete
 	 */
 	onWake?: (
-		c: ActorContext<
+		c: WakeContext<
 			TState,
 			TConnParams,
 			TConnState,
@@ -295,7 +296,7 @@ interface BaseActorConfig<
 	 * @returns Void or a Promise that resolves when shutdown is complete
 	 */
 	onSleep?: (
-		c: ActorContext<
+		c: SleepContext<
 			TState,
 			TConnParams,
 			TConnState,
@@ -317,7 +318,7 @@ interface BaseActorConfig<
 	 * @param newState The updated state
 	 */
 	onStateChange?: (
-		c: ActorContext<
+		c: StateChangeContext<
 			TState,
 			TConnParams,
 			TConnState,
@@ -339,7 +340,7 @@ interface BaseActorConfig<
 	 * @throws Throw an error to reject the connection
 	 */
 	onBeforeConnect?: (
-		c: OnBeforeConnectContext<TState, TVars, TInput, TDatabase>,
+		c: BeforeConnectContext<TState, TVars, TInput, TDatabase>,
 		params: TConnParams,
 	) => void | Promise<void>;
 
@@ -353,7 +354,7 @@ interface BaseActorConfig<
 	 * @returns Void or a Promise that resolves when connection handling is complete
 	 */
 	onConnect?: (
-		c: OnConnectContext<
+		c: ConnectContext<
 			TState,
 			TConnParams,
 			TConnState,
@@ -374,7 +375,7 @@ interface BaseActorConfig<
 	 * @returns Void or a Promise that resolves when disconnect handling is complete
 	 */
 	onDisconnect?: (
-		c: ActorContext<
+		c: DisconnectContext<
 			TState,
 			TConnParams,
 			TConnState,
@@ -398,7 +399,7 @@ interface BaseActorConfig<
 	 * @returns The modified output to send to the client
 	 */
 	onBeforeActionResponse?: <Out>(
-		c: ActorContext<
+		c: BeforeActionResponseContext<
 			TState,
 			TConnParams,
 			TConnState,

rivetkit-typescript/packages/rivetkit/src/actor/contexts/action.ts

@@ -1,7 +1,8 @@
 import type { Conn } from "../conn/mod";
 import type { AnyDatabaseProvider } from "../database";
+import type { ActorDefinition, AnyActorDefinition } from "../definition";
 import type { ActorInstance } from "../instance/mod";
-import { ConnContext } from "./conn";
+import { ConnContext } from "./base/conn";
 
 /**
  * Context for a remote procedure call.
@@ -21,3 +22,18 @@ export class ActionContext<
 	TInput,
 	TDatabase
 > {}
+
+/**
+ * Extracts the ActionContext type from an ActorDefinition.
+ */
+export type ActionContextOf<AD extends AnyActorDefinition> = AD extends ActorDefinition<
+	infer S,
+	infer CP,
+	infer CS,
+	infer V,
+	infer I,
+	infer DB extends AnyDatabaseProvider,
+	any
+>
+	? ActionContext<S, CP, CS, V, I, DB>
+	: never;

rivetkit-typescript/packages/rivetkit/src/actor/contexts/actor.ts → rivetkit-typescript/packages/rivetkit/src/actor/contexts/base/actor.ts

@@ -2,10 +2,14 @@ import type { ActorKey } from "@/actor/mod";
 import type { Client } from "@/client/client";
 import type { Logger } from "@/common/log";
 import type { Registry } from "@/registry/mod";
-import type { Conn, ConnId } from "../conn/mod";
-import type { AnyDatabaseProvider, InferDatabaseClient } from "../database";
-import type { ActorInstance, SaveStateOptions } from "../instance/mod";
-import type { Schedule } from "../schedule";
+import type { Conn, ConnId } from "../../conn/mod";
+import type { AnyDatabaseProvider, InferDatabaseClient } from "../../database";
+import type {
+	ActorDefinition,
+	AnyActorDefinition,
+} from "../../definition";
+import type { ActorInstance, SaveStateOptions } from "../../instance/mod";
+import type { Schedule } from "../../schedule";
 
 /**
  * ActorContext class that provides access to actor methods and state
@@ -42,16 +46,23 @@ export class ActorContext<
 
 	/**
 	 * Get the actor state
+	 *
+	 * @remarks
+	 * This property is not available in `createState` since the state hasn't been created yet.
 	 */
-	get state(): TState {
-		return this.#actor.state;
+	get state(): TState extends never ? never : TState {
+		return this.#actor.state as TState extends never ? never : TState;
 	}
 
 	/**
 	 * Get the actor variables
+	 *
+	 * @remarks
+	 * This property is not available in `createVars` since the variables haven't been created yet.
+	 * Variables are only available if you define `vars` or `createVars` in your actor config.
 	 */
-	get vars(): TVars {
-		return this.#actor.vars;
+	get vars(): TVars extends never ? never : TVars {
+		return this.#actor.vars as TVars extends never ? never : TVars;
 	}
 
 	/**
@@ -125,11 +136,18 @@ export class ActorContext<
 
 	/**
 	 * Gets the database.
+	 *
 	 * @experimental
+	 * @remarks
+	 * This property is only available if you define a `db` provider in your actor config.
 	 * @throws {DatabaseNotEnabled} If the database is not enabled.
 	 */
-	get db(): InferDatabaseClient<TDatabase> {
-		return this.#actor.db;
+	get db(): TDatabase extends never
+		? never
+		: InferDatabaseClient<TDatabase> {
+		return this.#actor.db as TDatabase extends never
+			? never
+			: InferDatabaseClient<TDatabase>;
 	}
 
 	/**
@@ -177,3 +195,15 @@ export class ActorContext<
 		this.#actor.startDestroy();
 	}
 }
+
+export type ActorContextOf<AD extends AnyActorDefinition> = AD extends ActorDefinition<
+	infer S,
+	infer CP,
+	infer CS,
+	infer V,
+	infer I,
+	infer DB extends AnyDatabaseProvider,
+	any
+>
+	? ActorContext<S, CP, CS, V, I, DB>
+	: never;

rivetkit-typescript/packages/rivetkit/src/actor/contexts/conn-init.ts → rivetkit-typescript/packages/rivetkit/src/actor/contexts/base/conn-init.ts

@@ -1,5 +1,9 @@
-import type { AnyDatabaseProvider } from "../database";
-import type { ActorInstance } from "../instance/mod";
+import type { AnyDatabaseProvider } from "../../database";
+import type {
+	ActorDefinition,
+	AnyActorDefinition,
+} from "../../definition";
+import type { ActorInstance } from "../../instance/mod";
 import { ActorContext } from "./actor";
 
 /**
@@ -11,7 +15,7 @@ export abstract class ConnInitContext<
 	TVars,
 	TInput,
 	TDatabase extends AnyDatabaseProvider,
-> extends ActorContext<TState, undefined, undefined, TVars, TInput, TDatabase> {
+> extends ActorContext<TState, never, never, TVars, TInput, TDatabase> {
 	/**
 	 * The incoming request that initiated the connection.
 	 * May be undefined for connections initiated without a direct HTTP request.
@@ -25,7 +29,20 @@ export abstract class ConnInitContext<
 		actor: ActorInstance<TState, any, any, TVars, TInput, TDatabase>,
 		request: Request | undefined,
 	) {
-		super(actor);
+		super(actor as any);
 		this.request = request;
 	}
 }
+
+export type ConnInitContextOf<AD extends AnyActorDefinition> =
+	AD extends ActorDefinition<
+		infer S,
+		any,
+		any,
+		infer V,
+		infer I,
+		infer DB extends AnyDatabaseProvider,
+		any
+	>
+		? ConnInitContext<S, V, I, DB>
+		: never;

rivetkit-typescript/packages/rivetkit/src/actor/contexts/conn.ts → rivetkit-typescript/packages/rivetkit/src/actor/contexts/base/conn.ts

@@ -1,6 +1,10 @@
-import type { Conn } from "../conn/mod";
-import type { AnyDatabaseProvider } from "../database";
-import type { ActorInstance } from "../instance/mod";
+import type { Conn } from "../../conn/mod";
+import type { AnyDatabaseProvider } from "../../database";
+import type {
+	ActorDefinition,
+	AnyActorDefinition,
+} from "../../definition";
+import type { ActorInstance } from "../../instance/mod";
 import { ActorContext } from "./actor";
 
 /**
@@ -46,3 +50,15 @@ export abstract class ConnContext<
 		super(actor);
 	}
 }
+
+export type ConnContextOf<AD extends AnyActorDefinition> = AD extends ActorDefinition<
+	infer S,
+	infer CP,
+	infer CS,
+	infer V,
+	infer I,
+	infer DB extends AnyDatabaseProvider,
+	any
+>
+	? ConnContext<S, CP, CS, V, I, DB>
+	: never;

rivetkit-typescript/packages/rivetkit/src/actor/contexts/before-action-response.ts

@@ -0,0 +1,35 @@
+import type { AnyDatabaseProvider } from "../database";
+import type { ActorDefinition, AnyActorDefinition } from "../definition";
+import { ActorContext } from "./base/actor";
+
+/**
+ * Context for the onBeforeActionResponse lifecycle hook.
+ */
+export class BeforeActionResponseContext<
+	TState,
+	TConnParams,
+	TConnState,
+	TVars,
+	TInput,
+	TDatabase extends AnyDatabaseProvider,
+> extends ActorContext<
+	TState,
+	TConnParams,
+	TConnState,
+	TVars,
+	TInput,
+	TDatabase
+> {}
+
+export type BeforeActionResponseContextOf<AD extends AnyActorDefinition> =
+	AD extends ActorDefinition<
+		infer S,
+		infer CP,
+		infer CS,
+		infer V,
+		infer I,
+		infer DB extends AnyDatabaseProvider,
+		any
+	>
+		? BeforeActionResponseContext<S, CP, CS, V, I, DB>
+		: never;