Skip to content

Commit f2f12a8

Browse files
committed
chore(rivetkit): specialize context types for better compiler errors (#3669)
1 parent efc24bf commit f2f12a8

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

41 files changed

+680
-267
lines changed

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

Lines changed: 37 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -1,24 +1,25 @@
11
import { z } from "zod";
22
import type { UniversalWebSocket } from "@/common/websocket-interface";
33
import type { Conn } from "./conn/mod";
4-
import type { ActionContext } from "./contexts/action";
5-
import type { ActorContext } from "./contexts/actor";
6-
import type { CreateConnStateContext } from "./contexts/create-conn-state";
7-
import type { OnBeforeConnectContext } from "./contexts/on-before-connect";
8-
import type { OnConnectContext } from "./contexts/on-connect";
9-
import type { RequestContext } from "./contexts/request";
10-
import type { WebSocketContext } from "./contexts/websocket";
4+
import type {
5+
ActionContext,
6+
ActorContext,
7+
BeforeActionResponseContext,
8+
BeforeConnectContext,
9+
ConnectContext,
10+
CreateConnStateContext,
11+
CreateContext,
12+
CreateVarsContext,
13+
DestroyContext,
14+
DisconnectContext,
15+
RequestContext,
16+
SleepContext,
17+
StateChangeContext,
18+
WakeContext,
19+
WebSocketContext,
20+
} from "./contexts";
1121
import type { AnyDatabaseProvider } from "./database";
1222

13-
export type InitContext = ActorContext<
14-
undefined,
15-
undefined,
16-
undefined,
17-
undefined,
18-
undefined,
19-
undefined
20-
>;
21-
2223
export interface ActorTypes<
2324
TState,
2425
TConnParams,
@@ -131,11 +132,11 @@ export const ActorConfigSchema = z
131132
// This must have only one or the other or else TState will not be able to be inferred
132133
//
133134
// Data returned from this handler will be available on `c.state`.
134-
type CreateState<TState, TConnParams, TConnState, TVars, TInput, TDatabase> =
135+
type CreateState<TState, TConnParams, TConnState, TVars, TInput, TDatabase extends AnyDatabaseProvider> =
135136
| { state: TState }
136137
| {
137138
createState: (
138-
c: InitContext,
139+
c: CreateContext<TState, TInput, TDatabase>,
139140
input: TInput,
140141
) => TState | Promise<TState>;
141142
}
@@ -169,7 +170,14 @@ type CreateConnState<
169170
/**
170171
* @experimental
171172
*/
172-
type CreateVars<TState, TConnParams, TConnState, TVars, TInput, TDatabase> =
173+
type CreateVars<
174+
TState,
175+
TConnParams,
176+
TConnState,
177+
TVars,
178+
TInput,
179+
TDatabase extends AnyDatabaseProvider,
180+
> =
173181
| {
174182
/**
175183
* @experimental
@@ -181,7 +189,7 @@ type CreateVars<TState, TConnParams, TConnState, TVars, TInput, TDatabase> =
181189
* @experimental
182190
*/
183191
createVars: (
184-
c: InitContext,
192+
c: CreateVarsContext<TState, TInput, TDatabase>,
185193
driverCtx: any,
186194
) => TVars | Promise<TVars>;
187195
}
@@ -241,22 +249,15 @@ interface BaseActorConfig<
241249
* This is called before any other lifecycle hooks.
242250
*/
243251
onCreate?: (
244-
c: ActorContext<
245-
TState,
246-
TConnParams,
247-
TConnState,
248-
TVars,
249-
TInput,
250-
TDatabase
251-
>,
252+
c: CreateContext<TState, TInput, TDatabase>,
252253
input: TInput,
253254
) => void | Promise<void>;
254255

255256
/**
256257
* Called when the actor is destroyed.
257258
*/
258259
onDestroy?: (
259-
c: ActorContext<
260+
c: DestroyContext<
260261
TState,
261262
TConnParams,
262263
TConnState,
@@ -275,7 +276,7 @@ interface BaseActorConfig<
275276
* @returns Void or a Promise that resolves when startup is complete
276277
*/
277278
onWake?: (
278-
c: ActorContext<
279+
c: WakeContext<
279280
TState,
280281
TConnParams,
281282
TConnState,
@@ -296,7 +297,7 @@ interface BaseActorConfig<
296297
* @returns Void or a Promise that resolves when shutdown is complete
297298
*/
298299
onSleep?: (
299-
c: ActorContext<
300+
c: SleepContext<
300301
TState,
301302
TConnParams,
302303
TConnState,
@@ -318,7 +319,7 @@ interface BaseActorConfig<
318319
* @param newState The updated state
319320
*/
320321
onStateChange?: (
321-
c: ActorContext<
322+
c: StateChangeContext<
322323
TState,
323324
TConnParams,
324325
TConnState,
@@ -340,7 +341,7 @@ interface BaseActorConfig<
340341
* @throws Throw an error to reject the connection
341342
*/
342343
onBeforeConnect?: (
343-
c: OnBeforeConnectContext<TState, TVars, TInput, TDatabase>,
344+
c: BeforeConnectContext<TState, TVars, TInput, TDatabase>,
344345
params: TConnParams,
345346
) => void | Promise<void>;
346347

@@ -354,7 +355,7 @@ interface BaseActorConfig<
354355
* @returns Void or a Promise that resolves when connection handling is complete
355356
*/
356357
onConnect?: (
357-
c: OnConnectContext<
358+
c: ConnectContext<
358359
TState,
359360
TConnParams,
360361
TConnState,
@@ -375,7 +376,7 @@ interface BaseActorConfig<
375376
* @returns Void or a Promise that resolves when disconnect handling is complete
376377
*/
377378
onDisconnect?: (
378-
c: ActorContext<
379+
c: DisconnectContext<
379380
TState,
380381
TConnParams,
381382
TConnState,
@@ -399,7 +400,7 @@ interface BaseActorConfig<
399400
* @returns The modified output to send to the client
400401
*/
401402
onBeforeActionResponse?: <Out>(
402-
c: ActorContext<
403+
c: BeforeActionResponseContext<
403404
TState,
404405
TConnParams,
405406
TConnState,

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

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
11
import type { Conn } from "../conn/mod";
22
import type { AnyDatabaseProvider } from "../database";
3+
import type { ActorDefinition, AnyActorDefinition } from "../definition";
34
import type { ActorInstance } from "../instance/mod";
4-
import { ConnContext } from "./conn";
5+
import { ConnContext } from "./base/conn";
56

67
/**
78
* Context for a remote procedure call.
@@ -21,3 +22,18 @@ export class ActionContext<
2122
TInput,
2223
TDatabase
2324
> {}
25+
26+
/**
27+
* Extracts the ActionContext type from an ActorDefinition.
28+
*/
29+
export type ActionContextOf<AD extends AnyActorDefinition> = AD extends ActorDefinition<
30+
infer S,
31+
infer CP,
32+
infer CS,
33+
infer V,
34+
infer I,
35+
infer DB extends AnyDatabaseProvider,
36+
any
37+
>
38+
? ActionContext<S, CP, CS, V, I, DB>
39+
: never;

rivetkit-typescript/packages/rivetkit/src/actor/contexts/actor.ts renamed to rivetkit-typescript/packages/rivetkit/src/actor/contexts/base/actor.ts

Lines changed: 36 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,11 @@ import type { ActorKey } from "@/actor/mod";
22
import type { Client } from "@/client/client";
33
import type { Logger } from "@/common/log";
44
import type { Registry } from "@/registry";
5-
import type { Conn, ConnId } from "../conn/mod";
6-
import type { AnyDatabaseProvider, InferDatabaseClient } from "../database";
7-
import type { ActorInstance, SaveStateOptions } from "../instance/mod";
8-
import type { Schedule } from "../schedule";
5+
import type { Conn, ConnId } from "../../conn/mod";
6+
import type { AnyDatabaseProvider, InferDatabaseClient } from "../../database";
7+
import type { ActorDefinition, AnyActorDefinition } from "../../definition";
8+
import type { ActorInstance, SaveStateOptions } from "../../instance/mod";
9+
import type { Schedule } from "../../schedule";
910

1011
/**
1112
* ActorContext class that provides access to actor methods and state
@@ -42,16 +43,23 @@ export class ActorContext<
4243

4344
/**
4445
* Get the actor state
46+
*
47+
* @remarks
48+
* This property is not available in `createState` since the state hasn't been created yet.
4549
*/
46-
get state(): TState {
47-
return this.#actor.state;
50+
get state(): TState extends never ? never : TState {
51+
return this.#actor.state as TState extends never ? never : TState;
4852
}
4953

5054
/**
5155
* Get the actor variables
56+
*
57+
* @remarks
58+
* This property is not available in `createVars` since the variables haven't been created yet.
59+
* Variables are only available if you define `vars` or `createVars` in your actor config.
5260
*/
53-
get vars(): TVars {
54-
return this.#actor.vars;
61+
get vars(): TVars extends never ? never : TVars {
62+
return this.#actor.vars as TVars extends never ? never : TVars;
5563
}
5664

5765
/**
@@ -125,11 +133,16 @@ export class ActorContext<
125133

126134
/**
127135
* Gets the database.
136+
*
128137
* @experimental
138+
* @remarks
139+
* This property is only available if you define a `db` provider in your actor config.
129140
* @throws {DatabaseNotEnabled} If the database is not enabled.
130141
*/
131-
get db(): InferDatabaseClient<TDatabase> {
132-
return this.#actor.db;
142+
get db(): TDatabase extends never ? never : InferDatabaseClient<TDatabase> {
143+
return this.#actor.db as TDatabase extends never
144+
? never
145+
: InferDatabaseClient<TDatabase>;
133146
}
134147

135148
/**
@@ -177,3 +190,16 @@ export class ActorContext<
177190
this.#actor.startDestroy();
178191
}
179192
}
193+
194+
export type ActorContextOf<AD extends AnyActorDefinition> =
195+
AD extends ActorDefinition<
196+
infer S,
197+
infer CP,
198+
infer CS,
199+
infer V,
200+
infer I,
201+
infer DB extends AnyDatabaseProvider,
202+
any
203+
>
204+
? ActorContext<S, CP, CS, V, I, DB>
205+
: never;

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

Lines changed: 21 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,9 @@
1-
import type { AnyDatabaseProvider } from "../database";
2-
import type { ActorInstance } from "../instance/mod";
1+
import type { AnyDatabaseProvider } from "../../database";
2+
import type {
3+
ActorDefinition,
4+
AnyActorDefinition,
5+
} from "../../definition";
6+
import type { ActorInstance } from "../../instance/mod";
37
import { ActorContext } from "./actor";
48

59
/**
@@ -11,7 +15,7 @@ export abstract class ConnInitContext<
1115
TVars,
1216
TInput,
1317
TDatabase extends AnyDatabaseProvider,
14-
> extends ActorContext<TState, undefined, undefined, TVars, TInput, TDatabase> {
18+
> extends ActorContext<TState, never, never, TVars, TInput, TDatabase> {
1519
/**
1620
* The incoming request that initiated the connection.
1721
* May be undefined for connections initiated without a direct HTTP request.
@@ -25,7 +29,20 @@ export abstract class ConnInitContext<
2529
actor: ActorInstance<TState, any, any, TVars, TInput, TDatabase>,
2630
request: Request | undefined,
2731
) {
28-
super(actor);
32+
super(actor as any);
2933
this.request = request;
3034
}
3135
}
36+
37+
export type ConnInitContextOf<AD extends AnyActorDefinition> =
38+
AD extends ActorDefinition<
39+
infer S,
40+
any,
41+
any,
42+
infer V,
43+
infer I,
44+
infer DB extends AnyDatabaseProvider,
45+
any
46+
>
47+
? ConnInitContext<S, V, I, DB>
48+
: never;

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

Lines changed: 19 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,10 @@
1-
import type { Conn } from "../conn/mod";
2-
import type { AnyDatabaseProvider } from "../database";
3-
import type { ActorInstance } from "../instance/mod";
1+
import type { Conn } from "../../conn/mod";
2+
import type { AnyDatabaseProvider } from "../../database";
3+
import type {
4+
ActorDefinition,
5+
AnyActorDefinition,
6+
} from "../../definition";
7+
import type { ActorInstance } from "../../instance/mod";
48
import { ActorContext } from "./actor";
59

610
/**
@@ -46,3 +50,15 @@ export abstract class ConnContext<
4650
super(actor);
4751
}
4852
}
53+
54+
export type ConnContextOf<AD extends AnyActorDefinition> = AD extends ActorDefinition<
55+
infer S,
56+
infer CP,
57+
infer CS,
58+
infer V,
59+
infer I,
60+
infer DB extends AnyDatabaseProvider,
61+
any
62+
>
63+
? ConnContext<S, CP, CS, V, I, DB>
64+
: never;
Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
import type { AnyDatabaseProvider } from "../database";
2+
import type { ActorDefinition, AnyActorDefinition } from "../definition";
3+
import { ActorContext } from "./base/actor";
4+
5+
/**
6+
* Context for the onBeforeActionResponse lifecycle hook.
7+
*/
8+
export class BeforeActionResponseContext<
9+
TState,
10+
TConnParams,
11+
TConnState,
12+
TVars,
13+
TInput,
14+
TDatabase extends AnyDatabaseProvider,
15+
> extends ActorContext<
16+
TState,
17+
TConnParams,
18+
TConnState,
19+
TVars,
20+
TInput,
21+
TDatabase
22+
> {}
23+
24+
export type BeforeActionResponseContextOf<AD extends AnyActorDefinition> =
25+
AD extends ActorDefinition<
26+
infer S,
27+
infer CP,
28+
infer CS,
29+
infer V,
30+
infer I,
31+
infer DB extends AnyDatabaseProvider,
32+
any
33+
>
34+
? BeforeActionResponseContext<S, CP, CS, V, I, DB>
35+
: never;

0 commit comments

Comments
 (0)