rivetkit 2.3.18 → 2.3.19

package/src/registry/native.ts

@@ -8,6 +8,8 @@ import {
 	type ActorCron,
 	type ActorCronEveryOptions,
 	type ActorCronSetOptions,
+	type ActorLogger,
+	type ActorRun,
 	type ActorSchedule,
 	CONN_STATE_MANAGER_SYMBOL,
 	type CronFire,
@@ -46,6 +48,7 @@ import {
 } from "@/client/client";
 import { convertRegistryConfigToClientConfig } from "@/client/config";
 import { HEADER_CONN_PARAMS } from "@/common/actor-router-consts";
+import type { ActorInvocationTraceContext } from "@/common/actor-telemetry-context";
 import type {
 	AnyDatabaseProvider,
 	SqliteProfilingOptions,
@@ -2706,6 +2709,9 @@ class NativeConnectionMap implements ReadonlyMap<string, NativeConnAdapter> {
 	readonly [Symbol.toStringTag] = "NativeConnectionMap";
 }
 
+/** Logger cache key for code that runs outside any invocation. */
+const NO_INVOCATION_SCOPE = {};
+
 export class ActorContextHandleAdapter {
 	#runtime: CoreRuntime;
 	#ctx: ActorContextHandle;
@@ -2720,10 +2726,11 @@ export class ActorContextHandleAdapter {
 	#db?: unknown;
 	#dispatchCancelToken?: CancellationTokenHandle;
 	#kv?: NativeKvAdapter;
+	#logByInvocationScope = new WeakMap<object, ActorLogger>();
 	#queue?: NativeQueueAdapter;
 	#request?: Request;
 	#schedule?: NativeScheduleAdapter;
-	#run?: { setWakeAt(timestamp: number | null): Promise<void> };
+	#run?: ActorRun;
 	#runHandlerConfigured: boolean;
 	#onStateChange?: NativeOnStateChangeHandler;
 	#stateEnabled: boolean;
@@ -2904,6 +2911,12 @@ export class ActorContextHandleAdapter {
 						this.#runtime.actorSetRunWakeAt(this.#ctx, timestamp),
 					);
 				},
+				startWorkflowSpan: () =>
+					callNative(() =>
+						this.#runtime.startWorkflowSpan(this.#ctx),
+					),
+				runOutsideWorkflowSpan: (body) =>
+					this.#runtime.runOutsideActorInvocationContext(body),
 			};
 		}
 		return this.#run;
@@ -2956,8 +2969,32 @@ export class ActorContextHandleAdapter {
 		return this.#connMap;
 	}
 
+	#invocationTraceContext(): ActorInvocationTraceContext | undefined {
+		return callNativeSync(() =>
+			this.#runtime.actorInvocationTraceContext(this.#ctx),
+		);
+	}
+
+	/** Cached per invocation, because the `run` context outlives workflow runs and steps. */
 	get log() {
-		return logger();
+		const scope =
+			this.#runtime.currentInvocationScope() ?? NO_INVOCATION_SCOPE;
+		const cached = this.#logByInvocationScope.get(scope);
+		if (cached) return cached;
+
+		const invocation = this.#invocationTraceContext();
+		const log = logger().child({
+			actorId: this.actorId,
+			actorName: this.name,
+			actorKey: this.key,
+			...(invocation?.rayId && { rayId: invocation.rayId }),
+			...(invocation?.span && {
+				traceId: invocation.span.traceId,
+				spanId: invocation.span.spanId,
+			}),
+		});
+		this.#logByInvocationScope.set(scope, log);
+		return log;
 	}
 
 	get abortSignal(): AbortSignal {
@@ -3995,12 +4032,22 @@ export function buildNativeFactory(
 		events: config.events,
 		queues: config.queues,
 	};
-	const createClient = () =>
+	const createClient = (ctx: ActorContextHandle) =>
 		createClientWithDriver(
 			new RemoteEngineControlClient(
 				convertRegistryConfigToClientConfig(registryConfig),
 			),
-			{ encoding: "bare" },
+			{
+				encoding: "bare",
+				currentActorInvocation: () =>
+					callNativeSync(() =>
+						runtime.actorInvocationTraceContext(ctx),
+					),
+				startCallSpan: (actorName, actionName) =>
+					callNativeSync(() =>
+						runtime.startCallSpan(ctx, actorName, actionName),
+					),
+			},
 		);
 	const run = getRunFunction(config.run);
 	const runHandlerCoordinator =
@@ -4044,7 +4091,7 @@ export function buildNativeFactory(
 		new ActorContextHandleAdapter(
 			runtime,
 			ctx,
-			createClient,
+			() => createClient(ctx),
 			schemaConfig,
 			databaseProvider,
 			request,
@@ -4063,7 +4110,7 @@ export function buildNativeFactory(
 			runtime,
 			ctx,
 			conn,
-			createClient,
+			() => createClient(ctx),
 			schemaConfig,
 			databaseProvider,
 			request,
@@ -5041,129 +5088,147 @@ export function buildNativeFactory(
 				try {
 					const { ctx, request, cancelToken, responseBodyStream } =
 						unwrapTsfnPayload(error, payload);
-					const inspectorResponse =
-						await maybeHandleNativeInspectorRequest(ctx, request);
-					if (inspectorResponse) {
-						await cancelNativeHttpRequestBody(request.bodyStream);
-						return (
-							await convertNativeHttpResponse(
-								inspectorResponse,
-								responseBodyStream,
-							)
-						).response;
-					}
-
-					if (typeof config.onRequest !== "function") {
-						await cancelNativeHttpRequestBody(request.bodyStream);
-						return (
-							await convertNativeHttpResponse(
-								new Response(null, { status: 404 }),
-								responseBodyStream,
-							)
-						).response;
-					}
+					return await runtime.runWithActorInvocationContext(
+						ctx,
+						async () => {
+							const inspectorResponse =
+								await maybeHandleNativeInspectorRequest(
+									ctx,
+									request,
+								);
+							if (inspectorResponse) {
+								await cancelNativeHttpRequestBody(
+									request.bodyStream,
+								);
+								return (
+									await convertNativeHttpResponse(
+										inspectorResponse,
+										responseBodyStream,
+									)
+								).response;
+							}
 
-					const requestAbortController = new AbortController();
-					const handlerRequest = buildNativeHttpRequest({
-						...request,
-						abortController: requestAbortController,
-					});
-					const rawConnParams =
-						handlerRequest.headers.get(HEADER_CONN_PARAMS);
-					let requestCtx:
-						| ReturnType<typeof withConnContext>
-						| undefined;
-					let conn: ConnHandle | undefined;
-					let removeRequestAbortListener: (() => void) | undefined;
-					let cleanupDeferredToBody = false;
-					let cleanedUp = false;
-					const cleanupRequest = async () => {
-						if (cleanedUp) return;
-						cleanedUp = true;
-						removeRequestAbortListener?.();
-						try {
-							await requestCtx?.dispose();
-						} finally {
-							if (conn) {
-								await runtime.connDisconnect(conn);
+							if (typeof config.onRequest !== "function") {
+								await cancelNativeHttpRequestBody(
+									request.bodyStream,
+								);
+								return (
+									await convertNativeHttpResponse(
+										new Response(null, { status: 404 }),
+										responseBodyStream,
+									)
+								).response;
 							}
-						}
-					};
-					try {
-						const connParams = validateConnParams(
-							schemaConfig.connParamsSchema,
-							rawConnParams
-								? JSON.parse(rawConnParams)
-								: undefined,
-						);
-						conn = await callNative(() =>
-							runtime.actorConnectConn(
-								ctx,
-								encodeValue(connParams),
-								request,
-							),
-						);
-						requestCtx = makeConnCtx(
-							ctx,
-							conn,
-							handlerRequest,
-							cancelToken,
-						);
-						const ctxAbortSignal = requestCtx.abortSignal;
-						const abortRequest = () =>
-							requestAbortController.abort(ctxAbortSignal.reason);
-						if (ctxAbortSignal.aborted) {
-							abortRequest();
-						} else {
-							ctxAbortSignal.addEventListener(
-								"abort",
-								abortRequest,
-								{ once: true },
-							);
-							removeRequestAbortListener = () =>
-								ctxAbortSignal.removeEventListener(
-									"abort",
-									abortRequest,
+
+							const requestAbortController =
+								new AbortController();
+							const handlerRequest = buildNativeHttpRequest({
+								...request,
+								abortController: requestAbortController,
+							});
+							const rawConnParams =
+								handlerRequest.headers.get(HEADER_CONN_PARAMS);
+							let requestCtx:
+								| ReturnType<typeof withConnContext>
+								| undefined;
+							let conn: ConnHandle | undefined;
+							let removeRequestAbortListener:
+								| (() => void)
+								| undefined;
+							let cleanupDeferredToBody = false;
+							let cleanedUp = false;
+							const cleanupRequest = async () => {
+								if (cleanedUp) return;
+								cleanedUp = true;
+								removeRequestAbortListener?.();
+								try {
+									await requestCtx?.dispose();
+								} finally {
+									if (conn) {
+										await runtime.connDisconnect(conn);
+									}
+								}
+							};
+							try {
+								const connParams = validateConnParams(
+									schemaConfig.connParamsSchema,
+									rawConnParams
+										? JSON.parse(rawConnParams)
+										: undefined,
 								);
-						}
-						const response = await config.onRequest(
-							requestCtx,
-							handlerRequest,
-						);
-						if (!isResponseLike(response)) {
-							throw new Error(
-								"onRequest handler must return a Response",
-							);
-						}
-						const conversion = await convertNativeHttpResponse(
-							response,
-							responseBodyStream,
-						);
-						if (conversion.bodyCompletion) {
-							cleanupDeferredToBody = true;
-							void conversion.bodyCompletion
-								.then(cleanupRequest)
-								.catch((cleanupError) => {
-									logger().error({
-										msg: "failed to clean up native streaming http request",
-										error: cleanupError,
-									});
-								});
-						}
-						return conversion.response;
-					} finally {
-						try {
-							// Handler completion ends upload ownership even when
-							// the Web Request body is locked or partly consumed.
-							await cancelNativeHttpRequestBody(
-								request.bodyStream,
-							);
-						} finally {
-							if (!cleanupDeferredToBody) {
-								await cleanupRequest();
+								conn = await callNative(() =>
+									runtime.actorConnectConn(
+										ctx,
+										encodeValue(connParams),
+										request,
+									),
+								);
+								requestCtx = makeConnCtx(
+									ctx,
+									conn,
+									handlerRequest,
+									cancelToken,
+								);
+								const ctxAbortSignal = requestCtx.abortSignal;
+								const abortRequest = () =>
+									requestAbortController.abort(
+										ctxAbortSignal.reason,
+									);
+								if (ctxAbortSignal.aborted) {
+									abortRequest();
+								} else {
+									ctxAbortSignal.addEventListener(
+										"abort",
+										abortRequest,
+										{ once: true },
+									);
+									removeRequestAbortListener = () =>
+										ctxAbortSignal.removeEventListener(
+											"abort",
+											abortRequest,
+										);
+								}
+								const response = await config.onRequest(
+									requestCtx,
+									handlerRequest,
+								);
+								if (!isResponseLike(response)) {
+									throw new Error(
+										"onRequest handler must return a Response",
+									);
+								}
+								const conversion =
+									await convertNativeHttpResponse(
+										response,
+										responseBodyStream,
+									);
+								if (conversion.bodyCompletion) {
+									cleanupDeferredToBody = true;
+									void conversion.bodyCompletion
+										.then(cleanupRequest)
+										.catch((cleanupError) => {
+											logger().error({
+												msg: "failed to clean up native streaming http request",
+												error: cleanupError,
+											});
+										});
+								}
+								return conversion.response;
+							} finally {
+								try {
+									// Handler completion ends upload ownership even when
+									// the Web Request body is locked or partly consumed.
+									await cancelNativeHttpRequestBody(
+										request.bodyStream,
+									);
+								} finally {
+									if (!cleanupDeferredToBody) {
+										await cleanupRequest();
+									}
+								}
 							}
-						}
-					}
+						},
+					);
 				} catch (error) {
 					logger().error({
 						msg: "native onRequest failed",
@@ -5312,21 +5377,29 @@ export function buildNativeFactory(
 							conn != null
 								? makeConnCtx(ctx, conn, undefined, cancelToken)
 								: makeActorCtx(ctx, undefined, cancelToken);
-						try {
-							return encodeValue(
-								await handler(
-									actorCtx,
-									...validateActionArgs(
-										schemaConfig.actionInputSchemas,
-										name,
-										decodeArgs(args),
+						const runAction = async () => {
+							try {
+								return encodeValue(
+									await handler(
+										actorCtx,
+										...validateActionArgs(
+											schemaConfig.actionInputSchemas,
+											name,
+											decodeArgs(args),
+										),
+										...(scheduledFire
+											? [scheduledFire]
+											: []),
 									),
-									...(scheduledFire ? [scheduledFire] : []),
-								),
-							);
-						} finally {
-							await actorCtx.dispose();
-						}
+								);
+							} finally {
+								await actorCtx.dispose();
+							}
+						};
+						return await runtime.runWithActorInvocationContext(
+							ctx,
+							runAction,
+						);
 					},
 				),
 			]),
@@ -5365,7 +5438,7 @@ export function buildNativeFactory(
 					runtime,
 					ctx,
 					conn,
-					createClient,
+					() => createClient(ctx),
 					schemaConfig,
 					databaseProvider,
 					jsRequest,
@@ -5374,62 +5447,75 @@ export function buildNativeFactory(
 					cancelToken,
 					run !== undefined,
 				);
-				try {
-					if (
-						!schemaConfig.queues ||
-						!hasSchemaConfigKey(schemaConfig.queues, name)
-					) {
-						return { status: "completed" };
-					}
-
-					const canPublish = getQueueCanPublish(
-						schemaConfig.queues,
-						name,
-					);
-					if (canPublish && !(await canPublish(actorCtx))) {
-						throw forbiddenError();
-					}
-
-					const decodedBody = decodeValue(body);
-					if (wait) {
+				return await runtime.runWithActorInvocationContext(
+					ctx,
+					async () => {
 						try {
-							const response =
-								await actorCtx.queue.enqueueAndWait(
-									name,
-									decodedBody,
-									{
-										timeout:
-											timeoutMs === undefined ||
-											timeoutMs === null
-												? undefined
-												: Number(timeoutMs),
-									},
-								);
-							return {
-								status: "completed",
-								response:
-									response === undefined
-										? undefined
-										: encodeValue(response),
-							};
-						} catch (error) {
 							if (
-								(error as { group?: string; code?: string })
-									.group === "queue" &&
-								(error as { group?: string; code?: string })
-									.code === "timed_out"
+								!schemaConfig.queues ||
+								!hasSchemaConfigKey(schemaConfig.queues, name)
 							) {
-								return { status: "timedOut" };
+								return { status: "completed" };
 							}
-							throw error;
-						}
-					}
 
-					await actorCtx.queue.send(name, decodedBody);
-					return { status: "completed" };
-				} finally {
-					await actorCtx.dispose();
-				}
+							const canPublish = getQueueCanPublish(
+								schemaConfig.queues,
+								name,
+							);
+							if (canPublish && !(await canPublish(actorCtx))) {
+								throw forbiddenError();
+							}
+
+							const decodedBody = decodeValue(body);
+							if (wait) {
+								try {
+									const response =
+										await actorCtx.queue.enqueueAndWait(
+											name,
+											decodedBody,
+											{
+												timeout:
+													timeoutMs === undefined ||
+													timeoutMs === null
+														? undefined
+														: Number(timeoutMs),
+											},
+										);
+									return {
+										status: "completed",
+										response:
+											response === undefined
+												? undefined
+												: encodeValue(response),
+									};
+								} catch (error) {
+									if (
+										(
+											error as {
+												group?: string;
+												code?: string;
+											}
+										).group === "queue" &&
+										(
+											error as {
+												group?: string;
+												code?: string;
+											}
+										).code === "timed_out"
+									) {
+										return { status: "timedOut" };
+									}
+									throw error;
+								}
+							}
+
+							await actorCtx.queue.send(name, decodedBody);
+							return { status: "completed" };
+						} finally {
+							await actorCtx.dispose();
+						}
+					},
+				);
 			},
 		),
 		serializeState: wrapNativeCallback(

package/src/registry/runtime.ts

@@ -1,3 +1,8 @@
+import type { WorkflowSpan } from "@/actor/config";
+import type {
+	ActorInvocationSpanContext,
+	ActorInvocationTraceContext,
+} from "@/common/actor-telemetry-context";
 import type {
 	SqliteNativeMetrics,
 	SqliteProfilingOptions,
@@ -29,6 +34,35 @@ export interface RuntimeActorKeySegment {
 	numberValue?: number;
 }
 
+/** Resolves correlation at operation time so retained clients cannot freeze stale context. */
+export type CurrentActorInvocation = () =>
+	| ActorInvocationTraceContext
+	| undefined;
+
+/** One open call from an actor out to another actor. */
+export interface RuntimeOutboundCall {
+	/**
+	 * W3C context of the call's own span, to send to the callee so it parents to
+	 * the call. Absent when tracing is disabled.
+	 */
+	readonly span?: ActorInvocationSpanContext;
+	/**
+	 * Records the call's outcome. `error` is the failure encoded the way the
+	 * bridge encodes errors, so a structured error keeps its group and code.
+	 */
+	finish(error?: string): void;
+}
+
+/**
+ * Opens the span covering one call out to another actor, or returns `undefined`
+ * outside an invocation or on a runtime without invocation telemetry. Callers
+ * send their own context when it returns nothing.
+ */
+export type StartCallSpan = (
+	actorName: string,
+	actionName: string,
+) => RuntimeOutboundCall | undefined;
+
 export interface RuntimeHttpRequest {
 	method: string;
 	uri: string;
@@ -452,6 +486,7 @@ export interface CoreRuntime {
 	): Promise<void>;
 	waitRegistryReady(registry: RegistryHandle): Promise<void>;
 	shutdownRegistry(registry: RegistryHandle): Promise<void>;
+	shutdownTelemetry?(): Promise<void>;
 	registryActorStopThresholdMs?(
 		registry: RegistryHandle,
 	): Promise<number | undefined>;
@@ -540,6 +575,37 @@ export interface CoreRuntime {
 		writes: RuntimeWorkflowKvWrite[],
 	): Promise<void>;
 	actorId(ctx: ActorContextHandle): string;
+	/**
+	 * Runs one actor callback with `ctx` as the current invocation: operations
+	 * on retained handles for the same actor resolve to it, and its Core span
+	 * is the active OpenTelemetry span for the duration of `run`. Only actions,
+	 * raw requests, and queue dispatch run this way; other callbacks have no
+	 * invocation span.
+	 */
+	runWithActorInvocationContext<T>(ctx: ActorContextHandle, run: () => T): T;
+	/**
+	 * Correlation of the invocation currently executing for this actor, or
+	 * `undefined` outside an invocation or after it finished. A sampled-out
+	 * invocation can still expose valid span context for propagation.
+	 */
+	actorInvocationTraceContext(
+		ctx: ActorContextHandle,
+	): ActorInvocationTraceContext | undefined;
+	/**
+	 * Opens the span covering one call this actor makes to another actor. See
+	 * `StartCallSpan` for when this returns nothing.
+	 */
+	startCallSpan(
+		ctx: ActorContextHandle,
+		actorName: string,
+		actionName: string,
+	): RuntimeOutboundCall | undefined;
+	/** Backs `ActorRun.startWorkflowSpan`. */
+	startWorkflowSpan(ctx: ActorContextHandle): Promise<WorkflowSpan>;
+	/** Backs `ActorRun.runOutsideWorkflowSpan`. */
+	runOutsideActorInvocationContext<T>(run: () => T): T;
+	/** Identity of the invocation the caller runs in, usable as a cache key. */
+	currentInvocationScope(): object | undefined;
 	actorName(ctx: ActorContextHandle): string;
 	actorKey(ctx: ActorContextHandle): RuntimeActorKeySegment[];
 	actorRegion(ctx: ActorContextHandle): string;