Skip to content
This repository was archived by the owner on Oct 22, 2025. It is now read-only.

Commit efce452

Browse files
NathanFlurryNathanFlurry
committed
chore: add generic types for input & auth data (#993)
1 parent 192e573 commit efce452

File tree

12 files changed

+226
-112
lines changed

12 files changed

+226
-112
lines changed

‎packages/core/src/client/worker-common.ts‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,17 @@
1-
import type { AnyWorkerDefinition, WorkerDefinition } from "@/worker/definition";
1+
import type {
2+
AnyWorkerDefinition,
3+
WorkerDefinition,
4+
} from "@/worker/definition";
25
import type * as protoHttpResolve from "@/worker/protocol/http/resolve";
36
import type { Encoding } from "@/worker/protocol/serde";
47
import type { WorkerQuery } from "@/manager/protocol/query";
58
import { logger } from "./log";
69
import * as errors from "./errors";
710
import { sendHttpRequest } from "./utils";
8-
import { HEADER_WORKER_QUERY, HEADER_ENCODING } from "@/worker/router-endpoints";
11+
import {
12+
HEADER_WORKER_QUERY,
13+
HEADER_ENCODING,
14+
} from "@/worker/router-endpoints";
915

1016
/**
1117
* Action function returned by Worker connections and handles.
@@ -27,11 +33,10 @@ export type WorkerActionFunction<
2733
* Maps action methods from worker definition to typed function signatures.
2834
*/
2935
export type WorkerDefinitionActions<AD extends AnyWorkerDefinition> =
30-
AD extends WorkerDefinition<any, any, any, any, infer R>
36+
AD extends WorkerDefinition<any, any, any, any, any, any, infer R>
3137
? {
3238
[K in keyof R]: R[K] extends (...args: infer Args) => infer Return
3339
? WorkerActionFunction<Args, Return>
3440
: never;
3541
}
3642
: never;
37-

‎packages/core/src/registry/config.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@ export type WorkerPeerConfig = z.infer<typeof WorkerPeerConfigSchema>;
4343

4444
export const WorkersSchema = z.record(
4545
z.string(),
46-
z.custom<WorkerDefinition<any, any, any, any, any>>(),
46+
z.custom<WorkerDefinition<any, any, any, any, any, any, any>>(),
4747
);
4848
export type Workers = z.infer<typeof WorkersSchema>;
4949

‎packages/core/src/worker/action.ts‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,8 @@ import { WorkerContext } from "./context";
1313
*
1414
* @typeParam A Worker this action belongs to
1515
*/
16-
export class ActionContext<S, CP, CS, V> {
17-
#workerContext: WorkerContext<S, CP, CS, V>;
16+
export class ActionContext<S, CP, CS, V, I, AD> {
17+
#workerContext: WorkerContext<S, CP, CS, V, I, AD>;
1818

1919
/**
2020
* Should not be called directly.
@@ -23,8 +23,8 @@ export class ActionContext<S, CP, CS, V> {
2323
* @param conn - The connection associated with the action
2424
*/
2525
constructor(
26-
workerContext: WorkerContext<S, CP, CS, V>,
27-
public readonly conn: Conn<S, CP, CS, V>,
26+
workerContext: WorkerContext<S, CP, CS, V, I, AD>,
27+
public readonly conn: Conn<S, CP, CS, V, I, AD>,
2828
) {
2929
this.#workerContext = workerContext;
3030
}
@@ -95,7 +95,7 @@ export class ActionContext<S, CP, CS, V> {
9595
/**
9696
* Gets the map of connections.
9797
*/
98-
get conns(): Map<ConnId, Conn<S, CP, CS, V>> {
98+
get conns(): Map<ConnId, Conn<S, CP, CS, V, I, AD>> {
9999
return this.#workerContext.conns;
100100
}
101101

‎packages/core/src/worker/config.ts‎

Lines changed: 71 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -75,12 +75,12 @@ export const WorkerConfigSchema = z
7575
},
7676
);
7777

78-
export interface OnCreateOptions {
79-
input?: unknown;
78+
export interface OnCreateOptions<I> {
79+
input?: I;
8080
}
8181

82-
export interface CreateStateOptions {
83-
input?: unknown;
82+
export interface CreateStateOptions<I> {
83+
input?: I;
8484
}
8585

8686
export interface OnConnectOptions<CP> {
@@ -98,12 +98,12 @@ export interface OnConnectOptions<CP> {
9898
// This must have only one or the other or else S will not be able to be inferred
9999
//
100100
// Data returned from this handler will be available on `c.state`.
101-
type CreateState<S, CP, CS, V> =
101+
type CreateState<S, CP, CS, V, I, AD> =
102102
| { state: S }
103103
| {
104104
createState: (
105-
c: WorkerContext<undefined, undefined, undefined, undefined>,
106-
opts: CreateStateOptions,
105+
c: WorkerContext<undefined, undefined, undefined, undefined, undefined, undefined>,
106+
opts: CreateStateOptions<I>,
107107
) => S | Promise<S>;
108108
}
109109
| Record<never, never>;
@@ -113,11 +113,11 @@ type CreateState<S, CP, CS, V> =
113113
// This must have only one or the other or else S will not be able to be inferred
114114
//
115115
// Data returned from this handler will be available on `c.conn.state`.
116-
type CreateConnState<S, CP, CS, V> =
116+
type CreateConnState<S, CP, CS, V, I, AD> =
117117
| { connState: CS }
118118
| {
119119
createConnState: (
120-
c: WorkerContext<undefined, undefined, undefined, undefined>,
120+
c: WorkerContext<undefined, undefined, undefined, undefined, undefined, undefined>,
121121
opts: OnConnectOptions<CP>,
122122
) => CS | Promise<CS>;
123123
}
@@ -129,7 +129,7 @@ type CreateConnState<S, CP, CS, V> =
129129
/**
130130
* @experimental
131131
*/
132-
type CreateVars<S, CP, CS, V> =
132+
type CreateVars<S, CP, CS, V, I, AD> =
133133
| {
134134
/**
135135
* @experimental
@@ -141,20 +141,23 @@ type CreateVars<S, CP, CS, V> =
141141
* @experimental
142142
*/
143143
createVars: (
144-
c: WorkerContext<undefined, undefined, undefined, undefined>,
144+
c: WorkerContext<undefined, undefined, undefined, undefined, undefined, undefined>,
145145
driverCtx: unknown,
146146
) => V | Promise<V>;
147147
}
148148
| Record<never, never>;
149149

150-
export interface Actions<S, CP, CS, V> {
151-
[Action: string]: (c: ActionContext<S, CP, CS, V>, ...args: any[]) => any;
150+
export interface Actions<S, CP, CS, V, I, AD> {
151+
[Action: string]: (
152+
c: ActionContext<S, CP, CS, V, I, AD>,
153+
...args: any[]
154+
) => any;
152155
}
153156

154-
//export type WorkerConfig<S, CP, CS, V> = BaseWorkerConfig<S, CP, CS, V> &
155-
// WorkerConfigLifecycle<S, CP, CS, V> &
156-
// CreateState<S, CP, CS, V> &
157-
// CreateConnState<S, CP, CS, V>;
157+
//export type WorkerConfig<S, CP, CS, V, I, AD> = BaseWorkerConfig<S, CP, CS, V, I, AD> &
158+
// WorkerConfigLifecycle<S, CP, CS, V, I, AD> &
159+
// CreateState<S, CP, CS, V, I, AD> &
160+
// CreateConnState<S, CP, CS, V, I, AD>;
158161

159162
/**
160163
* @experimental
@@ -170,7 +173,15 @@ interface OnAuthOptions<CP> {
170173
params: CP;
171174
}
172175

173-
interface BaseWorkerConfig<S, CP, CS, V, R extends Actions<S, CP, CS, V>> {
176+
interface BaseWorkerConfig<
177+
S,
178+
CP,
179+
CS,
180+
V,
181+
I,
182+
AD,
183+
R extends Actions<S, CP, CS, V, I, AD>,
184+
> {
174185
/**
175186
* Called on the HTTP server before clients can interact with the worker.
176187
*
@@ -196,7 +207,7 @@ interface BaseWorkerConfig<S, CP, CS, V, R extends Actions<S, CP, CS, V>> {
196207
* @returns Authentication data to attach to connections (must be serializable)
197208
* @throws Throw an error to deny access to the worker
198209
*/
199-
onAuth?: (opts: OnAuthOptions<CP>) => unknown | Promise<unknown>;
210+
onAuth?: (opts: OnAuthOptions<CP>) => AD | Promise<AD>;
200211

201212
/**
202213
* Called when the worker is first initialized.
@@ -205,8 +216,8 @@ interface BaseWorkerConfig<S, CP, CS, V, R extends Actions<S, CP, CS, V>> {
205216
* This is called before any other lifecycle hooks.
206217
*/
207218
onCreate?: (
208-
c: WorkerContext<S, CP, CS, V>,
209-
opts: OnCreateOptions,
219+
c: WorkerContext<S, CP, CS, V, I, AD>,
220+
opts: OnCreateOptions<I>,
210221
) => void | Promise<void>;
211222

212223
/**
@@ -217,7 +228,7 @@ interface BaseWorkerConfig<S, CP, CS, V, R extends Actions<S, CP, CS, V>> {
217228
*
218229
* @returns Void or a Promise that resolves when startup is complete
219230
*/
220-
onStart?: (c: WorkerContext<S, CP, CS, V>) => void | Promise<void>;
231+
onStart?: (c: WorkerContext<S, CP, CS, V, I, AD>) => void | Promise<void>;
221232

222233
/**
223234
* Called when the worker's state changes.
@@ -227,7 +238,7 @@ interface BaseWorkerConfig<S, CP, CS, V, R extends Actions<S, CP, CS, V>> {
227238
*
228239
* @param newState The updated state
229240
*/
230-
onStateChange?: (c: WorkerContext<S, CP, CS, V>, newState: S) => void;
241+
onStateChange?: (c: WorkerContext<S, CP, CS, V, I, AD>, newState: S) => void;
231242

232243
/**
233244
* Called before a client connects to the worker.
@@ -250,7 +261,7 @@ interface BaseWorkerConfig<S, CP, CS, V, R extends Actions<S, CP, CS, V>> {
250261
* @throws Throw an error to reject the connection
251262
*/
252263
onBeforeConnect?: (
253-
c: WorkerContext<S, CP, CS, V>,
264+
c: WorkerContext<S, CP, CS, V, I, AD>,
254265
opts: OnConnectOptions<CP>,
255266
) => void | Promise<void>;
256267

@@ -264,8 +275,8 @@ interface BaseWorkerConfig<S, CP, CS, V, R extends Actions<S, CP, CS, V>> {
264275
* @returns Void or a Promise that resolves when connection handling is complete
265276
*/
266277
onConnect?: (
267-
c: WorkerContext<S, CP, CS, V>,
268-
conn: Conn<S, CP, CS, V>,
278+
c: WorkerContext<S, CP, CS, V, I, AD>,
279+
conn: Conn<S, CP, CS, V, I, AD>,
269280
) => void | Promise<void>;
270281

271282
/**
@@ -278,8 +289,8 @@ interface BaseWorkerConfig<S, CP, CS, V, R extends Actions<S, CP, CS, V>> {
278289
* @returns Void or a Promise that resolves when disconnect handling is complete
279290
*/
280291
onDisconnect?: (
281-
c: WorkerContext<S, CP, CS, V>,
282-
conn: Conn<S, CP, CS, V>,
292+
c: WorkerContext<S, CP, CS, V, I, AD>,
293+
conn: Conn<S, CP, CS, V, I, AD>,
283294
) => void | Promise<void>;
284295

285296
/**
@@ -295,7 +306,7 @@ interface BaseWorkerConfig<S, CP, CS, V, R extends Actions<S, CP, CS, V>> {
295306
* @returns The modified output to send to the client
296307
*/
297308
onBeforeActionResponse?: <Out>(
298-
c: WorkerContext<S, CP, CS, V>,
309+
c: WorkerContext<S, CP, CS, V, I, AD>,
299310
name: string,
300311
args: unknown[],
301312
output: Out,
@@ -307,7 +318,7 @@ interface BaseWorkerConfig<S, CP, CS, V, R extends Actions<S, CP, CS, V>> {
307318
// 1. Infer schema
308319
// 2. Omit keys that we'll manually define (because of generics)
309320
// 3. Define our own types that have generic constraints
310-
export type WorkerConfig<S, CP, CS, V> = Omit<
321+
export type WorkerConfig<S, CP, CS, V, I, AD> = Omit<
311322
z.infer<typeof WorkerConfigSchema>,
312323
| "actions"
313324
| "onAuth"
@@ -325,18 +336,20 @@ export type WorkerConfig<S, CP, CS, V> = Omit<
325336
| "vars"
326337
| "createVars"
327338
> &
328-
BaseWorkerConfig<S, CP, CS, V, Actions<S, CP, CS, V>> &
329-
CreateState<S, CP, CS, V> &
330-
CreateConnState<S, CP, CS, V> &
331-
CreateVars<S, CP, CS, V>;
339+
BaseWorkerConfig<S, CP, CS, V, I, AD, Actions<S, CP, CS, V, I, AD>> &
340+
CreateState<S, CP, CS, V, I, AD> &
341+
CreateConnState<S, CP, CS, V, I, AD> &
342+
CreateVars<S, CP, CS, V, I, AD>;
332343

333344
// See description on `WorkerConfig`
334345
export type WorkerConfigInput<
335346
S,
336347
CP,
337348
CS,
338349
V,
339-
R extends Actions<S, CP, CS, V>,
350+
I,
351+
AD,
352+
R extends Actions<S, CP, CS, V, I, AD>,
340353
> = Omit<
341354
z.input<typeof WorkerConfigSchema>,
342355
| "actions"
@@ -355,16 +368,31 @@ export type WorkerConfigInput<
355368
| "vars"
356369
| "createVars"
357370
> &
358-
BaseWorkerConfig<S, CP, CS, V, R> &
359-
CreateState<S, CP, CS, V> &
360-
CreateConnState<S, CP, CS, V> &
361-
CreateVars<S, CP, CS, V>;
371+
BaseWorkerConfig<S, CP, CS, V, I, AD, R> &
372+
CreateState<S, CP, CS, V, I, AD> &
373+
CreateConnState<S, CP, CS, V, I, AD> &
374+
CreateVars<S, CP, CS, V, I, AD>;
362375

363376
// For testing type definitions:
364-
export function test<S, CP, CS, V, R extends Actions<S, CP, CS, V>>(
365-
input: WorkerConfigInput<S, CP, CS, V, R>,
366-
): WorkerConfig<S, CP, CS, V> {
367-
const config = WorkerConfigSchema.parse(input) as WorkerConfig<S, CP, CS, V>;
377+
export function test<
378+
S,
379+
CP,
380+
CS,
381+
V,
382+
I,
383+
AD,
384+
R extends Actions<S, CP, CS, V, I, AD>,
385+
>(
386+
input: WorkerConfigInput<S, CP, CS, V, I, AD, R>,
387+
): WorkerConfig<S, CP, CS, V, I, AD> {
388+
const config = WorkerConfigSchema.parse(input) as WorkerConfig<
389+
S,
390+
CP,
391+
CS,
392+
V,
393+
I,
394+
AD
395+
>;
368396
return config;
369397
}
370398

‎packages/core/src/worker/connection.ts‎

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ export function generateConnToken(): string {
1717

1818
export type ConnId = string;
1919

20-
export type AnyConn = Conn<any, any, any, any>;
20+
export type AnyConn = Conn<any, any, any, any, any, any>;
2121

2222
/**
2323
* Represents a client connection to a worker.
@@ -26,13 +26,13 @@ export type AnyConn = Conn<any, any, any, any>;
2626
*
2727
* @see {@link https://rivet.gg/docs/connections|Connection Documentation}
2828
*/
29-
export class Conn<S, CP, CS, V> {
29+
export class Conn<S, CP, CS, V, I, AD> {
3030
subscriptions: Set<string> = new Set<string>();
3131

3232
#stateEnabled: boolean;
3333

3434
// TODO: Remove this cyclical reference
35-
#worker: WorkerInstance<S, CP, CS, V>;
35+
#worker: WorkerInstance<S, CP, CS, V, I, AD>;
3636

3737
/**
3838
* The proxied state that notifies of changes automatically.
@@ -103,7 +103,7 @@ export class Conn<S, CP, CS, V> {
103103
* @protected
104104
*/
105105
public constructor(
106-
worker: WorkerInstance<S, CP, CS, V>,
106+
worker: WorkerInstance<S, CP, CS, V, I, AD>,
107107
persist: PersistedConn<CP, CS>,
108108
driver: ConnDriver,
109109
stateEnabled: boolean,
@@ -157,6 +157,11 @@ export class Conn<S, CP, CS, V> {
157157
* @param reason - The reason for disconnection.
158158
*/
159159
public async disconnect(reason?: string) {
160-
await this.#driver.disconnect(this.#worker, this, this.__persist.ds, reason);
160+
await this.#driver.disconnect(
161+
this.#worker,
162+
this,
163+
this.__persist.ds,
164+
reason,
165+
);
161166
}
162167
}

‎packages/core/src/worker/context.ts‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,10 +8,10 @@ import { Schedule } from "./schedule";
88
/**
99
* WorkerContext class that provides access to worker methods and state
1010
*/
11-
export class WorkerContext<S, CP, CS, V> {
12-
#worker: WorkerInstance<S, CP, CS, V>;
11+
export class WorkerContext<S, CP, CS, V, I, AD> {
12+
#worker: WorkerInstance<S, CP, CS, V, I, AD>;
1313

14-
constructor(worker: WorkerInstance<S, CP, CS, V>) {
14+
constructor(worker: WorkerInstance<S, CP, CS, V, I, AD>) {
1515
this.#worker = worker;
1616
}
1717

@@ -84,7 +84,7 @@ export class WorkerContext<S, CP, CS, V> {
8484
/**
8585
* Gets the map of connections.
8686
*/
87-
get conns(): Map<ConnId, Conn<S, CP, CS, V>> {
87+
get conns(): Map<ConnId, Conn<S, CP, CS, V, I, AD>> {
8888
return this.#worker.conns;
8989
}
9090

0 commit comments

Comments
 (0)