Skip to content

Commit 4040925

Browse files
NathanFlurryNathanFlurry
committed
chore(rivetkit): specialize context types for better compiler errors
1 parent 96cb137 commit 4040925

28 files changed

+565
-155
lines changed

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

Lines changed: 36 additions & 35 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,
@@ -134,7 +135,7 @@ type CreateState<TState, TConnParams, TConnState, TVars, TInput, TDatabase> =
134135
| { state: TState }
135136
| {
136137
createState: (
137-
c: InitContext,
138+
c: CreateContext<TState, TInput, TDatabase>,
138139
input: TInput,
139140
) => TState | Promise<TState>;
140141
}
@@ -168,7 +169,14 @@ type CreateConnState<
168169
/**
169170
* @experimental
170171
*/
171-
type CreateVars<TState, TConnParams, TConnState, TVars, TInput, TDatabase> =
172+
type CreateVars<
173+
TState,
174+
TConnParams,
175+
TConnState,
176+
TVars,
177+
TInput,
178+
TDatabase extends AnyDatabaseProvider,
179+
> =
172180
| {
173181
/**
174182
* @experimental
@@ -180,7 +188,7 @@ type CreateVars<TState, TConnParams, TConnState, TVars, TInput, TDatabase> =
180188
* @experimental
181189
*/
182190
createVars: (
183-
c: InitContext,
191+
c: CreateVarsContext<TState, TInput, TDatabase>,
184192
driverCtx: any,
185193
) => TVars | Promise<TVars>;
186194
}
@@ -240,22 +248,15 @@ interface BaseActorConfig<
240248
* This is called before any other lifecycle hooks.
241249
*/
242250
onCreate?: (
243-
c: ActorContext<
244-
TState,
245-
TConnParams,
246-
TConnState,
247-
TVars,
248-
TInput,
249-
TDatabase
250-
>,
251+
c: CreateContext<TState, TInput, TDatabase>,
251252
input: TInput,
252253
) => void | Promise<void>;
253254

254255
/**
255256
* Called when the actor is destroyed.
256257
*/
257258
onDestroy?: (
258-
c: ActorContext<
259+
c: DestroyContext<
259260
TState,
260261
TConnParams,
261262
TConnState,
@@ -274,7 +275,7 @@ interface BaseActorConfig<
274275
* @returns Void or a Promise that resolves when startup is complete
275276
*/
276277
onWake?: (
277-
c: ActorContext<
278+
c: WakeContext<
278279
TState,
279280
TConnParams,
280281
TConnState,
@@ -295,7 +296,7 @@ interface BaseActorConfig<
295296
* @returns Void or a Promise that resolves when shutdown is complete
296297
*/
297298
onSleep?: (
298-
c: ActorContext<
299+
c: SleepContext<
299300
TState,
300301
TConnParams,
301302
TConnState,
@@ -317,7 +318,7 @@ interface BaseActorConfig<
317318
* @param newState The updated state
318319
*/
319320
onStateChange?: (
320-
c: ActorContext<
321+
c: StateChangeContext<
321322
TState,
322323
TConnParams,
323324
TConnState,
@@ -339,7 +340,7 @@ interface BaseActorConfig<
339340
* @throws Throw an error to reject the connection
340341
*/
341342
onBeforeConnect?: (
342-
c: OnBeforeConnectContext<TState, TVars, TInput, TDatabase>,
343+
c: BeforeConnectContext<TState, TVars, TInput, TDatabase>,
343344
params: TConnParams,
344345
) => void | Promise<void>;
345346

@@ -353,7 +354,7 @@ interface BaseActorConfig<
353354
* @returns Void or a Promise that resolves when connection handling is complete
354355
*/
355356
onConnect?: (
356-
c: OnConnectContext<
357+
c: ConnectContext<
357358
TState,
358359
TConnParams,
359360
TConnState,
@@ -374,7 +375,7 @@ interface BaseActorConfig<
374375
* @returns Void or a Promise that resolves when disconnect handling is complete
375376
*/
376377
onDisconnect?: (
377-
c: ActorContext<
378+
c: DisconnectContext<
378379
TState,
379380
TConnParams,
380381
TConnState,
@@ -398,7 +399,7 @@ interface BaseActorConfig<
398399
* @returns The modified output to send to the client
399400
*/
400401
onBeforeActionResponse?: <Out>(
401-
c: ActorContext<
402+
c: BeforeActionResponseContext<
402403
TState,
403404
TConnParams,
404405
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: 40 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,14 @@ 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/mod";
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 {
8+
ActorDefinition,
9+
AnyActorDefinition,
10+
} from "../../definition";
11+
import type { ActorInstance, SaveStateOptions } from "../../instance/mod";
12+
import type { Schedule } from "../../schedule";
913

1014
/**
1115
* ActorContext class that provides access to actor methods and state
@@ -42,16 +46,23 @@ export class ActorContext<
4246

4347
/**
4448
* Get the actor state
49+
*
50+
* @remarks
51+
* This property is not available in `createState` since the state hasn't been created yet.
4552
*/
46-
get state(): TState {
47-
return this.#actor.state;
53+
get state(): TState extends never ? never : TState {
54+
return this.#actor.state as TState extends never ? never : TState;
4855
}
4956

5057
/**
5158
* Get the actor variables
59+
*
60+
* @remarks
61+
* This property is not available in `createVars` since the variables haven't been created yet.
62+
* Variables are only available if you define `vars` or `createVars` in your actor config.
5263
*/
53-
get vars(): TVars {
54-
return this.#actor.vars;
64+
get vars(): TVars extends never ? never : TVars {
65+
return this.#actor.vars as TVars extends never ? never : TVars;
5566
}
5667

5768
/**
@@ -125,11 +136,18 @@ export class ActorContext<
125136

126137
/**
127138
* Gets the database.
139+
*
128140
* @experimental
141+
* @remarks
142+
* This property is only available if you define a `db` provider in your actor config.
129143
* @throws {DatabaseNotEnabled} If the database is not enabled.
130144
*/
131-
get db(): InferDatabaseClient<TDatabase> {
132-
return this.#actor.db;
145+
get db(): TDatabase extends never
146+
? never
147+
: InferDatabaseClient<TDatabase> {
148+
return this.#actor.db as TDatabase extends never
149+
? never
150+
: InferDatabaseClient<TDatabase>;
133151
}
134152

135153
/**
@@ -177,3 +195,15 @@ export class ActorContext<
177195
this.#actor.startDestroy();
178196
}
179197
}
198+
199+
export type ActorContextOf<AD extends AnyActorDefinition> = AD extends ActorDefinition<
200+
infer S,
201+
infer CP,
202+
infer CS,
203+
infer V,
204+
infer I,
205+
infer DB extends AnyDatabaseProvider,
206+
any
207+
>
208+
? ActorContext<S, CP, CS, V, I, DB>
209+
: 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;

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

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)