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

Commit cd9b2cd

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

File tree

18 files changed

+535
-511
lines changed

18 files changed

+535
-511
lines changed

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

Lines changed: 24 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,10 @@ import {
1111
HEADER_ENCODING,
1212
type ManagerDisplayInformation,
1313
type ManagerDriver,
14+
WS_PROTOCOL_ACTOR,
15+
WS_PROTOCOL_CONN_PARAMS,
16+
WS_PROTOCOL_ENCODING,
17+
WS_PROTOCOL_TARGET,
1418
} from "rivetkit/driver-helpers";
1519
import { ActorAlreadyExists, InternalError } from "rivetkit/errors";
1620
import { getCloudflareAmbientEnv } from "./handler";
@@ -81,16 +85,21 @@ export class CloudflareActorsManagerDriver implements ManagerDriver {
8185
const id = env.ACTOR_DO.idFromString(actorId);
8286
const stub = env.ACTOR_DO.get(id);
8387

88+
const protocols: string[] = [];
89+
protocols.push(`${WS_PROTOCOL_TARGET}actor`);
90+
protocols.push(`${WS_PROTOCOL_ACTOR}${actorId}`);
91+
protocols.push(`${WS_PROTOCOL_ENCODING}${encoding}`);
92+
if (params) {
93+
protocols.push(
94+
`${WS_PROTOCOL_CONN_PARAMS}${encodeURIComponent(JSON.stringify(params))}`,
95+
);
96+
}
97+
8498
const headers: Record<string, string> = {
8599
Upgrade: "websocket",
86100
Connection: "Upgrade",
87-
[HEADER_ENCODING]: encoding,
101+
"sec-websocket-protocol": protocols.join(", "),
88102
};
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";
94103

95104
// Use the path parameter to determine the URL
96105
const normalizedPath = path.startsWith("/") ? path : `/${path}`;
@@ -152,7 +161,6 @@ export class CloudflareActorsManagerDriver implements ManagerDriver {
152161
actorId: string,
153162
encoding: Encoding,
154163
params: unknown,
155-
authData: unknown,
156164
): Promise<Response> {
157165
logger().debug({
158166
msg: "forwarding websocket to durable object",
@@ -188,14 +196,17 @@ export class CloudflareActorsManagerDriver implements ManagerDriver {
188196
}
189197
}
190198

191-
// Add RivetKit headers
192-
actorRequest.headers.set(HEADER_ENCODING, encoding);
199+
// Build protocols for WebSocket connection
200+
const protocols: string[] = [];
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/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: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -12,15 +12,26 @@ export const HEADER_ENCODING = "X-RivetKit-Encoding";
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.
1313
export const HEADER_CONN_PARAMS = "X-RivetKit-Conn-Params";
1414

15-
// Internal header
16-
export const HEADER_AUTH_DATA = "X-RivetKit-Auth-Data";
17-
1815
export const HEADER_ACTOR_ID = "X-RivetKit-Actor";
1916

2017
export const HEADER_CONN_ID = "X-RivetKit-Conn";
2118

2219
export const HEADER_CONN_TOKEN = "X-RivetKit-Conn-Token";
2320

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+
export const WS_PROTOCOL_TARGET = "target.";
27+
export const WS_PROTOCOL_ACTOR = "actor.";
28+
export const WS_PROTOCOL_ENCODING = "encoding.";
29+
export const WS_PROTOCOL_CONN_PARAMS = "conn_params.";
30+
31+
// MARK: WebSocket Inline Test Protocol Prefixes
32+
export const WS_PROTOCOL_TRANSPORT = "transport.";
33+
export const WS_PROTOCOL_PATH = "path.";
34+
2435
/**
2536
* Headers that publics can send from public clients.
2637
*
@@ -35,4 +46,6 @@ export const ALLOWED_PUBLIC_HEADERS = [
3546
HEADER_ACTOR_ID,
3647
HEADER_CONN_ID,
3748
HEADER_CONN_TOKEN,
49+
HEADER_RIVET_TARGET,
50+
HEADER_RIVET_ACTOR,
3851
];

‎packages/rivetkit/src/driver-helpers/mod.ts‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@ export type { ActorInstance, AnyActorInstance } from "@/actor/instance";
33
export {
44
HEADER_ACTOR_ID,
55
HEADER_ACTOR_QUERY,
6-
HEADER_AUTH_DATA,
76
HEADER_CONN_ID,
87
HEADER_CONN_PARAMS,
98
HEADER_CONN_TOKEN,

0 commit comments

Comments
 (0)