rivetkit 2.3.18-rc.4 → 2.3.19

package/src/client/actor-handle.ts

@@ -1,5 +1,6 @@
 import type { AnyActorDefinition } from "@/actor/definition";
-import type { ActorSpecifier } from "@/actor/errors";
+import { type ActorSpecifier, encodeErrorForBridge } from "@/actor/errors";
+import type { ActorInvocationSpanContext } from "@/common/actor-telemetry-context";
 import {
 	HEADER_CONN_PARAMS,
 	HEADER_ENCODING,
@@ -24,6 +25,7 @@ import { AsyncMutex } from "@/common/database/shared";
 import type { Encoding, JsonCompatValue } from "@/common/encoding";
 import { deconstructError } from "@/common/utils";
 import type { EngineControlClient } from "@/engine-client/driver";
+import type { StartCallSpan, CurrentActorInvocation } from "@/registry/runtime";
 import {
 	decodeCborCompat,
 	deserializeWithEncoding,
@@ -53,6 +55,7 @@ import { type ClientRaw, CREATE_ACTOR_CONN_PROXY } from "./client";
 import { ActorError, isSchedulingError } from "./errors";
 import { retryOnLifecycleBoundary } from "./lifecycle-errors";
 import { logger } from "./log";
+import { outboundTelemetryHeaders } from "./outbound-telemetry";
 import {
 	createQueueSender,
 	type QueueSendNoWaitOptions,
@@ -82,6 +85,8 @@ export class ActorHandleRaw {
 	#resolvedActorId?: string;
 	#resolvingActorId?: Promise<string>;
 	#queueSendMutex = new AsyncMutex();
+	#currentActorInvocation?: CurrentActorInvocation;
+	#startCallSpan?: StartCallSpan;
 
 	/**
 	 * Do not call this directly.
@@ -99,6 +104,8 @@ export class ActorHandleRaw {
 		actorResolutionState: ActorResolutionState,
 		gatewayOptions: ActorGatewayOptions = {},
 		signal?: AbortSignal,
+		currentActorInvocation?: CurrentActorInvocation,
+		startCallSpan?: StartCallSpan,
 	) {
 		this.#client = client;
 		this.#driver = driver;
@@ -108,6 +115,8 @@ export class ActorHandleRaw {
 		this.#params = params;
 		this.#getParams = getParams;
 		this.#signal = signal;
+		this.#currentActorInvocation = currentActorInvocation;
+		this.#startCallSpan = startCallSpan;
 	}
 
 	async #resolveConnectionParams(): Promise<unknown> {
@@ -166,6 +175,10 @@ export class ActorHandleRaw {
 					return await createQueueSender({
 						encoding: this.#encoding,
 						params: this.#params,
+						telemetryHeaders: () =>
+							outboundTelemetryHeaders(
+								this.#currentActorInvocation?.(),
+							),
 						customFetch: async (request: Request) => {
 							return await this.#driver.sendRequest(
 								target,
@@ -275,20 +288,41 @@ export class ActorHandleRaw {
 		// when no per-call signal is provided.
 		const signal = opts.signal ?? this.#signal;
 		const optsWithSignal = { ...opts, signal };
+		// Open before retries so all attempts share one span and parent context.
+		const call = this.#startCallSpan?.(
+			getActorNameFromQuery(this.#actorResolutionState),
+			opts.name,
+		);
 		const run = async () =>
-			(await this.#sendActionNow(optsWithSignal)) as Response;
-		if (opts.name === "destroy") {
-			return await run();
+			(await this.#sendActionAttempts(
+				optsWithSignal,
+				call?.span,
+			)) as Response;
+		const send = async () => {
+			if (opts.name === "destroy") {
+				return await run();
+			}
+			return await retryOnLifecycleBoundary(run, { signal });
+		};
+		if (!call) {
+			return await send();
+		}
+		try {
+			const output = await send();
+			call.finish();
+			return output;
+		} catch (error) {
+			call.finish(encodeErrorForBridge(error));
+			throw error;
 		}
-
-		return await retryOnLifecycleBoundary(run, { signal });
 	}
 
-	async #sendActionNow(
+	async #sendActionAttempts(
 		opts: {
 			name: string;
 			args: unknown[];
 		} & ActorActionOptions,
+		callSpan?: ActorInvocationSpanContext,
 	): Promise<unknown> {
 		const maxAttempts = this.#getDynamicQueryMaxAttempts();
 		let useQueryTarget = isDynamicActorQuery(this.#actorResolutionState);
@@ -320,6 +354,16 @@ export class ActorHandleRaw {
 					name: opts.name,
 					encoding: this.#encoding,
 				});
+				const headers: Record<string, string> = {
+					[HEADER_ENCODING]: this.#encoding,
+					...outboundTelemetryHeaders(
+						this.#currentActorInvocation?.(),
+						callSpan,
+					),
+				};
+				if (this.#params !== undefined) {
+					headers[HEADER_CONN_PARAMS] = JSON.stringify(this.#params);
+				}
 				const output = await sendHttpRequest<
 					protocol.HttpActionRequest,
 					protocol.HttpActionResponse,
@@ -330,16 +374,7 @@ export class ActorHandleRaw {
 				>({
 					url: `http://actor/action/${encodeURIComponent(opts.name)}`,
 					method: "POST",
-					headers: {
-						[HEADER_ENCODING]: this.#encoding,
-						...(this.#params !== undefined
-							? {
-									[HEADER_CONN_PARAMS]: JSON.stringify(
-										this.#params,
-									),
-								}
-							: {}),
-					},
+					headers,
 					body: opts.args,
 					encoding: this.#encoding,
 					customFetch: async (request) =>
@@ -643,6 +678,7 @@ export class ActorHandleRaw {
 			this.#encoding,
 			this.#actorResolutionState,
 			resolveActorGatewayOptions(this.#gatewayOptions, options),
+			this.#currentActorInvocation,
 		);
 
 		return this.#client[CREATE_ACTOR_CONN_PROXY](
@@ -680,6 +716,9 @@ export class ActorHandleRaw {
 				skipReadyWait,
 			},
 		);
+		const telemetryHeaders = outboundTelemetryHeaders(
+			this.#currentActorInvocation?.(),
+		);
 
 		for (let attempt = 0; attempt < maxAttempts; attempt++) {
 			let actorId: string | undefined;
@@ -698,6 +737,7 @@ export class ActorHandleRaw {
 					clonesInputBody ? input.clone() : input,
 					requestInit,
 					gatewayOptions,
+					telemetryHeaders,
 				);
 				const retry = await this.#shouldRetryRawFetchResponse(
 					response,

package/src/client/client.ts

@@ -3,6 +3,7 @@ import type { ActorQuery } from "@/client/query";
 import type { Encoding } from "@/common/encoding";
 import type { EngineControlClient } from "@/engine-client/driver";
 import type { Registry } from "@/registry";
+import type { StartCallSpan, CurrentActorInvocation } from "@/registry/runtime";
 import type { ActorActionFunction, ActorGatewayOptions } from "./actor-common";
 import {
 	type ActorConn,
@@ -175,6 +176,20 @@ export interface Region {
 export const ACTOR_CONNS_SYMBOL = Symbol("actorConns");
 export const CREATE_ACTOR_CONN_PROXY = Symbol("createActorConnProxy");
 
+/** Options for constructing a raw actor client. */
+export interface ClientRawOptions {
+	encoding?: Encoding;
+	gateway?: ActorGatewayOptions;
+}
+
+/** Runtime-supplied wiring for clients created inside an actor. Set only by the native runtime. */
+export interface ActorClientRuntimeOptions {
+	/** Supplies the calling actor's invocation so actor-to-actor clients propagate its trace and ray ID. */
+	currentActorInvocation?: CurrentActorInvocation;
+	/** Opens and finishes the Core span covering an actor-to-actor call. */
+	startCallSpan?: StartCallSpan;
+}
+
 /**
  * Client for managing & connecting to actors.
  *
@@ -189,19 +204,22 @@ export class ClientRaw {
 	#driver: EngineControlClient;
 	#encodingKind: Encoding;
 	#gatewayOptions: ActorGatewayOptions;
+	#currentActorInvocation?: CurrentActorInvocation;
+	#startCallSpan?: StartCallSpan;
 
 	/**
 	 * Creates an instance of Client.
 	 */
 	public constructor(
 		driver: EngineControlClient,
-		encoding: Encoding | undefined,
-		gatewayOptions: ActorGatewayOptions = {},
+		options: ClientRawOptions & ActorClientRuntimeOptions = {},
 	) {
 		this.#driver = driver;
 
-		this.#encodingKind = encoding ?? "bare";
-		this.#gatewayOptions = gatewayOptions;
+		this.#encodingKind = options.encoding ?? "bare";
+		this.#gatewayOptions = options.gateway ?? {};
+		this.#currentActorInvocation = options.currentActorInvocation;
+		this.#startCallSpan = options.startCallSpan;
 	}
 
 	/**
@@ -403,6 +421,8 @@ export class ClientRaw {
 			actorQuery,
 			this.#gatewayOptions,
 			signal,
+			this.#currentActorInvocation,
+			this.#startCallSpan,
 		);
 	}
 
@@ -459,9 +479,9 @@ export type AnyClient = Client<Registry<any>>;
 
 export function createClientWithDriver<A extends Registry<any>>(
 	driver: EngineControlClient,
-	config: { encoding?: Encoding; gateway?: ActorGatewayOptions } = {},
+	options: ClientRawOptions & ActorClientRuntimeOptions = {},
 ): Client<A> {
-	const client = new ClientRaw(driver, config.encoding, config.gateway);
+	const client = new ClientRaw(driver, options);
 
 	// Create proxy for accessing actors by name
 	return new Proxy(client, {

package/src/client/outbound-telemetry.ts

@@ -0,0 +1,66 @@
+import {
+	HEADER_RIVET_RAY_ID,
+	HEADER_TRACEPARENT,
+	HEADER_TRACESTATE,
+} from "@/common/actor-router-consts";
+import type { ActorInvocationTraceContext } from "@/common/actor-telemetry-context";
+import {
+	actorInvocationTraceHeaders,
+	readActiveRayId,
+	readActiveTraceHeaders,
+} from "@/common/otel-context";
+
+/**
+ * Headers that carry a caller's ray ID and trace context into an actor. One
+ * rule for every kind of outbound call, so an action call and a queue send
+ * made from the same place land in the same trace under the same ray ID.
+ *
+ * The ray ID is the calling invocation's when the caller is itself inside an
+ * actor, else the one placed in OpenTelemetry baggage by the surrounding
+ * request handler. The trace context is `callSpan` when the runtime opened a
+ * span for this call, else the application span active in this JavaScript
+ * context, else the calling actor's own Core invocation span.
+ */
+export function outboundTelemetryHeaders(
+	invocation: ActorInvocationTraceContext | undefined,
+	callSpan?: ActorInvocationTraceContext["span"],
+): Record<string, string> {
+	const headers: Record<string, string> = {};
+	const rayId = invocation?.rayId ?? readActiveRayId();
+	if (rayId) {
+		headers[HEADER_RIVET_RAY_ID] = rayId;
+	}
+	const traceHeaders =
+		actorInvocationTraceHeaders(callSpan) ??
+		readActiveTraceHeaders() ??
+		actorInvocationTraceHeaders(invocation?.span);
+	if (traceHeaders) {
+		headers[HEADER_TRACEPARENT] = traceHeaders.traceparent;
+		if (traceHeaders.tracestate) {
+			headers[HEADER_TRACESTATE] = traceHeaders.tracestate;
+		}
+	}
+	return headers;
+}
+
+/**
+ * Adds outbound telemetry headers to a request whose caller may have set
+ * some already. A header the caller set wins. `traceparent` and `tracestate`
+ * describe one span between them, so when the caller set either, both stay
+ * as the caller set them and neither is added.
+ */
+export function addOutboundTelemetryHeaders(
+	headers: Headers,
+	telemetry: Record<string, string>,
+): void {
+	const callerSetTraceContext =
+		headers.has(HEADER_TRACEPARENT) || headers.has(HEADER_TRACESTATE);
+	for (const [name, value] of Object.entries(telemetry)) {
+		const isTraceContext =
+			name === HEADER_TRACEPARENT || name === HEADER_TRACESTATE;
+		if (isTraceContext ? callerSetTraceContext : headers.has(name)) {
+			continue;
+		}
+		headers.set(name, value);
+	}
+}

package/src/client/queue.ts

@@ -54,6 +54,8 @@ export interface QueueSendResult<TResponse = unknown> {
 interface QueueSenderOptions {
 	encoding: Encoding;
 	params: unknown;
+	/** Ray ID and trace context headers, read per send. */
+	telemetryHeaders: () => Record<string, string>;
 	customFetch: (request: Request) => Promise<Response>;
 }
 
@@ -90,6 +92,7 @@ export function createQueueSender(
 			method: "POST",
 			headers: {
 				[HEADER_ENCODING]: senderOptions.encoding,
+				...senderOptions.telemetryHeaders(),
 				...(senderOptions.params !== undefined
 					? {
 							[HEADER_CONN_PARAMS]: JSON.stringify(

package/src/client/raw-utils.ts

@@ -12,6 +12,7 @@ import type {
 } from "@/engine-client/driver";
 import { ActorError } from "./errors";
 import { logger } from "./log";
+import { addOutboundTelemetryHeaders } from "./outbound-telemetry";
 
 /** Buffer a one-shot ReadableStream body to bytes so every retry attempt can re-send it. */
 export async function prepareRetryableInit(
@@ -37,6 +38,7 @@ export async function rawHttpFetch(
 	input: string | URL | Request,
 	init?: RequestInit,
 	options: GatewayRequestOptions = {},
+	telemetryHeaders: Record<string, string> = {},
 ): Promise<Response> {
 	// Extract path and merge init options
 	let path: string;
@@ -111,6 +113,7 @@ export async function rawHttpFetch(
 		if (params) {
 			proxyRequestHeaders.set(HEADER_CONN_PARAMS, JSON.stringify(params));
 		}
+		addOutboundTelemetryHeaders(proxyRequestHeaders, telemetryHeaders);
 
 		// Forward the request to the actor
 		const proxyRequest = new Request(url, {

package/src/common/actor-router-consts.ts

@@ -20,6 +20,9 @@ export const HEADER_ACTOR_GENERATION = "x-rivet-actor-generation";
 export const HEADER_ACTOR_KEY = "x-rivet-actor-key";
 
 export const HEADER_RIVET_TOKEN = "x-rivet-token";
+export const HEADER_RIVET_RAY_ID = "x-rivet-ray-id";
+export const HEADER_TRACEPARENT = "traceparent";
+export const HEADER_TRACESTATE = "tracestate";
 
 // MARK: Manager Gateway Headers
 export const HEADER_RIVET_TARGET = "x-rivet-target";

package/src/common/actor-telemetry-context.ts

@@ -0,0 +1,19 @@
+/** Correlation owned by the currently executing Core actor invocation. */
+export interface ActorInvocationTraceContext {
+	/** Ray ID the invocation received, absent when nothing upstream supplied one. */
+	readonly rayId?: string;
+	/** Core-owned invocation span context, absent when tracing is disabled. */
+	readonly span?: ActorInvocationSpanContext;
+}
+
+/** W3C span context of the Core invocation span. */
+export interface ActorInvocationSpanContext {
+	/** W3C trace identifier. */
+	readonly traceId: string;
+	/** W3C span identifier for the Core invocation. */
+	readonly spanId: string;
+	/** OpenTelemetry trace flags encoded as an integer. */
+	readonly traceFlags: number;
+	/** Optional vendor trace state inherited by the Core invocation. */
+	readonly tracestate?: string;
+}

package/src/common/otel-context.ts

@@ -0,0 +1,99 @@
+import {
+	type Context,
+	context,
+	createTraceState,
+	isSpanContextValid,
+	propagation,
+	type SpanContext,
+	trace,
+} from "@opentelemetry/api";
+import type { ActorInvocationSpanContext } from "./actor-telemetry-context";
+
+/** W3C headers derived from the active JavaScript OTel context. */
+export interface ActiveTraceHeaders {
+	/** W3C Trace Context identifying the active trace and span. */
+	readonly traceparent: string;
+	/** Optional vendor trace state associated with the active span. */
+	readonly tracestate?: string;
+}
+
+/** Returns the active W3C trace context, when an OTel provider has installed one. */
+export function readActiveTraceHeaders(): ActiveTraceHeaders | undefined {
+	return traceHeaders(trace.getSpanContext(context.active()));
+}
+
+/** W3C Baggage key that carries a ray ID through application code. */
+export const RAY_BAGGAGE_KEY = "rivet.ray.id";
+
+/**
+ * Returns the ray ID carried in the active OpenTelemetry baggage, so a request
+ * handler that received a ray ID can pass it to the actors it calls. The bound
+ * is the one Core applies on receipt, so an invalid value is never sent.
+ */
+export function readActiveRayId(): string | undefined {
+	const value = propagation
+		.getBaggage(context.active())
+		?.getEntry(RAY_BAGGAGE_KEY)?.value;
+	if (value && value.length <= 30 && /^[A-Za-z0-9_-]+$/.test(value)) {
+		return value;
+	}
+	return undefined;
+}
+
+/** W3C headers for the Core invocation span, or nothing when it is absent or invalid. */
+export function actorInvocationTraceHeaders(
+	invocation: ActorInvocationSpanContext | undefined,
+): ActiveTraceHeaders | undefined {
+	const spanContext = invocationSpanContext(invocation);
+	if (!spanContext) return undefined;
+	return traceHeaders(spanContext);
+}
+
+function traceHeaders(
+	spanContext: SpanContext | undefined,
+): ActiveTraceHeaders | undefined {
+	if (!spanContext || !isSpanContextValid(spanContext)) return undefined;
+	const tracestate = spanContext.traceState?.serialize();
+	return {
+		traceparent: `00-${spanContext.traceId}-${spanContext.spanId}-${spanContext.traceFlags.toString(16).padStart(2, "0")}`,
+		...(tracestate ? { tracestate } : {}),
+	};
+}
+
+/**
+ * Runs `run` with the Core invocation span as the active OpenTelemetry span,
+ * so application spans started inside an actor callback nest under it. With
+ * no span, or an invalid one, `run` executes unchanged.
+ */
+export function runWithActorInvocationSpan<T>(
+	invocation: ActorInvocationSpanContext | undefined,
+	run: () => T,
+): T {
+	if (!invocation) return run();
+
+	const spanContext = invocationSpanContext(invocation);
+	if (!spanContext) return run();
+	const parent: Context = trace.setSpanContext(context.active(), spanContext);
+
+	return context.with(parent, run);
+}
+
+function invocationSpanContext(
+	invocation: ActorInvocationSpanContext | undefined,
+): SpanContext | undefined {
+	if (!invocation) return undefined;
+	try {
+		const spanContext: SpanContext = {
+			traceId: invocation.traceId,
+			spanId: invocation.spanId,
+			traceFlags: invocation.traceFlags,
+			traceState: invocation.tracestate
+				? createTraceState(invocation.tracestate)
+				: undefined,
+			isRemote: false,
+		};
+		return isSpanContextValid(spanContext) ? spanContext : undefined;
+	} catch {
+		return undefined;
+	}
+}

package/src/common/utils.ts

@@ -48,24 +48,6 @@ export interface DeconstructedError {
 	actor?: errors.ActorSpecifier;
 }
 
-function isCanonicalStructuredRivetError(
-	error: unknown,
-): error is errors.RivetErrorLike {
-	return (
-		error instanceof errors.RivetError ||
-		(typeof error === "object" &&
-			error !== null &&
-			"__type" in error &&
-			error.__type === "RivetError" &&
-			"group" in error &&
-			typeof error.group === "string" &&
-			"code" in error &&
-			typeof error.code === "string" &&
-			"message" in error &&
-			typeof error.message === "string")
-	);
-}
-
 /**
  * Deconstructs errors into response fields. Bridge callback errors that cross
  * into rivetkit-core are sanitized there; this only classifies JS-local errors.
@@ -85,7 +67,7 @@ export function deconstructError(
 	let actor: errors.ActorSpecifier | undefined;
 	// Structured errors from core or from pre-built `RivetError` instances are canonical.
 	// Only unstructured errors go through the classifier below.
-	if (isCanonicalStructuredRivetError(error)) {
+	if (errors.isCanonicalStructuredRivetError(error)) {
 		statusCode = (
 			typeof error.statusCode === "number"
 				? error.statusCode

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

@@ -1,9 +1,12 @@
 import type { ClientConfig } from "@/client/config";
 import {
 	HEADER_RIVET_ACTOR,
+	HEADER_RIVET_RAY_ID,
 	HEADER_RIVET_SKIP_READY_WAIT,
 	HEADER_RIVET_TARGET,
 	HEADER_RIVET_TOKEN,
+	HEADER_TRACEPARENT,
+	HEADER_TRACESTATE,
 } from "@/common/actor-router-consts";
 import { type GatewayRequestOptions, shouldSkipReadyWait } from "./driver";
 
@@ -55,6 +58,18 @@ function buildGuardHeaders(
 	for (const [key, value] of Object.entries(runConfig.headers)) {
 		headers.set(key, value as string);
 	}
+	// Configured trace context is never sent: it would pin every call to one span.
+	for (const name of [HEADER_TRACEPARENT, HEADER_TRACESTATE]) {
+		headers.delete(name);
+		const value = actorRequest.headers.get(name);
+		if (value !== null) {
+			headers.set(name, value);
+		}
+	}
+	const requestRayId = actorRequest.headers.get(HEADER_RIVET_RAY_ID);
+	if (requestRayId !== null) {
+		headers.set(HEADER_RIVET_RAY_ID, requestRayId);
+	}
 	// Add guard-specific headers
 	if (runConfig.token) {
 		headers.set(HEADER_RIVET_TOKEN, runConfig.token);

package/src/registry/index.ts

@@ -732,6 +732,7 @@ export class Registry<A extends RegistryActors> {
 	async #drain(config: RegistryConfig): Promise<void> {
 		const modeAPromise = this.#runtimeServeConfiguredPromise;
 		const modeBPromise = this.#runtimeServerlessPromise;
+		const runtimes = new Set<CoreRuntime>();
 
 		const gracePeriodMs =
 			config.shutdown?.gracePeriodMs ??
@@ -751,6 +752,7 @@ export class Registry<A extends RegistryActors> {
 					(async () => {
 						try {
 							const { runtime, registry } = await modeAPromise;
+							runtimes.add(runtime);
 							await runtime.shutdownRegistry(registry);
 						} catch (err) {
 							logger().warn(
@@ -766,6 +768,7 @@ export class Registry<A extends RegistryActors> {
 					(async () => {
 						try {
 							const { runtime, registry } = await modeBPromise;
+							runtimes.add(runtime);
 							await runtime.shutdownRegistry(registry);
 						} catch (err) {
 							logger().warn(
@@ -795,6 +798,9 @@ export class Registry<A extends RegistryActors> {
 				setTimeout(resolve, gracePeriodMs).unref?.(),
 			),
 		]);
+		await Promise.all(
+			[...runtimes].map((runtime) => runtime.shutdownTelemetry?.()),
+		);
 	}
 
 	async #actorStopThresholdMs(