@rivetkit/engine-runner 2.0.24-rc.1 → 2.0.24

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
package/dist/mod.d.cts CHANGED
@@ -1,19 +1,147 @@
1
1
  import * as protocol from '@rivetkit/engine-runner-protocol';
2
+ import { GatewayId, RequestId } from '@rivetkit/engine-runner-protocol';
2
3
  import { Logger } from 'pino';
4
+ import WebSocket from 'ws';
3
5
 
4
- interface ActorInstance {
5
- actorId: string;
6
- generation: number;
7
- config: ActorConfig;
8
- requests: Set<string>;
9
- webSockets: Set<string>;
6
+ interface PendingRequest {
7
+ resolve: (response: Response) => void;
8
+ reject: (error: Error) => void;
9
+ streamController?: ReadableStreamDefaultController<Uint8Array>;
10
+ actorId?: string;
11
+ gatewayId?: GatewayId;
12
+ requestId?: RequestId;
13
+ clientMessageIndex: number;
14
+ }
15
+ interface HibernatingWebSocketMetadata {
16
+ gatewayId: GatewayId;
17
+ requestId: RequestId;
18
+ clientMessageIndex: number;
19
+ serverMessageIndex: number;
20
+ path: string;
21
+ headers: Record<string, string>;
22
+ }
23
+ declare class Tunnel {
24
+ #private;
25
+ get log(): Logger | undefined;
26
+ constructor(runner: Runner);
27
+ start(): void;
28
+ resendBufferedEvents(): void;
29
+ shutdown(): void;
30
+ restoreHibernatingRequests(actorId: string, metaEntries: HibernatingWebSocketMetadata[]): Promise<void>;
31
+ addRequestToActor(gatewayId: GatewayId, requestId: RequestId, actorId: string): void;
32
+ getRequestActor(gatewayId: GatewayId, requestId: RequestId): RunnerActor | undefined;
33
+ getAndWaitForRequestActor(gatewayId: GatewayId, requestId: RequestId): Promise<RunnerActor | undefined>;
34
+ closeActiveRequests(actor: RunnerActor): void;
35
+ handleTunnelMessage(message: protocol.ToClientTunnelMessage): Promise<void>;
36
+ sendHibernatableWebSocketMessageAck(gatewayId: ArrayBuffer, requestId: ArrayBuffer, clientMessageIndex: number): void;
37
+ }
38
+
39
+ /**
40
+ * Polyfill for Promise.withResolvers().
41
+ *
42
+ * This is specifically for Cloudflare Workers. Their implementation of Promise.withResolvers does not work correctly.
43
+ */
44
+ declare function promiseWithResolvers<T>(): {
45
+ promise: Promise<T>;
46
+ resolve: (value: T | PromiseLike<T>) => void;
47
+ reject: (reason?: any) => void;
48
+ };
49
+ declare function idToStr(id: ArrayBuffer): string;
50
+
51
+ declare const HIBERNATABLE_SYMBOL: unique symbol;
52
+ declare class WebSocketTunnelAdapter {
53
+ #private;
54
+ /** @experimental */
55
+ readonly request: Request;
56
+ get [HIBERNATABLE_SYMBOL](): boolean;
57
+ constructor(tunnel: Tunnel, actorId: string, requestId: string, serverMessageIndex: number, hibernatable: boolean, isRestoringHibernatable: boolean,
58
+ /** @experimental */
59
+ request: Request, sendCallback: (data: ArrayBuffer | string, isBinary: boolean) => void, closeCallback: (code?: number, reason?: string) => void);
60
+ get bufferedAmount(): number;
61
+ _handleOpen(requestId: ArrayBuffer): void;
62
+ _handleMessage(requestId: ArrayBuffer, data: string | Uint8Array, serverMessageIndex: number, isBinary: boolean): boolean;
63
+ _handleClose(_requestId: ArrayBuffer, code?: number, reason?: string): void;
64
+ _handleError(error: Error): void;
65
+ _closeWithoutCallback(code?: number, reason?: string): void;
66
+ get readyState(): number;
67
+ get binaryType(): string;
68
+ set binaryType(value: string);
69
+ get extensions(): string;
70
+ get protocol(): string;
71
+ get url(): string;
72
+ get onopen(): ((this: any, ev: any) => any) | null;
73
+ set onopen(value: ((this: any, ev: any) => any) | null);
74
+ get onclose(): ((this: any, ev: any) => any) | null;
75
+ set onclose(value: ((this: any, ev: any) => any) | null);
76
+ get onerror(): ((this: any, ev: any) => any) | null;
77
+ set onerror(value: ((this: any, ev: any) => any) | null);
78
+ get onmessage(): ((this: any, ev: any) => any) | null;
79
+ set onmessage(value: ((this: any, ev: any) => any) | null);
80
+ send(data: string | ArrayBuffer | ArrayBufferView | Blob | Buffer): void;
81
+ close(code?: number, reason?: string): void;
82
+ addEventListener(type: string, listener: (event: any) => void, options?: boolean | any): void;
83
+ removeEventListener(type: string, listener: (event: any) => void, options?: boolean | any): void;
84
+ dispatchEvent(event: any): boolean;
85
+ static readonly CONNECTING = 0;
86
+ static readonly OPEN = 1;
87
+ static readonly CLOSING = 2;
88
+ static readonly CLOSED = 3;
89
+ readonly CONNECTING = 0;
90
+ readonly OPEN = 1;
91
+ readonly CLOSING = 2;
92
+ readonly CLOSED = 3;
93
+ ping(data?: any, mask?: boolean, cb?: (err: Error) => void): void;
94
+ pong(data?: any, mask?: boolean, cb?: (err: Error) => void): void;
95
+ /** @experimental */
96
+ terminate(): void;
10
97
  }
98
+
11
99
  interface ActorConfig {
12
100
  name: string;
13
101
  key: string | null;
14
102
  createTs: bigint;
15
103
  input: Uint8Array | null;
16
104
  }
105
+ declare class RunnerActor {
106
+ /**
107
+ * List of hibernating requests provided by the gateway on actor start.
108
+ * This represents the WebSocket connections that the gateway knows about.
109
+ **/
110
+ hibernatingRequests: readonly protocol.HibernatingRequest[];
111
+ actorId: string;
112
+ generation: number;
113
+ config: ActorConfig;
114
+ pendingRequests: Array<{
115
+ gatewayId: protocol.GatewayId;
116
+ requestId: protocol.RequestId;
117
+ request: PendingRequest;
118
+ }>;
119
+ webSockets: Array<{
120
+ gatewayId: protocol.GatewayId;
121
+ requestId: protocol.RequestId;
122
+ ws: WebSocketTunnelAdapter;
123
+ }>;
124
+ actorStartPromise: ReturnType<typeof promiseWithResolvers<void>>;
125
+ /**
126
+ * If restoreHibernatingRequests has been called. This is used to assert
127
+ * that the caller is implemented correctly.
128
+ **/
129
+ hibernationRestored: boolean;
130
+ constructor(actorId: string, generation: number, config: ActorConfig,
131
+ /**
132
+ * List of hibernating requests provided by the gateway on actor start.
133
+ * This represents the WebSocket connections that the gateway knows about.
134
+ **/
135
+ hibernatingRequests: readonly protocol.HibernatingRequest[]);
136
+ getPendingRequest(gatewayId: protocol.GatewayId, requestId: protocol.RequestId): PendingRequest | undefined;
137
+ createPendingRequest(gatewayId: protocol.GatewayId, requestId: protocol.RequestId, clientMessageIndex: number): void;
138
+ createPendingRequestWithStreamController(gatewayId: protocol.GatewayId, requestId: protocol.RequestId, clientMessageIndex: number, streamController: ReadableStreamDefaultController<Uint8Array>): void;
139
+ deletePendingRequest(gatewayId: protocol.GatewayId, requestId: protocol.RequestId): void;
140
+ getWebSocket(gatewayId: protocol.GatewayId, requestId: protocol.RequestId): WebSocketTunnelAdapter | undefined;
141
+ setWebSocket(gatewayId: protocol.GatewayId, requestId: protocol.RequestId, ws: WebSocketTunnelAdapter): void;
142
+ deleteWebSocket(gatewayId: protocol.GatewayId, requestId: protocol.RequestId): void;
143
+ }
144
+
17
145
  interface RunnerConfig {
18
146
  logger?: Logger;
19
147
  version: number;
@@ -32,17 +160,105 @@ interface RunnerConfig {
32
160
  onConnected: () => void;
33
161
  onDisconnected: (code: number, reason: string) => void;
34
162
  onShutdown: () => void;
35
- fetch: (runner: Runner, actorId: string, requestId: protocol.RequestId, request: Request) => Promise<Response>;
36
- websocket?: (runner: Runner, actorId: string, ws: any, requestId: protocol.RequestId, request: Request) => Promise<void>;
163
+ /** Called when receiving a network request. */
164
+ fetch: (runner: Runner, actorId: string, gatewayId: protocol.GatewayId, requestId: protocol.RequestId, request: Request) => Promise<Response>;
165
+ /**
166
+ * Called when receiving a WebSocket connection.
167
+ *
168
+ * All event listeners must be added synchronously inside this function or
169
+ * else events may be missed. The open event will fire immediately after
170
+ * this function finishes.
171
+ *
172
+ * Any errors thrown here will disconnect the WebSocket immediately.
173
+ *
174
+ * While `path` and `headers` are partially redundant to the data in the
175
+ * `Request`, they may vary slightly from the actual content of `Request`.
176
+ * Prefer to persist the `path` and `headers` properties instead of the
177
+ * `Request` itself.
178
+ *
179
+ * ## Hibernating Web Sockets
180
+ *
181
+ * ### Implementation Requirements
182
+ *
183
+ * **Requirement 1: Persist HWS Immediately**
184
+ *
185
+ * This is responsible for persisting hibernatable WebSockets immediately
186
+ * (do not wait for open event). It is not time sensitive to flush the
187
+ * connection state. If this fails to persist the HWS, the client's
188
+ * WebSocket will be disconnected on next wake in the call to
189
+ * `Tunnel::restoreHibernatingRequests` since the connection entry will not
190
+ * exist.
191
+ *
192
+ * **Requirement 2: Persist Message Index On `message`**
193
+ *
194
+ * In the `message` event listener, this handler must persist the message
195
+ * index from the event. The request ID is available at
196
+ * `event.rivetRequestId` and message index at `event.rivetMessageIndex`.
197
+ *
198
+ * The message index should not be flushed immediately. Instead, this
199
+ * should:
200
+ *
201
+ * - Debounce calls to persist the message index
202
+ * - After each persist, call
203
+ * `Runner::sendHibernatableWebSocketMessageAck` to acknowledge the
204
+ * message
205
+ *
206
+ * This mechanism allows us to buffer messages on the gateway so we can
207
+ * batch-persist events on our end on a given interval.
208
+ *
209
+ * If this fails to persist, then the gateway will replay unacked
210
+ * messages when the actor starts again.
211
+ *
212
+ * **Requirement 3: Remove HWS From Storage On `close`**
213
+ *
214
+ * This handler should add an event listener for `close` to remove the
215
+ * connection from storage.
216
+ *
217
+ * If the connection remove fails to persist, the close event will be
218
+ * called again on the next actor start in
219
+ * `Tunnel::restoreHibernatingRequests` since there will be no request for
220
+ * the given connection.
221
+ *
222
+ * ### Restoring Connections
223
+ *
224
+ * The user of this library is responsible for:
225
+ * 1. Loading all persisted hibernatable WebSocket metadata for an actor
226
+ * 2. Calling `Runner::restoreHibernatingRequests` with this metadata at
227
+ * the end of `onActorStart`
228
+ *
229
+ * `restoreHibernatingRequests` will restore all connections and attach
230
+ * the appropriate event listeners.
231
+ *
232
+ * ### No Open Event On Restoration
233
+ *
234
+ * When restoring a HWS, the open event will not be called again. It will
235
+ * go straight to the message or close event.
236
+ */
237
+ websocket: (runner: Runner, actorId: string, ws: any, gatewayId: protocol.GatewayId, requestId: protocol.RequestId, request: Request, path: string, headers: Record<string, string>, isHibernatable: boolean, isRestoringHibernatable: boolean) => Promise<void>;
238
+ hibernatableWebSocket: {
239
+ /**
240
+ * Determines if a WebSocket can continue to live while an actor goes to
241
+ * sleep.
242
+ */
243
+ canHibernate: (actorId: string, gatewayId: ArrayBuffer, requestId: ArrayBuffer, request: Request) => boolean;
244
+ };
245
+ /**
246
+ * Called when an actor starts.
247
+ *
248
+ * This callback is responsible for:
249
+ * 1. Initializing the actor instance
250
+ * 2. Loading all persisted hibernatable WebSocket metadata for this actor
251
+ * 3. Calling `Runner::restoreHibernatingRequests` with the loaded metadata
252
+ * to restore hibernatable WebSocket connections
253
+ *
254
+ * The actor should not be marked as "ready" until after
255
+ * `restoreHibernatingRequests` completes to ensure all hibernatable
256
+ * connections are fully restored before the actor processes new requests.
257
+ */
37
258
  onActorStart: (actorId: string, generation: number, config: ActorConfig) => Promise<void>;
38
259
  onActorStop: (actorId: string, generation: number) => Promise<void>;
39
- getActorHibernationConfig: (actorId: string, requestId: ArrayBuffer, request: Request) => HibernationConfig;
40
260
  noAutoShutdown?: boolean;
41
261
  }
42
- interface HibernationConfig {
43
- enabled: boolean;
44
- lastMsgIndex: number | undefined;
45
- }
46
262
  interface KvListOptions {
47
263
  reverse?: boolean;
48
264
  limit?: number;
@@ -50,14 +266,17 @@ interface KvListOptions {
50
266
  declare class Runner {
51
267
  #private;
52
268
  get config(): RunnerConfig;
269
+ __pegboardWebSocket?: WebSocket;
53
270
  runnerId?: string;
54
271
  get log(): Logger | undefined;
55
272
  constructor(config: RunnerConfig);
56
273
  sleepActor(actorId: string, generation?: number): void;
57
274
  stopActor(actorId: string, generation?: number): Promise<void>;
58
275
  forceStopActor(actorId: string, generation?: number): Promise<void>;
59
- getActor(actorId: string, generation?: number): ActorInstance | undefined;
276
+ getActor(actorId: string, generation?: number): RunnerActor | undefined;
277
+ getAndWaitForActor(actorId: string, generation?: number): Promise<RunnerActor | undefined>;
60
278
  hasActor(actorId: string, generation?: number): boolean;
279
+ get actors(): Map<string, RunnerActor>;
61
280
  start(): Promise<void>;
62
281
  shutdown(immediate: boolean, exit?: boolean): Promise<void>;
63
282
  get pegboardEndpoint(): string;
@@ -71,10 +290,37 @@ declare class Runner {
71
290
  kvDrop(actorId: string): Promise<void>;
72
291
  setAlarm(actorId: string, alarmTs: number | null, generation?: number): void;
73
292
  clearAlarm(actorId: string, generation?: number): void;
74
- __webSocketReady(): boolean;
293
+ /** Asserts WebSocket exists and is ready. */
294
+ __webSocketReady(): this is this & {
295
+ __pegboardWebSocket: NonNullable<Runner["__pegboardWebSocket"]>;
296
+ };
75
297
  __sendToServer(message: protocol.ToServer): void;
76
- sendWebsocketMessageAck(requestId: ArrayBuffer, index: number): void;
298
+ sendHibernatableWebSocketMessageAck(gatewayId: ArrayBuffer, requestId: ArrayBuffer, index: number): void;
299
+ /**
300
+ * Restores hibernatable WebSocket connections for an actor.
301
+ *
302
+ * This method should be called at the end of `onActorStart` after the
303
+ * actor instance is fully initialized.
304
+ *
305
+ * This method will:
306
+ * - Restore all provided hibernatable WebSocket connections
307
+ * - Attach event listeners to the restored WebSockets
308
+ * - Close any WebSocket connections that failed to restore
309
+ *
310
+ * The provided metadata list should include all hibernatable WebSockets
311
+ * that were persisted for this actor. The gateway will automatically
312
+ * close any connections that are not restored (i.e., not included in
313
+ * this list).
314
+ *
315
+ * **Important:** This method must be called after `onActorStart` completes
316
+ * and before marking the actor as "ready" to ensure all hibernatable
317
+ * connections are fully restored.
318
+ *
319
+ * @param actorId - The ID of the actor to restore connections for
320
+ * @param metaEntries - Array of hibernatable WebSocket metadata to restore
321
+ */
322
+ restoreHibernatingRequests(actorId: string, metaEntries: HibernatingWebSocketMetadata[]): Promise<void>;
77
323
  getServerlessInitPacket(): string | undefined;
78
324
  }
79
325
 
80
- export { type ActorConfig, type ActorInstance, type HibernationConfig, type KvListOptions, Runner, type RunnerConfig };
326
+ export { type ActorConfig, type HibernatingWebSocketMetadata, type KvListOptions, Runner, RunnerActor, type RunnerConfig, idToStr };
package/dist/mod.d.ts CHANGED
@@ -1,19 +1,147 @@
1
1
  import * as protocol from '@rivetkit/engine-runner-protocol';
2
+ import { GatewayId, RequestId } from '@rivetkit/engine-runner-protocol';
2
3
  import { Logger } from 'pino';
4
+ import WebSocket from 'ws';
3
5
 
4
- interface ActorInstance {
5
- actorId: string;
6
- generation: number;
7
- config: ActorConfig;
8
- requests: Set<string>;
9
- webSockets: Set<string>;
6
+ interface PendingRequest {
7
+ resolve: (response: Response) => void;
8
+ reject: (error: Error) => void;
9
+ streamController?: ReadableStreamDefaultController<Uint8Array>;
10
+ actorId?: string;
11
+ gatewayId?: GatewayId;
12
+ requestId?: RequestId;
13
+ clientMessageIndex: number;
14
+ }
15
+ interface HibernatingWebSocketMetadata {
16
+ gatewayId: GatewayId;
17
+ requestId: RequestId;
18
+ clientMessageIndex: number;
19
+ serverMessageIndex: number;
20
+ path: string;
21
+ headers: Record<string, string>;
22
+ }
23
+ declare class Tunnel {
24
+ #private;
25
+ get log(): Logger | undefined;
26
+ constructor(runner: Runner);
27
+ start(): void;
28
+ resendBufferedEvents(): void;
29
+ shutdown(): void;
30
+ restoreHibernatingRequests(actorId: string, metaEntries: HibernatingWebSocketMetadata[]): Promise<void>;
31
+ addRequestToActor(gatewayId: GatewayId, requestId: RequestId, actorId: string): void;
32
+ getRequestActor(gatewayId: GatewayId, requestId: RequestId): RunnerActor | undefined;
33
+ getAndWaitForRequestActor(gatewayId: GatewayId, requestId: RequestId): Promise<RunnerActor | undefined>;
34
+ closeActiveRequests(actor: RunnerActor): void;
35
+ handleTunnelMessage(message: protocol.ToClientTunnelMessage): Promise<void>;
36
+ sendHibernatableWebSocketMessageAck(gatewayId: ArrayBuffer, requestId: ArrayBuffer, clientMessageIndex: number): void;
37
+ }
38
+
39
+ /**
40
+ * Polyfill for Promise.withResolvers().
41
+ *
42
+ * This is specifically for Cloudflare Workers. Their implementation of Promise.withResolvers does not work correctly.
43
+ */
44
+ declare function promiseWithResolvers<T>(): {
45
+ promise: Promise<T>;
46
+ resolve: (value: T | PromiseLike<T>) => void;
47
+ reject: (reason?: any) => void;
48
+ };
49
+ declare function idToStr(id: ArrayBuffer): string;
50
+
51
+ declare const HIBERNATABLE_SYMBOL: unique symbol;
52
+ declare class WebSocketTunnelAdapter {
53
+ #private;
54
+ /** @experimental */
55
+ readonly request: Request;
56
+ get [HIBERNATABLE_SYMBOL](): boolean;
57
+ constructor(tunnel: Tunnel, actorId: string, requestId: string, serverMessageIndex: number, hibernatable: boolean, isRestoringHibernatable: boolean,
58
+ /** @experimental */
59
+ request: Request, sendCallback: (data: ArrayBuffer | string, isBinary: boolean) => void, closeCallback: (code?: number, reason?: string) => void);
60
+ get bufferedAmount(): number;
61
+ _handleOpen(requestId: ArrayBuffer): void;
62
+ _handleMessage(requestId: ArrayBuffer, data: string | Uint8Array, serverMessageIndex: number, isBinary: boolean): boolean;
63
+ _handleClose(_requestId: ArrayBuffer, code?: number, reason?: string): void;
64
+ _handleError(error: Error): void;
65
+ _closeWithoutCallback(code?: number, reason?: string): void;
66
+ get readyState(): number;
67
+ get binaryType(): string;
68
+ set binaryType(value: string);
69
+ get extensions(): string;
70
+ get protocol(): string;
71
+ get url(): string;
72
+ get onopen(): ((this: any, ev: any) => any) | null;
73
+ set onopen(value: ((this: any, ev: any) => any) | null);
74
+ get onclose(): ((this: any, ev: any) => any) | null;
75
+ set onclose(value: ((this: any, ev: any) => any) | null);
76
+ get onerror(): ((this: any, ev: any) => any) | null;
77
+ set onerror(value: ((this: any, ev: any) => any) | null);
78
+ get onmessage(): ((this: any, ev: any) => any) | null;
79
+ set onmessage(value: ((this: any, ev: any) => any) | null);
80
+ send(data: string | ArrayBuffer | ArrayBufferView | Blob | Buffer): void;
81
+ close(code?: number, reason?: string): void;
82
+ addEventListener(type: string, listener: (event: any) => void, options?: boolean | any): void;
83
+ removeEventListener(type: string, listener: (event: any) => void, options?: boolean | any): void;
84
+ dispatchEvent(event: any): boolean;
85
+ static readonly CONNECTING = 0;
86
+ static readonly OPEN = 1;
87
+ static readonly CLOSING = 2;
88
+ static readonly CLOSED = 3;
89
+ readonly CONNECTING = 0;
90
+ readonly OPEN = 1;
91
+ readonly CLOSING = 2;
92
+ readonly CLOSED = 3;
93
+ ping(data?: any, mask?: boolean, cb?: (err: Error) => void): void;
94
+ pong(data?: any, mask?: boolean, cb?: (err: Error) => void): void;
95
+ /** @experimental */
96
+ terminate(): void;
10
97
  }
98
+
11
99
  interface ActorConfig {
12
100
  name: string;
13
101
  key: string | null;
14
102
  createTs: bigint;
15
103
  input: Uint8Array | null;
16
104
  }
105
+ declare class RunnerActor {
106
+ /**
107
+ * List of hibernating requests provided by the gateway on actor start.
108
+ * This represents the WebSocket connections that the gateway knows about.
109
+ **/
110
+ hibernatingRequests: readonly protocol.HibernatingRequest[];
111
+ actorId: string;
112
+ generation: number;
113
+ config: ActorConfig;
114
+ pendingRequests: Array<{
115
+ gatewayId: protocol.GatewayId;
116
+ requestId: protocol.RequestId;
117
+ request: PendingRequest;
118
+ }>;
119
+ webSockets: Array<{
120
+ gatewayId: protocol.GatewayId;
121
+ requestId: protocol.RequestId;
122
+ ws: WebSocketTunnelAdapter;
123
+ }>;
124
+ actorStartPromise: ReturnType<typeof promiseWithResolvers<void>>;
125
+ /**
126
+ * If restoreHibernatingRequests has been called. This is used to assert
127
+ * that the caller is implemented correctly.
128
+ **/
129
+ hibernationRestored: boolean;
130
+ constructor(actorId: string, generation: number, config: ActorConfig,
131
+ /**
132
+ * List of hibernating requests provided by the gateway on actor start.
133
+ * This represents the WebSocket connections that the gateway knows about.
134
+ **/
135
+ hibernatingRequests: readonly protocol.HibernatingRequest[]);
136
+ getPendingRequest(gatewayId: protocol.GatewayId, requestId: protocol.RequestId): PendingRequest | undefined;
137
+ createPendingRequest(gatewayId: protocol.GatewayId, requestId: protocol.RequestId, clientMessageIndex: number): void;
138
+ createPendingRequestWithStreamController(gatewayId: protocol.GatewayId, requestId: protocol.RequestId, clientMessageIndex: number, streamController: ReadableStreamDefaultController<Uint8Array>): void;
139
+ deletePendingRequest(gatewayId: protocol.GatewayId, requestId: protocol.RequestId): void;
140
+ getWebSocket(gatewayId: protocol.GatewayId, requestId: protocol.RequestId): WebSocketTunnelAdapter | undefined;
141
+ setWebSocket(gatewayId: protocol.GatewayId, requestId: protocol.RequestId, ws: WebSocketTunnelAdapter): void;
142
+ deleteWebSocket(gatewayId: protocol.GatewayId, requestId: protocol.RequestId): void;
143
+ }
144
+
17
145
  interface RunnerConfig {
18
146
  logger?: Logger;
19
147
  version: number;
@@ -32,17 +160,105 @@ interface RunnerConfig {
32
160
  onConnected: () => void;
33
161
  onDisconnected: (code: number, reason: string) => void;
34
162
  onShutdown: () => void;
35
- fetch: (runner: Runner, actorId: string, requestId: protocol.RequestId, request: Request) => Promise<Response>;
36
- websocket?: (runner: Runner, actorId: string, ws: any, requestId: protocol.RequestId, request: Request) => Promise<void>;
163
+ /** Called when receiving a network request. */
164
+ fetch: (runner: Runner, actorId: string, gatewayId: protocol.GatewayId, requestId: protocol.RequestId, request: Request) => Promise<Response>;
165
+ /**
166
+ * Called when receiving a WebSocket connection.
167
+ *
168
+ * All event listeners must be added synchronously inside this function or
169
+ * else events may be missed. The open event will fire immediately after
170
+ * this function finishes.
171
+ *
172
+ * Any errors thrown here will disconnect the WebSocket immediately.
173
+ *
174
+ * While `path` and `headers` are partially redundant to the data in the
175
+ * `Request`, they may vary slightly from the actual content of `Request`.
176
+ * Prefer to persist the `path` and `headers` properties instead of the
177
+ * `Request` itself.
178
+ *
179
+ * ## Hibernating Web Sockets
180
+ *
181
+ * ### Implementation Requirements
182
+ *
183
+ * **Requirement 1: Persist HWS Immediately**
184
+ *
185
+ * This is responsible for persisting hibernatable WebSockets immediately
186
+ * (do not wait for open event). It is not time sensitive to flush the
187
+ * connection state. If this fails to persist the HWS, the client's
188
+ * WebSocket will be disconnected on next wake in the call to
189
+ * `Tunnel::restoreHibernatingRequests` since the connection entry will not
190
+ * exist.
191
+ *
192
+ * **Requirement 2: Persist Message Index On `message`**
193
+ *
194
+ * In the `message` event listener, this handler must persist the message
195
+ * index from the event. The request ID is available at
196
+ * `event.rivetRequestId` and message index at `event.rivetMessageIndex`.
197
+ *
198
+ * The message index should not be flushed immediately. Instead, this
199
+ * should:
200
+ *
201
+ * - Debounce calls to persist the message index
202
+ * - After each persist, call
203
+ * `Runner::sendHibernatableWebSocketMessageAck` to acknowledge the
204
+ * message
205
+ *
206
+ * This mechanism allows us to buffer messages on the gateway so we can
207
+ * batch-persist events on our end on a given interval.
208
+ *
209
+ * If this fails to persist, then the gateway will replay unacked
210
+ * messages when the actor starts again.
211
+ *
212
+ * **Requirement 3: Remove HWS From Storage On `close`**
213
+ *
214
+ * This handler should add an event listener for `close` to remove the
215
+ * connection from storage.
216
+ *
217
+ * If the connection remove fails to persist, the close event will be
218
+ * called again on the next actor start in
219
+ * `Tunnel::restoreHibernatingRequests` since there will be no request for
220
+ * the given connection.
221
+ *
222
+ * ### Restoring Connections
223
+ *
224
+ * The user of this library is responsible for:
225
+ * 1. Loading all persisted hibernatable WebSocket metadata for an actor
226
+ * 2. Calling `Runner::restoreHibernatingRequests` with this metadata at
227
+ * the end of `onActorStart`
228
+ *
229
+ * `restoreHibernatingRequests` will restore all connections and attach
230
+ * the appropriate event listeners.
231
+ *
232
+ * ### No Open Event On Restoration
233
+ *
234
+ * When restoring a HWS, the open event will not be called again. It will
235
+ * go straight to the message or close event.
236
+ */
237
+ websocket: (runner: Runner, actorId: string, ws: any, gatewayId: protocol.GatewayId, requestId: protocol.RequestId, request: Request, path: string, headers: Record<string, string>, isHibernatable: boolean, isRestoringHibernatable: boolean) => Promise<void>;
238
+ hibernatableWebSocket: {
239
+ /**
240
+ * Determines if a WebSocket can continue to live while an actor goes to
241
+ * sleep.
242
+ */
243
+ canHibernate: (actorId: string, gatewayId: ArrayBuffer, requestId: ArrayBuffer, request: Request) => boolean;
244
+ };
245
+ /**
246
+ * Called when an actor starts.
247
+ *
248
+ * This callback is responsible for:
249
+ * 1. Initializing the actor instance
250
+ * 2. Loading all persisted hibernatable WebSocket metadata for this actor
251
+ * 3. Calling `Runner::restoreHibernatingRequests` with the loaded metadata
252
+ * to restore hibernatable WebSocket connections
253
+ *
254
+ * The actor should not be marked as "ready" until after
255
+ * `restoreHibernatingRequests` completes to ensure all hibernatable
256
+ * connections are fully restored before the actor processes new requests.
257
+ */
37
258
  onActorStart: (actorId: string, generation: number, config: ActorConfig) => Promise<void>;
38
259
  onActorStop: (actorId: string, generation: number) => Promise<void>;
39
- getActorHibernationConfig: (actorId: string, requestId: ArrayBuffer, request: Request) => HibernationConfig;
40
260
  noAutoShutdown?: boolean;
41
261
  }
42
- interface HibernationConfig {
43
- enabled: boolean;
44
- lastMsgIndex: number | undefined;
45
- }
46
262
  interface KvListOptions {
47
263
  reverse?: boolean;
48
264
  limit?: number;
@@ -50,14 +266,17 @@ interface KvListOptions {
50
266
  declare class Runner {
51
267
  #private;
52
268
  get config(): RunnerConfig;
269
+ __pegboardWebSocket?: WebSocket;
53
270
  runnerId?: string;
54
271
  get log(): Logger | undefined;
55
272
  constructor(config: RunnerConfig);
56
273
  sleepActor(actorId: string, generation?: number): void;
57
274
  stopActor(actorId: string, generation?: number): Promise<void>;
58
275
  forceStopActor(actorId: string, generation?: number): Promise<void>;
59
- getActor(actorId: string, generation?: number): ActorInstance | undefined;
276
+ getActor(actorId: string, generation?: number): RunnerActor | undefined;
277
+ getAndWaitForActor(actorId: string, generation?: number): Promise<RunnerActor | undefined>;
60
278
  hasActor(actorId: string, generation?: number): boolean;
279
+ get actors(): Map<string, RunnerActor>;
61
280
  start(): Promise<void>;
62
281
  shutdown(immediate: boolean, exit?: boolean): Promise<void>;
63
282
  get pegboardEndpoint(): string;
@@ -71,10 +290,37 @@ declare class Runner {
71
290
  kvDrop(actorId: string): Promise<void>;
72
291
  setAlarm(actorId: string, alarmTs: number | null, generation?: number): void;
73
292
  clearAlarm(actorId: string, generation?: number): void;
74
- __webSocketReady(): boolean;
293
+ /** Asserts WebSocket exists and is ready. */
294
+ __webSocketReady(): this is this & {
295
+ __pegboardWebSocket: NonNullable<Runner["__pegboardWebSocket"]>;
296
+ };
75
297
  __sendToServer(message: protocol.ToServer): void;
76
- sendWebsocketMessageAck(requestId: ArrayBuffer, index: number): void;
298
+ sendHibernatableWebSocketMessageAck(gatewayId: ArrayBuffer, requestId: ArrayBuffer, index: number): void;
299
+ /**
300
+ * Restores hibernatable WebSocket connections for an actor.
301
+ *
302
+ * This method should be called at the end of `onActorStart` after the
303
+ * actor instance is fully initialized.
304
+ *
305
+ * This method will:
306
+ * - Restore all provided hibernatable WebSocket connections
307
+ * - Attach event listeners to the restored WebSockets
308
+ * - Close any WebSocket connections that failed to restore
309
+ *
310
+ * The provided metadata list should include all hibernatable WebSockets
311
+ * that were persisted for this actor. The gateway will automatically
312
+ * close any connections that are not restored (i.e., not included in
313
+ * this list).
314
+ *
315
+ * **Important:** This method must be called after `onActorStart` completes
316
+ * and before marking the actor as "ready" to ensure all hibernatable
317
+ * connections are fully restored.
318
+ *
319
+ * @param actorId - The ID of the actor to restore connections for
320
+ * @param metaEntries - Array of hibernatable WebSocket metadata to restore
321
+ */
322
+ restoreHibernatingRequests(actorId: string, metaEntries: HibernatingWebSocketMetadata[]): Promise<void>;
77
323
  getServerlessInitPacket(): string | undefined;
78
324
  }
79
325
 
80
- export { type ActorConfig, type ActorInstance, type HibernationConfig, type KvListOptions, Runner, type RunnerConfig };
326
+ export { type ActorConfig, type HibernatingWebSocketMetadata, type KvListOptions, Runner, RunnerActor, type RunnerConfig, idToStr };