rivetkit 2.3.20 → 2.3.22

package/src/engine-client/mod.ts

@@ -2,6 +2,12 @@ import type { Context as HonoContext } from "hono";
 import invariant from "invariant";
 import { deserializeActorKey, serializeActorKey } from "@/actor/keys";
 import type { ClientConfig } from "@/client/client";
+import type { IssuedToken, IssueTokenOptions } from "@/client/auth";
+import {
+	isInvalidToken,
+	isInvalidTokenResponse,
+	TokenProvider,
+} from "@/client/token-provider";
 import {
 	PATH_CONNECT,
 	PATH_WEBSOCKET_BASE,
@@ -40,6 +46,7 @@ import {
 	getActor,
 	getActorByKey,
 	getOrCreateActor,
+	issueToken as issueEngineToken,
 	listActorsByName,
 } from "./api-endpoints";
 import { EngineApiError, getEndpoint } from "./api-utils";
@@ -50,6 +57,13 @@ import { createWebSocketProxy } from "./ws-proxy";
 export class RemoteEngineControlClient implements EngineControlClient {
 	#config: ClientConfig;
 	#metadataPromise: Promise<void> | undefined;
+	#tokenProvider?: TokenProvider;
+	#webSocketTokens = new WeakMap<
+		UniversalWebSocket,
+		{
+			refresh?: Promise<string>;
+		}
+	>();
 
 	constructor(runConfig: ClientConfig) {
 		// Disable health check if in Next.js build phase since there is no `/metadata` endpoint
@@ -65,6 +79,8 @@ export class RemoteEngineControlClient implements EngineControlClient {
 		// Clone config so we can mutate the endpoint in #metadataPromise
 		// NOTE: This is a shallow clone, so mutating nested properties will not do anything
 		this.#config = { ...runConfig };
+		if (runConfig.getToken)
+			this.#tokenProvider = new TokenProvider(runConfig.getToken);
 
 		// Perform metadata check if enabled
 		if (!runConfig.disableMetadataLookup) {
@@ -80,7 +96,7 @@ export class RemoteEngineControlClient implements EngineControlClient {
 							this.#config.namespace =
 								metadataData.clientNamespace;
 						}
-						if (metadataData.clientToken) {
+						if (metadataData.clientToken && !this.#tokenProvider) {
 							this.#config.token = metadataData.clientToken;
 						}
 
@@ -88,7 +104,6 @@ export class RemoteEngineControlClient implements EngineControlClient {
 							msg: "overriding client endpoint",
 							endpoint: metadataData.clientEndpoint,
 							namespace: metadataData.clientNamespace,
-							token: metadataData.clientToken,
 						});
 					}
 
@@ -103,6 +118,51 @@ export class RemoteEngineControlClient implements EngineControlClient {
 		}
 	}
 
+	async #authorizedConfig(): Promise<ClientConfig> {
+		if (!this.#tokenProvider) return this.#config;
+		return { ...this.#config, token: await this.#tokenProvider.current() };
+	}
+
+	async #withCredential<T>(
+		request: (config: ClientConfig) => Promise<T>,
+		readOnly = false,
+	): Promise<T> {
+		const config = await this.#authorizedConfig();
+		try {
+			return await request(config);
+		} catch (error) {
+			if (
+				this.#tokenProvider &&
+				config.token &&
+				error instanceof EngineApiError &&
+				error.statusCode === 401 &&
+				isInvalidToken(error.group, error.code)
+			) {
+				if (readOnly) {
+					const token = await this.#tokenProvider.refreshIfCurrent(
+						config.token,
+					);
+					return request({ ...this.#config, token });
+				}
+				// Mutations may have executed; refresh the next call, not this one.
+				await this.#tokenProvider.refreshIfCurrent(config.token);
+			}
+			throw error;
+		}
+	}
+
+	async issueToken(options: IssueTokenOptions): Promise<IssuedToken> {
+		await this.#metadataPromise;
+		const response = await this.#withCredential((config) =>
+			issueEngineToken(config, options),
+		);
+		return {
+			token: response.token,
+			issuedAt: response.issued_ts,
+			expiresAt: response.expires_ts,
+		};
+	}
+
 	async getForId({
 		name,
 		actorId,
@@ -110,7 +170,10 @@ export class RemoteEngineControlClient implements EngineControlClient {
 		await this.#metadataPromise;
 
 		// Fetch from API if not in cache
-		const response = await getActor(this.#config, name, actorId);
+		const response = await this.#withCredential(
+			(config) => getActor(config, name, actorId),
+			true,
+		);
 		const actor = response.actors[0];
 		if (!actor) return undefined;
 
@@ -138,7 +201,10 @@ export class RemoteEngineControlClient implements EngineControlClient {
 
 		// If not in local cache, fetch by key from API
 		try {
-			const response = await getActorByKey(this.#config, name, key);
+			const response = await this.#withCredential(
+				(config) => getActorByKey(config, name, key),
+				true,
+			);
 			const actor = response.actors[0];
 			if (!actor) return undefined;
 
@@ -183,18 +249,20 @@ export class RemoteEngineControlClient implements EngineControlClient {
 		});
 
 		try {
-			const { actor, created } = await getOrCreateActor(this.#config, {
-				datacenter: region,
-				name,
-				key: serializeActorKey(key),
-				runner_name_selector: poolName ?? this.#config.poolName,
-				input: actorInput
-					? uint8ArrayToBase64(
-							encodeCborCompat(actorInput as JsonCompatValue),
-						)
-					: undefined,
-				crash_policy: crashPolicy ?? "sleep",
-			});
+			const { actor, created } = await this.#withCredential((config) =>
+				getOrCreateActor(config, {
+					datacenter: region,
+					name,
+					key: serializeActorKey(key),
+					runner_name_selector: poolName ?? this.#config.poolName,
+					input: actorInput
+						? uint8ArrayToBase64(
+								encodeCborCompat(actorInput as JsonCompatValue),
+							)
+						: undefined,
+					crash_policy: crashPolicy ?? "sleep",
+				}),
+			);
 
 			logger().info({
 				msg: "getOrCreateWithKey: actor ready",
@@ -220,7 +288,10 @@ export class RemoteEngineControlClient implements EngineControlClient {
 					key,
 				});
 
-				const response = await getActorByKey(this.#config, name, key);
+				const response = await this.#withCredential(
+					(config) => getActorByKey(config, name, key),
+					true,
+				);
 				const existing = response.actors[0];
 				if (!existing) throw error;
 
@@ -250,16 +321,20 @@ export class RemoteEngineControlClient implements EngineControlClient {
 		logger().info({ msg: "creating actor via engine api", name, key });
 
 		// Create actor via engine API
-		const result = await createActor(this.#config, {
-			datacenter: region,
-			name,
-			runner_name_selector: poolName ?? this.#config.poolName,
-			key: serializeActorKey(key),
-			input: input
-				? uint8ArrayToBase64(encodeCborCompat(input as JsonCompatValue))
-				: undefined,
-			crash_policy: crashPolicy ?? "sleep",
-		});
+		const result = await this.#withCredential((config) =>
+			createActor(config, {
+				datacenter: region,
+				name,
+				runner_name_selector: poolName ?? this.#config.poolName,
+				key: serializeActorKey(key),
+				input: input
+					? uint8ArrayToBase64(
+							encodeCborCompat(input as JsonCompatValue),
+						)
+					: undefined,
+				crash_policy: crashPolicy ?? "sleep",
+			}),
+		);
 
 		logger().info({
 			msg: "actor created",
@@ -276,7 +351,10 @@ export class RemoteEngineControlClient implements EngineControlClient {
 
 		logger().debug({ msg: "listing actors via engine api", name });
 
-		const response = await listActorsByName(this.#config, name);
+		const response = await this.#withCredential(
+			(config) => listActorsByName(config, name),
+			true,
+		);
 
 		return response.actors.map(apiActorToOutput);
 	}
@@ -286,7 +364,7 @@ export class RemoteEngineControlClient implements EngineControlClient {
 
 		logger().info({ msg: "destroying actor via engine api", actorId });
 
-		await destroyActor(this.#config, actorId);
+		await this.#withCredential((config) => destroyActor(config, actorId));
 
 		logger().info({ msg: "actor destroyed", actorId });
 	}
@@ -297,13 +375,7 @@ export class RemoteEngineControlClient implements EngineControlClient {
 		options: GatewayRequestOptions = {},
 	): Promise<Response> {
 		await this.#metadataPromise;
-
-		const path = requestPath(actorRequest);
-		const gatewayUrl = this.#buildGatewayUrlForTarget(
-			target,
-			path,
-			options,
-		);
+		const config = await this.#authorizedConfig();
 		const httpOptions = {
 			...options,
 			directActorId: shouldSkipReadyWait(options)
@@ -311,12 +383,37 @@ export class RemoteEngineControlClient implements EngineControlClient {
 				: undefined,
 		};
 
-		return sendHttpRequestToGateway(
-			this.#config,
-			gatewayUrl,
-			actorRequest,
-			httpOptions,
-		);
+		const send = (current: ClientConfig) =>
+			sendHttpRequestToGateway(
+				current,
+				this.#buildGatewayUrlForTarget(
+					current,
+					target,
+					requestPath(actorRequest),
+					options,
+				),
+				actorRequest,
+				httpOptions,
+			);
+		const response = await send(config);
+		if (
+			!this.#tokenProvider ||
+			!config.token ||
+			!isInvalidTokenResponse(response)
+		)
+			return response;
+		if (
+			(actorRequest.method === "GET" || actorRequest.method === "HEAD") &&
+			!("getOrCreateForKey" in target)
+		) {
+			const token = await this.#tokenProvider.refreshIfCurrent(
+				config.token,
+			);
+			await response.body?.cancel();
+			return send({ ...this.#config, token });
+		}
+		await this.#tokenProvider.refreshIfCurrent(config.token);
+		return response;
 	}
 
 	async openWebSocket(
@@ -327,15 +424,17 @@ export class RemoteEngineControlClient implements EngineControlClient {
 		options: GatewayRequestOptions = {},
 	): Promise<UniversalWebSocket> {
 		await this.#metadataPromise;
+		const config = await this.#authorizedConfig();
 
 		const gatewayUrl = this.#buildGatewayUrlForTarget(
+			config,
 			target,
 			path,
 			options,
 		);
 
-		return openWebSocketToGateway(
-			this.#config,
+		const ws = await openWebSocketToGateway(
+			config,
 			gatewayUrl,
 			encoding,
 			params,
@@ -346,6 +445,45 @@ export class RemoteEngineControlClient implements EngineControlClient {
 					: undefined,
 			},
 		);
+		this.#watchWebSocketAuth(ws, config.token);
+		return ws;
+	}
+
+	async refreshAuthToken(ws: UniversalWebSocket): Promise<boolean> {
+		const context = this.#webSocketTokens.get(ws);
+		if (!context?.refresh) return false;
+		await context.refresh;
+		return true;
+	}
+
+	#watchWebSocketAuth(ws: UniversalWebSocket, token?: string): void {
+		const tokenProvider = this.#tokenProvider;
+		if (!tokenProvider || !token) return;
+		const context: { refresh?: Promise<string> } = {};
+		this.#webSocketTokens.set(ws, context);
+		ws.addEventListener(
+			"close",
+			(event: { code?: number; reason?: string }) => {
+				if (event.code !== 1008 || !event.reason) return;
+				const separator = event.reason.indexOf(".");
+				if (separator < 0) return;
+				const code = event.reason
+					.slice(separator + 1)
+					.split(/[\s:#]/, 1)[0];
+				if (isInvalidToken(event.reason.slice(0, separator), code)) {
+					context.refresh = tokenProvider
+						.refreshIfCurrent(token)
+						.catch((error) => {
+							logger().warn(
+								"failed to renew authentication token",
+							);
+							throw error;
+						});
+					// Raw websockets have no built-in reconnection; preempt unhandled rejection.
+					context.refresh.catch(() => {});
+				}
+			},
+		);
 	}
 
 	async buildGatewayUrl(
@@ -353,7 +491,12 @@ export class RemoteEngineControlClient implements EngineControlClient {
 		options: GatewayRequestOptions = {},
 	): Promise<string> {
 		await this.#metadataPromise;
-		return this.#buildGatewayUrlForTarget(target, "", options);
+		return this.#buildGatewayUrlForTarget(
+			await this.#authorizedConfig(),
+			target,
+			"",
+			options,
+		);
 	}
 
 	async proxyRequest(
@@ -361,14 +504,7 @@ export class RemoteEngineControlClient implements EngineControlClient {
 		actorRequest: Request,
 		actorId: string,
 	): Promise<Response> {
-		await this.#metadataPromise;
-
-		const gatewayUrl = this.#buildGatewayUrlForTarget(
-			{ directId: actorId },
-			requestPath(actorRequest),
-		);
-
-		return sendHttpRequestToGateway(this.#config, gatewayUrl, actorRequest);
+		return this.sendRequest({ directId: actorId }, actorRequest);
 	}
 
 	async proxyWebSocket(
@@ -379,11 +515,12 @@ export class RemoteEngineControlClient implements EngineControlClient {
 		params: unknown,
 	): Promise<Response> {
 		await this.#metadataPromise;
+		const config = await this.#authorizedConfig();
 
-		const upgradeWebSocket = this.#config.getUpgradeWebSocket?.();
+		const upgradeWebSocket = config.getUpgradeWebSocket?.();
 		invariant(upgradeWebSocket, "missing getUpgradeWebSocket");
 
-		const endpoint = getEndpoint(this.#config);
+		const endpoint = getEndpoint(config);
 		const guardUrl = combineUrlPath(endpoint, path);
 		const wsGuardUrl = guardUrl.replace("http://", "ws://");
 
@@ -396,7 +533,7 @@ export class RemoteEngineControlClient implements EngineControlClient {
 
 		// Build protocols
 		const protocols = buildWebSocketProtocols(
-			this.#config,
+			config,
 			encoding,
 			params,
 			undefined,
@@ -419,11 +556,12 @@ export class RemoteEngineControlClient implements EngineControlClient {
 	}
 
 	#buildGatewayUrlForTarget(
+		config: ClientConfig,
 		target: GatewayTarget,
 		path: string,
 		options: GatewayRequestOptions = {},
 	): string {
-		const endpoint = getEndpoint(this.#config);
+		const endpoint = getEndpoint(config);
 
 		if (
 			shouldSkipReadyWait(options) &&
@@ -437,7 +575,7 @@ export class RemoteEngineControlClient implements EngineControlClient {
 			return buildActorGatewayUrl(
 				endpoint,
 				target.directId,
-				this.#config.token,
+				config.token,
 				path,
 			);
 		}
@@ -446,7 +584,7 @@ export class RemoteEngineControlClient implements EngineControlClient {
 			return buildActorGatewayUrl(
 				endpoint,
 				target.getForId.actorId,
-				this.#config.token,
+				config.token,
 				path,
 			);
 		}
@@ -454,15 +592,14 @@ export class RemoteEngineControlClient implements EngineControlClient {
 		if ("getForKey" in target || "getOrCreateForKey" in target) {
 			return buildActorQueryGatewayUrl(
 				endpoint,
-				this.#config.namespace,
+				config.namespace,
 				target,
-				this.#config.token,
+				config.token,
 				path,
-				this.#config.maxInputSize,
+				config.maxInputSize,
 				undefined,
 				"getOrCreateForKey" in target
-					? (target.getOrCreateForKey.poolName ??
-							this.#config.poolName)
+					? (target.getOrCreateForKey.poolName ?? config.poolName)
 					: undefined,
 				options,
 			);

package/src/registry/config/envoy.ts

@@ -1,29 +1,13 @@
 import { z } from "zod/v4";
-import { getLogger } from "@/common/log";
 import {
-	getNodeEnv,
 	getRivetEnvoyVersion,
 	getRivetPool,
 	getRivetTotalSlots,
 } from "@/utils/env-vars";
 
-let warnedMissingVersion = false;
-
 export const EnvoyConfigSchema = z.object({
 	poolName: z.string().default(() => getRivetPool() ?? "default"),
-	version: z.number().default(() => {
-		const version = getRivetEnvoyVersion();
-		if (version !== undefined) return version;
-
-		if (getNodeEnv() === "production" && !warnedMissingVersion) {
-			warnedMissingVersion = true;
-			getLogger("rivetkit").error(
-				"RIVET_ENVOY_VERSION is not set. Actors will not be versioned, which means they won't be drained on deploy. This is only needed when self-hosting or using a custom envoy (not needed for Rivet Compute). Set this as a build arg in your Dockerfile. See https://rivet.dev/docs/actors/versions",
-			);
-		}
-
-		return 1;
-	}),
+	version: z.number().default(() => getRivetEnvoyVersion() ?? 1),
 
 	// Deprecated.
 	totalSlots: z.number().default(() => getRivetTotalSlots() ?? 100000),

package/src/registry/config/index.test.ts

@@ -12,7 +12,7 @@ afterEach(() => {
 });
 
 describe("RegistryConfigSchema", () => {
-	test("defaults token to dev", () => {
+	test("defaults local Engine token", () => {
 		delete process.env.RIVET_TOKEN;
 
 		const config = RegistryConfigSchema.parse({
@@ -20,6 +20,6 @@ describe("RegistryConfigSchema", () => {
 			startEngine: true,
 		});
 
-		expect(config.token).toBe("dev");
+		expect(config.token).toBe("default");
 	});
 });

package/src/registry/config/index.ts

@@ -379,8 +379,13 @@ export const RegistryConfigSchema = z
 		// Namespace priority: parsed from endpoint URL > config value (includes env var) > "default"
 		const namespace =
 			parsedEndpoint?.namespace ?? config.namespace ?? "default";
-		// Token priority: parsed from endpoint URL > config value (includes env var) > "dev"
-		const token = parsedEndpoint?.token ?? config.token ?? "dev";
+		// Token priority: parsed from endpoint URL > config value (includes env var)
+		const token =
+			parsedEndpoint?.token ??
+			config.token ??
+			(endpoint && isLocalEngineEndpoint(endpoint)
+				? "default"
+				: undefined);
 
 		// Parse publicEndpoint string (env var fallback is applied via transform in serverless schema)
 		const parsedPublicEndpoint = config.serverless.publicEndpoint

package/src/registry/index.ts

@@ -4,6 +4,8 @@ import { configureBaseLogger } from "@/common/log";
 import { configureServerlessPool } from "@/serverless/configure";
 import { VERSION } from "@/utils";
 import {
+	getNodeEnv,
+	getRivetEnvoyVersion,
 	getRivetkitPublicDir,
 	getRivetkitRuntimeMode,
 	parsePortEnv,
@@ -624,6 +626,32 @@ export class Registry<A extends RegistryActors> {
 	 */
 	#startEnvoy(config: RegistryConfig, printWelcome: boolean) {
 		if (!this.#runtimeServePromise) {
+			const configuredEnvoyVersion = this.#config.envoy?.version;
+			const envoyVersionSource =
+				configuredEnvoyVersion !== undefined
+					? "config"
+					: getRivetEnvoyVersion() !== undefined
+						? "environment"
+						: "default";
+
+			logger().info({
+				msg: "starting rivetkit envoy",
+				rivetkitVersion: VERSION,
+				envoyVersion: config.envoy.version,
+				envoyVersionSource,
+			});
+
+			if (
+				envoyVersionSource === "default" &&
+				getNodeEnv() === "production"
+			) {
+				logger().error({
+					msg: "envoy version defaulted because neither config.envoy.version nor RIVET_ENVOY_VERSION is set; actors will not drain automatically across deploys. See https://rivet.dev/docs/actors/versions",
+					envoyVersion: config.envoy.version,
+					envoyVersionSource,
+				});
+			}
+
 			const configuredRegistryPromise =
 				this.#buildConfiguredRegistry(config);
 			this.#runtimeServeConfiguredPromise = configuredRegistryPromise;

package/src/registry/native.ts

@@ -872,7 +872,19 @@ function encodeNativeCallbackError(error: unknown): Error {
 		? error
 		: deconstructError(error, true);
 
-	const bridgeError = new Error(encodeBridgeRivetError(structuredError), {
+	const { encoded, droppedMetadataCause } =
+		encodeBridgeRivetError(structuredError);
+	if (droppedMetadataCause !== undefined) {
+		logger().warn(
+			{
+				group: structuredError.group,
+				code: structuredError.code,
+				error: String(droppedMetadataCause),
+			},
+			"dropped error metadata that JSON cannot encode",
+		);
+	}
+	const bridgeError = new Error(encoded, {
 		cause: error instanceof Error ? error : undefined,
 	});
 	return Object.assign(bridgeError, {
@@ -2771,6 +2783,11 @@ export class ActorContextHandleAdapter {
 			}
 		)[ACTOR_CONTEXT_INTERNAL_SYMBOL] = new NativeWorkflowRuntimeAdapter(
 			this,
+			() =>
+				this.#logByInvocationScope.delete(
+					this.#runtime.currentInvocationScope() ??
+						NO_INVOCATION_SCOPE,
+				),
 		);
 	}
 
@@ -2978,7 +2995,10 @@ export class ActorContextHandleAdapter {
 		);
 	}
 
-	/** Cached per invocation, because the `run` context outlives workflow runs and steps. */
+	/**
+	 * Cached per invocation, because the `run` context outlives workflow runs and steps.
+	 * A workflow queue receive clears the entry, because the run's ray becomes the message's ray.
+	 */
 	get log() {
 		const scope =
 			this.#runtime.currentInvocationScope() ?? NO_INVOCATION_SCOPE;
@@ -3635,6 +3655,7 @@ type NativeWorkflowQueueMessage = Awaited<
 
 class NativeWorkflowRuntimeAdapter {
 	#ctx: ActorContextHandleAdapter;
+	#onMessagesReceived: () => void;
 	#completions = new Map<string, (response?: unknown) => Promise<void>>();
 
 	readonly id: string;
@@ -3702,8 +3723,12 @@ class NativeWorkflowRuntimeAdapter {
 		) => Promise<void>;
 	};
 
-	constructor(ctx: ActorContextHandleAdapter) {
+	constructor(
+		ctx: ActorContextHandleAdapter,
+		onMessagesReceived: () => void,
+	) {
 		this.#ctx = ctx;
+		this.#onMessagesReceived = onMessagesReceived;
 		this.id = ctx.actorId;
 		this.driver = {
 			kvBatchGet: async (actorId, keys) => {
@@ -3749,6 +3774,9 @@ class NativeWorkflowRuntimeAdapter {
 					timeout: timeout ?? 0,
 					completable,
 				});
+				if (messages.length > 0) {
+					this.#onMessagesReceived();
+				}
 				return messages.map((message) =>
 					this.#wrapQueueMessage(message),
 				);