rivetkit 2.3.20 → 2.3.22

package/src/client/auth.ts

@@ -0,0 +1,37 @@
+import type { Rivet } from "@rivetkit/engine-api-full";
+
+/** A namespace-scoped Engine grant. The caller must provide every grant explicitly. */
+export type TokenGrant = Rivet.AuthTokenGrant;
+
+/** Options for `client.auth.issueToken()`. `expiresIn` is in seconds. */
+export type IssueTokenOptions = Pick<
+	Rivet.AuthTokenCreateRequest,
+	"subject" | "grants"
+> & {
+	expiresIn?: Rivet.AuthTokenCreateRequest["duration"];
+};
+
+/** Token timestamps are Unix milliseconds. */
+export interface IssuedToken {
+	token: Rivet.AuthTokenCreateResponse["token"];
+	issuedAt: Rivet.AuthTokenCreateResponse["issuedTs"];
+	expiresAt: Rivet.AuthTokenCreateResponse["expiresTs"];
+}
+
+/** Resources that can be scoped to a single actor ID. */
+export type ActorTokenResource = Extract<
+	Rivet.AuthTokenGrant["resource"],
+	"actor_gateway" | "actor" | "actor_kv"
+>;
+
+export type ActorTokenPermissions = Partial<
+	Record<ActorTokenResource, Rivet.AuthTokenGrant["operations"]>
+>;
+
+/** Options for `actor.issueToken()`. Omit permissions for gateway read access only. */
+export type ActorIssueTokenOptions = Pick<
+	IssueTokenOptions,
+	"subject" | "expiresIn"
+> & {
+	permissions?: ActorTokenPermissions;
+};

package/src/client/client.ts

@@ -1,4 +1,5 @@
 import type { AnyActorDefinition } from "@/actor/definition";
+import type { IssuedToken, IssueTokenOptions } from "@/client/auth";
 import type { ActorQuery } from "@/client/query";
 import type { Encoding } from "@/common/encoding";
 import type { EngineControlClient } from "@/engine-client/driver";
