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

Commit 9ad97a7

Browse files
NathanFlurryNathanFlurry
committed
chore(core): migrate gateway to use websocket protocols
1 parent ecc0c9d commit 9ad97a7

File tree

19 files changed

+565
-521
lines changed

19 files changed

+565
-521
lines changed

‎packages/cloudflare-workers/src/manager-driver.ts‎

Lines changed: 27 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -6,11 +6,13 @@ import {
66
type GetForIdInput,
77
type GetOrCreateWithKeyInput,
88
type GetWithKeyInput,
9-
HEADER_AUTH_DATA,
10-
HEADER_CONN_PARAMS,
11-
HEADER_ENCODING,
129
type ManagerDisplayInformation,
1310
type ManagerDriver,
11+
WS_PROTOCOL_ACTOR,
12+
WS_PROTOCOL_CONN_PARAMS,
13+
WS_PROTOCOL_ENCODING,
14+
WS_PROTOCOL_STANDARD,
15+
WS_PROTOCOL_TARGET,
1416
} from "rivetkit/driver-helpers";
1517
import { ActorAlreadyExists, InternalError } from "rivetkit/errors";
1618
import { getCloudflareAmbientEnv } from "./handler";
@@ -81,16 +83,22 @@ export class CloudflareActorsManagerDriver implements ManagerDriver {
8183
const id = env.ACTOR_DO.idFromString(actorId);
8284
const stub = env.ACTOR_DO.get(id);
8385

86+
const protocols: string[] = [];
87+
protocols.push(WS_PROTOCOL_STANDARD);
88+
protocols.push(`${WS_PROTOCOL_TARGET}actor`);
89+
protocols.push(`${WS_PROTOCOL_ACTOR}${actorId}`);
90+
protocols.push(`${WS_PROTOCOL_ENCODING}${encoding}`);
91+
if (params) {
92+
protocols.push(
93+
`${WS_PROTOCOL_CONN_PARAMS}${encodeURIComponent(JSON.stringify(params))}`,
94+
);
95+
}
96+
8497
const headers: Record<string, string> = {
8598
Upgrade: "websocket",
8699
Connection: "Upgrade",
87-
[HEADER_ENCODING]: encoding,
100+
"sec-websocket-protocol": protocols.join(", "),
88101
};
89-
if (params) {
90-
headers[HEADER_CONN_PARAMS] = JSON.stringify(params);
91-
}
92-
// HACK: See packages/drivers/cloudflare-workers/src/websocket.ts
93-
headers["sec-websocket-protocol"] = "rivetkit";
94102

95103
// Use the path parameter to determine the URL
96104
const normalizedPath = path.startsWith("/") ? path : `/${path}`;
@@ -152,7 +160,6 @@ export class CloudflareActorsManagerDriver implements ManagerDriver {
152160
actorId: string,
153161
encoding: Encoding,
154162
params: unknown,
155-
authData: unknown,
156163
): Promise<Response> {
157164
logger().debug({
158165
msg: "forwarding websocket to durable object",
@@ -188,14 +195,18 @@ export class CloudflareActorsManagerDriver implements ManagerDriver {
188195
}
189196
}
190197

191-
// Add RivetKit headers
192-
actorRequest.headers.set(HEADER_ENCODING, encoding);
198+
// Build protocols for WebSocket connection
199+
const protocols: string[] = [];
200+
protocols.push(WS_PROTOCOL_STANDARD);
201+
protocols.push(`${WS_PROTOCOL_TARGET}actor`);
202+
protocols.push(`${WS_PROTOCOL_ACTOR}${actorId}`);
203+
protocols.push(`${WS_PROTOCOL_ENCODING}${encoding}`);
193204
if (params) {
194-
actorRequest.headers.set(HEADER_CONN_PARAMS, JSON.stringify(params));
195-
}
196-
if (authData) {
197-
actorRequest.headers.set(HEADER_AUTH_DATA, JSON.stringify(authData));
205+
protocols.push(
206+
`${WS_PROTOCOL_CONN_PARAMS}${encodeURIComponent(JSON.stringify(params))}`,
207+
);
198208
}
209+
actorRequest.headers.set("sec-websocket-protocol", protocols.join(", "));
199210

200211
const id = c.env.ACTOR_DO.idFromString(actorId);
201212
const stub = c.env.ACTOR_DO.get(id);

‎packages/cloudflare-workers/src/websocket.ts‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44

55
import type { UpgradeWebSocket, WSEvents, WSReadyState } from "hono/ws";
66
import { defineWebSocketHelper, WSContext } from "hono/ws";
7+
import { WS_PROTOCOL_STANDARD } from "rivetkit/driver-helpers";
78

89
// Based on https://github.com/honojs/hono/issues/1153#issuecomment-1767321332
910
export const upgradeWebSocket: UpgradeWebSocket<
@@ -62,8 +63,11 @@ export const upgradeWebSocket: UpgradeWebSocket<
6263

6364
// Set Sec-WebSocket-Protocol if does not exist
6465
const protocols = c.req.header("Sec-WebSocket-Protocol");
65-
if (typeof protocols === "string" && protocols.includes("rivetkit")) {
66-
headers["Sec-WebSocket-Protocol"] = "rivetkit";
66+
if (
67+
typeof protocols === "string" &&
68+
protocols.includes(WS_PROTOCOL_STANDARD)
69+
) {
70+
headers["Sec-WebSocket-Protocol"] = WS_PROTOCOL_STANDARD;
6771
}
6872

6973
return new Response(null, {

‎packages/rivetkit/fixtures/driver-test-suite/raw-websocket.ts‎

Lines changed: 0 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -48,16 +48,6 @@ export const rawWebSocketActor = actor({
4848
messageCount: ctx.state.messageCount,
4949
}),
5050
);
51-
} else if (parsed.type === "getAuthData") {
52-
// Auth data is not directly available in raw WebSocket handler
53-
// Send a message indicating this limitation
54-
websocket.send(
55-
JSON.stringify({
56-
type: "authData",
57-
authData: null,
58-
message: "Auth data not available in raw WebSocket handler",
59-
}),
60-
);
6151
} else if (parsed.type === "getRequestInfo") {
6252
// Send back the request URL info
6353
websocket.send(

‎packages/rivetkit/src/actor/instance.ts‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -918,7 +918,6 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
918918
state: CS,
919919
driverId: ConnectionDriver,
920920
driverState: unknown,
921-
authData: unknown,
922921
): Promise<Conn<S, CP, CS, V, I, DB>> {
923922
this.#assertReady();
924923

@@ -935,7 +934,6 @@ export class ActorInstance<S, CP, CS, V, I, DB extends AnyDatabaseProvider> {
935934
connDriverState: driverState,
936935
params: params,
937936
state: state,
938-
authData: authData,
939937
lastSeen: Date.now(),
940938
subscriptions: [],
941939
};

‎packages/rivetkit/src/actor/persisted.ts‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@ export interface PersistedConn<CP, CS> {
1717
connDriverState: unknown;
1818
params: CP;
1919
state: CS;
20-
authData?: unknown;
2120
subscriptions: PersistedSubscription[];
2221
lastSeen: number;
2322
}

‎packages/rivetkit/src/actor/router-endpoints.ts‎

Lines changed: 0 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,6 @@ export interface ConnectWebSocketOpts {
5353
encoding: Encoding;
5454
actorId: string;
5555
params: unknown;
56-
authData: unknown;
5756
}
5857

5958
export interface ConnectWebSocketOutput {
@@ -67,7 +66,6 @@ export interface ConnectSseOpts {
6766
encoding: Encoding;
6867
params: unknown;
6968
actorId: string;
70-
authData: unknown;
7169
}
7270

7371
export interface ConnectSseOutput {
@@ -81,7 +79,6 @@ export interface ActionOpts {
8179
actionName: string;
8280
actionArgs: unknown[];
8381
actorId: string;
84-
authData: unknown;
8582
}
8683

8784
export interface ActionOutput {
@@ -99,14 +96,12 @@ export interface ConnsMessageOpts {
9996
export interface FetchOpts {
10097
request: Request;
10198
actorId: string;
102-
authData: unknown;
10399
}
104100

105101
export interface WebSocketOpts {
106102
request: Request;
107103
websocket: UniversalWebSocket;
108104
actorId: string;
109-
authData: unknown;
110105
}
111106

112107
/**
@@ -119,7 +114,6 @@ export async function handleWebSocketConnect(
119114
actorId: string,
120115
encoding: Encoding,
121116
parameters: unknown,
122-
authData: unknown,
123117
): Promise<UpgradeWebSocketArgs> {
124118
const exposeInternalError = req ? getRequestExposeInternalError(req) : false;
125119

@@ -189,7 +183,6 @@ export async function handleWebSocketConnect(
189183
connState,
190184
CONNECTION_DRIVER_WEBSOCKET,
191185
{ encoding } satisfies GenericWebSocketDriverState,
192-
authData,
193186
);
194187

195188
// Unblock other handlers
@@ -339,7 +332,6 @@ export async function handleSseConnect(
339332
_runConfig: RunConfig,
340333
actorDriver: ActorDriver,
341334
actorId: string,
342-
authData: unknown,
343335
) {
344336
c.header("Content-Encoding", "Identity");
345337

@@ -376,7 +368,6 @@ export async function handleSseConnect(
376368
connState,
377369
CONNECTION_DRIVER_SSE,
378370
{ encoding } satisfies GenericSseDriverState,
379-
authData,
380371
);
381372

382373
// Wait for close
@@ -459,7 +450,6 @@ export async function handleAction(
459450
actorDriver: ActorDriver,
460451
actionName: string,
461452
actorId: string,
462-
authData: unknown,
463453
) {
464454
const encoding = getRequestEncoding(c.req);
465455
const parameters = getRequestConnParams(c.req);
@@ -491,7 +481,6 @@ export async function handleAction(
491481
connState,
492482
CONNECTION_DRIVER_HTTP,
493483
{} satisfies GenericHttpDriverState,
494-
authData,
495484
);
496485

497486
// Call action
@@ -562,7 +551,6 @@ export async function handleRawWebSocketHandler(
562551
path: string,
563552
actorDriver: ActorDriver,
564553
actorId: string,
565-
authData: unknown,
566554
): Promise<UpgradeWebSocketArgs> {
567555
const actor = await actorDriver.loadActor(actorId);
568556

‎packages/rivetkit/src/actor/router.ts‎

Lines changed: 3 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@ import {
1717
handleWebSocketConnect,
1818
} from "@/actor/router-endpoints";
1919
import {
20-
HEADER_AUTH_DATA,
2120
HEADER_CONN_ID,
2221
HEADER_CONN_PARAMS,
2322
HEADER_CONN_TOKEN,
@@ -84,13 +83,11 @@ export function createActorRouter(
8483
return upgradeWebSocket(async (c) => {
8584
const encodingRaw = c.req.header(HEADER_ENCODING);
8685
const connParamsRaw = c.req.header(HEADER_CONN_PARAMS);
87-
const authDataRaw = c.req.header(HEADER_AUTH_DATA);
8886

8987
const encoding = EncodingSchema.parse(encodingRaw);
9088
const connParams = connParamsRaw
9189
? JSON.parse(connParamsRaw)
9290
: undefined;
93-
const authData = authDataRaw ? JSON.parse(authDataRaw) : undefined;
9491

9592
return await handleWebSocketConnect(
9693
c.req.raw,
@@ -99,7 +96,6 @@ export function createActorRouter(
9996
c.env.actorId,
10097
encoding,
10198
connParams,
102-
authData,
10399
);
104100
})(c, noopNext());
105101
} else {
@@ -111,32 +107,13 @@ export function createActorRouter(
111107
});
112108

113109
router.get("/connect/sse", async (c) => {
114-
const authDataRaw = c.req.header(HEADER_AUTH_DATA);
115-
let authData: unknown;
116-
if (authDataRaw) {
117-
authData = JSON.parse(authDataRaw);
118-
}
119-
120-
return handleSseConnect(c, runConfig, actorDriver, c.env.actorId, authData);
110+
return handleSseConnect(c, runConfig, actorDriver, c.env.actorId);
121111
});
122112

123113
router.post("/action/:action", async (c) => {
124114
const actionName = c.req.param("action");
125115

126-
const authDataRaw = c.req.header(HEADER_AUTH_DATA);
127-
let authData: unknown;
128-
if (authDataRaw) {
129-
authData = JSON.parse(authDataRaw);
130-
}
131-
132-
return handleAction(
133-
c,
134-
runConfig,
135-
actorDriver,
136-
actionName,
137-
c.env.actorId,
138-
authData,
139-
);
116+
return handleAction(c, runConfig, actorDriver, actionName, c.env.actorId);
140117
});
141118

142119
router.post("/connections/message", async (c) => {
@@ -157,12 +134,6 @@ export function createActorRouter(
157134

158135
// Raw HTTP endpoints - /http/*
159136
router.all("/raw/http/*", async (c) => {
160-
const authDataRaw = c.req.header(HEADER_AUTH_DATA);
161-
let authData: unknown;
162-
if (authDataRaw) {
163-
authData = JSON.parse(authDataRaw);
164-
}
165-
166137
const actor = await actorDriver.loadActor(c.env.actorId);
167138

168139
// TODO: This is not a clean way of doing this since `/http/` might exist mid-path
@@ -186,9 +157,7 @@ export function createActorRouter(
186157
});
187158

188159
// Call the actor's onFetch handler - it will throw appropriate errors
189-
const response = await actor.handleFetch(correctedRequest, {
190-
auth: authData,
191-
});
160+
const response = await actor.handleFetch(correctedRequest, {});
192161

193162
// This should never happen now since handleFetch throws errors
194163
if (!response) {
@@ -205,13 +174,11 @@ export function createActorRouter(
205174
return upgradeWebSocket(async (c) => {
206175
const encodingRaw = c.req.header(HEADER_ENCODING);
207176
const connParamsRaw = c.req.header(HEADER_CONN_PARAMS);
208-
const authDataRaw = c.req.header(HEADER_AUTH_DATA);
209177

210178
const encoding = EncodingSchema.parse(encodingRaw);
211179
const connParams = connParamsRaw
212180
? JSON.parse(connParamsRaw)
213181
: undefined;
214-
const authData = authDataRaw ? JSON.parse(authDataRaw) : undefined;
215182

216183
const url = new URL(c.req.url);
217184
const pathWithQuery = c.req.path + url.search;
@@ -229,7 +196,6 @@ export function createActorRouter(
229196
pathWithQuery,
230197
actorDriver,
231198
c.env.actorId,
232-
authData,
233199
);
234200
})(c, noopNext());
235201
} else {

‎packages/rivetkit/src/common/actor-router-consts.ts‎

Lines changed: 23 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -5,21 +5,34 @@ export const PATH_CONNECT_WEBSOCKET = "/connect/websocket";
55
export const PATH_RAW_WEBSOCKET_PREFIX = "/raw/websocket/";
66

77
// MARK: Headers
8-
export const HEADER_ACTOR_QUERY = "X-RivetKit-Query";
8+
export const HEADER_ACTOR_QUERY = "x-rivet-query";
99

10-
export const HEADER_ENCODING = "X-RivetKit-Encoding";
10+
export const HEADER_ENCODING = "x-rivet-encoding";
1111

1212
// IMPORTANT: Params must be in headers or in an E2EE part of the request (i.e. NOT the URL or query string) in order to ensure that tokens can be securely passed in params.
13-
export const HEADER_CONN_PARAMS = "X-RivetKit-Conn-Params";
13+
export const HEADER_CONN_PARAMS = "x-rivet-conn-params";
1414

15-
// Internal header
16-
export const HEADER_AUTH_DATA = "X-RivetKit-Auth-Data";
15+
export const HEADER_ACTOR_ID = "x-rivet-actor";
1716

18-
export const HEADER_ACTOR_ID = "X-RivetKit-Actor";
17+
export const HEADER_CONN_ID = "x-rivet-conn";
1918

20-
export const HEADER_CONN_ID = "X-RivetKit-Conn";
19+
export const HEADER_CONN_TOKEN = "x-rivet-conn-token";
2120

22-
export const HEADER_CONN_TOKEN = "X-RivetKit-Conn-Token";
21+
// MARK: Manager Gateway Headers
22+
export const HEADER_RIVET_TARGET = "x-rivet-target";
23+
export const HEADER_RIVET_ACTOR = "x-rivet-actor";
24+
25+
// MARK: WebSocket Protocol Prefixes
26+
/** Some servers (such as node-ws & Cloudflare) require explicitly match a certain WebSocket protocol. This gives us a static protocol to match against. */
27+
export const WS_PROTOCOL_STANDARD = "rivet";
28+
export const WS_PROTOCOL_TARGET = "rivet_target.";
29+
export const WS_PROTOCOL_ACTOR = "rivet_actor.";
30+
export const WS_PROTOCOL_ENCODING = "rivet_encoding.";
31+
export const WS_PROTOCOL_CONN_PARAMS = "rivet_conn_params.";
32+
33+
// MARK: WebSocket Inline Test Protocol Prefixes
34+
export const WS_PROTOCOL_TRANSPORT = "test_transport.";
35+
export const WS_PROTOCOL_PATH = "test_path.";
2336

2437
/**
2538
* Headers that publics can send from public clients.
@@ -35,4 +48,6 @@ export const ALLOWED_PUBLIC_HEADERS = [
3548
HEADER_ACTOR_ID,
3649
HEADER_CONN_ID,
3750
HEADER_CONN_TOKEN,
51+
HEADER_RIVET_TARGET,
52+
HEADER_RIVET_ACTOR,
3853
];

0 commit comments

Comments
 (0)