@@ -198,6 +199,10 @@ export interface ActorClientRuntimeOptions {
  */
 export class ClientRaw {
 	#disposed = false;
+	/** Issue namespace-scoped tokens with an explicit grant list. */
+	readonly auth: {
+		issueToken(options: IssueTokenOptions): Promise<IssuedToken>;
+	};
 
 	[ACTOR_CONNS_SYMBOL] = new Set<ActorConnRaw>();
 
@@ -215,6 +220,9 @@ export class ClientRaw {
 		options: ClientRawOptions & ActorClientRuntimeOptions = {},
 	) {
 		this.#driver = driver;
+		this.auth = {
+			issueToken: (issueOptions) => this.#driver.issueToken(issueOptions),
+		};
 
 		this.#encodingKind = options.encoding ?? "bare";
 		this.#gatewayOptions = options.gateway ?? {};
@@ -469,11 +477,11 @@ export class ClientRaw {
  *
  * @template A The actor registry type.
  */
-export type Client<A extends Registry<any>> = ClientRaw & {
+export type Client<A extends Registry<any>> = ClientRaw & Omit<{
 	[K in keyof ExtractActorsFromRegistry<A>]: ActorAccessor<
 		ExtractActorsFromRegistry<A>[K]
 	>;
-};
+}, keyof ClientRaw>;
 
 export type AnyClient = Client<Registry<any>>;
 

package/src/client/config.ts

@@ -1,5 +1,6 @@
 import z from "zod/v4";
 import { EncodingSchema } from "@/common/encoding";
+import { isLocalEngineEndpoint } from "@/common/engine";
 import type { RegistryConfig } from "@/registry/config";
 import type { GetUpgradeWebSocket } from "@/utils";
 import { tryParseEndpoint } from "@/utils/endpoint-parser";
@@ -10,6 +11,7 @@ import {
 	getRivetPool,
 	getRivetToken,
 } from "@/utils/env-vars";
+import type { GetToken } from "./token-provider";
 
 const DEFAULT_ENDPOINT = "http://localhost:6420";
 export const DEFAULT_MAX_QUERY_INPUT_SIZE = 4 * 1024;
@@ -54,6 +56,9 @@ export const ClientConfigSchemaBase = z.object({
 		.optional()
 		.transform((val) => val ?? getRivetToken()),
 
+	/** Issue and renew a client token on demand. Overrides the static token when provided. */
+	getToken: z.custom<GetToken>().optional(),
+
 	/** Namespace to connect to. */
 	namespace: z
 		.string()
@@ -131,7 +136,12 @@ export function transformClientConfig(
 		...config,
 		endpoint: parsedEndpoint?.endpoint,
 		namespace: parsedEndpoint?.namespace ?? config.namespace ?? "default",
-		token: parsedEndpoint?.token ?? config.token,
+		token:
+			parsedEndpoint?.token ??
+			config.token ??
+			(config.endpoint && isLocalEngineEndpoint(config.endpoint)
+				? "default"
+				: undefined),
 	};
 }
 
@@ -151,6 +161,7 @@ export function convertRegistryConfigToClientConfig(
 	return {
 		endpoint: config.endpoint,
 		token: config.token,
+		getToken: undefined,
 		namespace: config.namespace,
 		poolName: config.envoy.poolName,
 		headers: config.headers,

package/src/client/mod.ts

@@ -19,6 +19,14 @@ export {
 	UserError,
 } from "@/client/errors";
 export type { CreateRequest } from "@/client/query";
+export type {
+	ActorIssueTokenOptions,
+	ActorTokenPermissions,
+	ActorTokenResource,
+	IssuedToken,
+	IssueTokenOptions,
+	TokenGrant,
+} from "./auth";
 export type { Encoding } from "@/common/encoding";
 export type {
 	ActorActionFunction,
@@ -49,6 +57,7 @@ export type {
 	QueryOptions,
 	Region,
 } from "./client";
+export type { GetToken } from "./token-provider";
 
 /**
  * Creates a client with the actor accessor proxy.

package/src/client/token-provider.test.ts

@@ -0,0 +1,86 @@
+import { describe, expect, test, vi } from "vitest";
+import {
+	isInvalidToken,
+	isInvalidTokenResponse,
+	TokenProvider,
+} from "./token-provider";
+
+describe("renewable client tokens", () => {
+	test("single-flights initial issuance and concurrent rejection refreshes", async () => {
+		const getToken = vi.fn(
+			async ({ forceRefresh }: { forceRefresh: boolean }) =>
+				forceRefresh ? "renewed" : "initial",
+		);
+		const provider = new TokenProvider(getToken);
+		expect(
+			await Promise.all([provider.current(), provider.current()]),
+		).toEqual(["initial", "initial"]);
+		expect(getToken).toHaveBeenCalledTimes(1);
+		expect(
+			await Promise.all([
+				provider.refreshIfCurrent("initial"),
+				provider.refreshIfCurrent("initial"),
+			]),
+		).toEqual(["renewed", "renewed"]);
+		expect(getToken).toHaveBeenCalledTimes(2);
+		expect(await provider.refreshIfCurrent("initial")).toBe("renewed");
+		expect(getToken).toHaveBeenCalledTimes(2);
+	});
+
+	test("retries failed issuance and refuses empty credentials", async () => {
+		const getToken = vi
+			.fn()
+			.mockRejectedValueOnce(new Error("issuer down"))
+			.mockResolvedValueOnce("ok");
+		const provider = new TokenProvider(getToken);
+		await expect(provider.current()).rejects.toThrow("issuer down");
+		expect(await provider.current()).toBe("ok");
+		await expect(
+			new TokenProvider(async () => "").current(),
+		).rejects.toThrow("nonempty");
+	});
+
+	test("refreshes before a known JWT expiry and treats opaque tokens as unexpiring hints", async () => {
+		const payload = btoa(
+			JSON.stringify({ exp: Math.floor(Date.now() / 1_000) + 1 }),
+		).replaceAll("=", "");
+		const getToken = vi.fn(async () => `header.${payload}.signature`);
+		const provider = new TokenProvider(getToken);
+		await provider.current();
+		await provider.current();
+		expect(getToken).toHaveBeenCalledTimes(2);
+		const opaque = vi.fn(async () => "opaque");
+		const opaqueProvider = new TokenProvider(opaque);
+		await opaqueProvider.current();
+		await opaqueProvider.current();
+		expect(opaque).toHaveBeenCalledTimes(1);
+	});
+
+	test("only confirmed 401 invalid/expired errors permit renewal", () => {
+		expect(isInvalidToken("auth", "invalid_token")).toBe(true);
+		expect(isInvalidToken("auth", "token_expired")).toBe(true);
+		for (const code of [
+			"insufficient_permissions",
+			"verification_unavailable",
+			"issuance_unavailable",
+		]) {
+			expect(isInvalidToken("auth", code)).toBe(false);
+		}
+		expect(
+			isInvalidTokenResponse(
+				new Response(null, {
+					status: 401,
+					headers: { "x-rivet-error": "auth.token_expired" },
+				}),
+			),
+		).toBe(true);
+		expect(
+			isInvalidTokenResponse(
+				new Response(null, {
+					status: 503,
+					headers: { "x-rivet-error": "auth.token_expired" },
+				}),
+			),
+		).toBe(false);
+	});
+});

package/src/client/token-provider.ts

@@ -0,0 +1,93 @@
+/** Called when RivetKit needs a credential. Keep issuer secrets on your backend. */
+export type GetToken = (options: { forceRefresh: boolean }) => Promise<string>;
+
+/** Only a confirmed invalid/expired credential can trigger an automatic renewal. */
+export function isInvalidToken(group: string, code: string): boolean {
+	return (
+		group === "auth" &&
+		(code === "invalid_token" || code === "token_expired")
+	);
+}
+
+export function isInvalidTokenResponse(response: Response): boolean {
+	if (response.status !== 401) return false;
+	const error = response.headers.get("x-rivet-error");
+	if (!error) return false;
+	const separator = error.indexOf(".");
+	return (
+		separator > 0 &&
+		isInvalidToken(error.slice(0, separator), error.slice(separator + 1))
+	);
+}
+
+/** One token cache per client; concurrent requests and refreshes share the same issuance call. */
+export class TokenProvider {
+	#cached?: { token: string; expiresAt?: number };
+	#pending?: Promise<string>;
+	#pendingForce = false;
+
+	constructor(private readonly getToken: GetToken) {}
+
+	async current(): Promise<string> {
+		if (
+			this.#cached &&
+			(this.#cached.expiresAt === undefined ||
+				Date.now() < this.#cached.expiresAt - 5_000)
+		) {
+			return this.#cached.token;
+		}
+		return this.#load(false);
+	}
+
+	/** A late rejection for an older token must never erase a newer credential. */
+	async refreshIfCurrent(failedToken: string): Promise<string> {
+		if (this.#pending !== undefined) {
+			const wasForced = this.#pendingForce;
+			const refreshed = await this.#pending;
+			if (wasForced || refreshed !== failedToken) return refreshed;
+		}
+		if (this.#cached?.token !== failedToken) {
+			return this.current();
+		}
+		return this.#load(true);
+	}
+
+	#load(forceRefresh: boolean): Promise<string> {
+		if (this.#pending !== undefined) return this.#pending;
+
+		const pending = Promise.resolve()
+			.then(() => this.getToken({ forceRefresh }))
+			.then((token) => {
+				if (typeof token !== "string" || token.length === 0) {
+					throw new Error(
+						"getToken must resolve to a nonempty string",
+					);
+				}
+				this.#cached = { token, expiresAt: jwtExpiration(token) };
+				return token;
+			})
+			.finally(() => {
+				if (this.#pending === pending) this.#pending = undefined;
+			});
+		this.#pending = pending;
+		this.#pendingForce = forceRefresh;
+		return pending;
+	}
+}
+
+// This is only a cache hint, never a signature or authorization check. The Engine verifies JWTs.
+function jwtExpiration(token: string): number | undefined {
+	try {
+		const payload = token.split(".")[1];
+		if (!payload) return undefined;
+		const claims = JSON.parse(
+			atob(payload.replace(/-/g, "+").replace(/_/g, "/")),
+		) as { exp?: unknown };
+		return typeof claims.exp === "number" &&
+			Number.isSafeInteger(claims.exp)
+			? claims.exp * 1_000
+			: undefined;
+	} catch {
+		return undefined;
+	}
+}

package/src/client/utils.ts

@@ -229,6 +229,7 @@ export async function sendHttpRequest<
 					metadata: responseData.metadata,
 					rayId: rayId ?? undefined,
 					actor: responseData.actor,
+					statusCode: response.status,
 				},
 			);
 		} catch (error) {

package/src/engine-client/actor-websocket-client.ts

@@ -336,7 +336,8 @@ export async function openWebSocketToGateway(
 
 	logger().debug({
 		msg: "opening websocket to actor via guard",
-		gatewayUrl,
+		// Gateway URLs can embed credentials; never include the path or query in logs.
+		origin: new URL(gatewayUrl).origin,
 	});
 
 	// Create WebSocket connection

package/src/engine-client/api-endpoints.ts

@@ -1,5 +1,6 @@
 import { serializeActorKey } from "@/actor/keys";
 import type { ClientConfig } from "@/client/client";
+import type { IssuedToken, IssueTokenOptions } from "@/client/auth";
 import type { MetadataResponse } from "@/common/router";
 import type {
 	ActorsCreateRequest,
@@ -12,6 +13,33 @@ import type {
 import type { RivetId } from "@/engine-api/common";
 import { apiCall } from "./api-utils";
 
+/** Engine's JSON response uses snake_case; the public result uses camelCase. */
+export type IssueTokenResponse = {
+	token: IssuedToken["token"];
+	issued_ts: IssuedToken["issuedAt"];
+	expires_ts: IssuedToken["expiresAt"];
+};
+
+export async function issueToken(
+	config: ClientConfig,
+	options: IssueTokenOptions,
+): Promise<IssueTokenResponse> {
+	return apiCall<
+		{
+			namespace: string;
+			subject?: IssueTokenOptions["subject"];
+			duration?: IssueTokenOptions["expiresIn"];
+			grants: IssueTokenOptions["grants"];
+		},
+		IssueTokenResponse
+	>(config, "POST", "/auth/tokens", {
+		namespace: config.namespace,
+		subject: options.subject,
+		duration: options.expiresIn,
+		grants: options.grants,
+	});
+}
+
 // MARK: Get actor
 export async function getActor(
 	config: ClientConfig,

package/src/engine-client/driver.ts

@@ -1,6 +1,7 @@
 import type { Hono, Context as HonoContext } from "hono";
 import type { ActorKey, Encoding, UniversalWebSocket } from "@/actor/mod";
 import type { ActorQuery, CrashPolicy } from "@/client/query";
+import type { IssuedToken, IssueTokenOptions } from "@/client/auth";
 import type { RegistryConfig } from "@/registry/config";
 import type { GetUpgradeWebSocket } from "@/utils";
 
@@ -17,6 +18,7 @@ export function shouldSkipReadyWait(
 }
 
 export interface EngineControlClient {
+	issueToken(options: IssueTokenOptions): Promise<IssuedToken>;
 	getForId(input: GetForIdInput): Promise<ActorOutput | undefined>;
 	getWithKey(input: GetWithKeyInput): Promise<ActorOutput | undefined>;
 	getOrCreateWithKey(input: GetOrCreateWithKeyInput): Promise<ActorOutput>;
@@ -35,6 +37,8 @@ export interface EngineControlClient {
 		params: unknown,
 		options?: GatewayRequestOptions,
 	): Promise<UniversalWebSocket>;
+	/** Refresh a credential rejected during this WebSocket handshake, if supported. */
+	refreshAuthToken?(ws: UniversalWebSocket): Promise<boolean>;
 	proxyRequest(
 		c: HonoContext,
 		actorRequest: Request,