rivetkit 2.3.18 → 2.3.19

package/dist/tsup/mod.js

@@ -1,6 +1,3 @@
-import {
-  decodeWorkflowHistoryTransport
-} from "./chunk-QIFZELUQ.js";
 import {
   ACTOR_CONTEXT_INTERNAL_SYMBOL,
   ActorConfigSchema,
@@ -16,7 +13,10 @@ import {
   getRunInspectorKind,
   getRunMetadata,
   hasRunInspectorConfig
-} from "./chunk-6W5VGLFT.js";
+} from "./chunk-EB77NZL4.js";
+import {
+  decodeWorkflowHistoryTransport
+} from "./chunk-4SC2NSAV.js";
 import {
   ALLOWED_PUBLIC_HEADERS,
   HEADER_CONN_PARAMS,
@@ -26,20 +26,22 @@ import {
   createClientWithDriver,
   getDatacenters,
   isResponseLike,
+  readActiveTraceHeaders,
+  runWithActorInvocationSpan,
   tryParseEndpoint,
   updateRunnerConfig
-} from "./chunk-PXKNBNVU.js";
+} from "./chunk-BKX7UDTA.js";
 import {
   makePrefixedKey,
   removePrefixFromKey
 } from "./chunk-JI6GZ2C2.js";
+import "./chunk-X7RMJQGY.js";
 import {
   bindNativeStateTransactionContext,
   db,
   registerNativeStateTransactionOpener
-} from "./chunk-TNSLNSQH.js";
+} from "./chunk-WME6KOXW.js";
 import "./chunk-JTHHCZCZ.js";
-import "./chunk-LAITPJRA.js";
 import {
   LogLevelSchema,
   VERSION,
@@ -73,19 +75,20 @@ import {
   parsePortEnv,
   stringifyError,
   toUint8Array
-} from "./chunk-7EPLADKM.js";
+} from "./chunk-DN6MD3KK.js";
 import {
   INTERNAL_ERROR_CODE,
   RivetError,
   UserError,
   decodeBridgeRivetError,
   encodeBridgeRivetError,
+  encodeErrorForBridge,
   forbiddenError,
   isActorAbortedError,
   isRivetErrorLike,
   toRivetError,
   unsupportedFeature
-} from "./chunk-OUQUIBVW.js";
+} from "./chunk-P2OFXH4G.js";
 
 // src/actor/log.ts
 function loggerWithoutContext() {
@@ -1545,20 +1548,43 @@ function toNapiQueueMessage(message) {
 var NapiCoreRuntime = class {
   kind = "napi";
   #bindings;
-  #sql = /* @__PURE__ */ new WeakMap();
-  constructor(bindings) {
+  #sqlOpened = /* @__PURE__ */ new WeakSet();
+  #invocationContext;
+  constructor(bindings, invocationContext) {
     this.#bindings = bindings;
+    this.#invocationContext = invocationContext;
   }
-  #actorSql(ctx) {
-    const nativeCtx = asNativeActorContext(ctx);
-    let database = this.#sql.get(nativeCtx);
-    if (!database) {
-      database = nativeCtx.sql();
-      this.#sql.set(nativeCtx, database);
+  #actorContextForOperation(owner) {
+    const ownerCtx = asNativeActorContext(owner);
+    const active = this.#invocationContext.getStore();
+    if (!(active == null ? void 0 : active.sameActorInstance(ownerCtx))) {
+      return ownerCtx;
     }
-    return database;
+    const applicationSpan = readActiveTraceHeaders();
+    if (!applicationSpan) {
+      return active;
+    }
+    return active.withApplicationSpan(
+      applicationSpan.traceparent,
+      applicationSpan.tracestate ?? null
+    );
+  }
+  #actorSql(ctx) {
+    this.#sqlOpened.add(asNativeActorContext(ctx));
+    return this.#actorContextForOperation(ctx).sql();
   }
   createRegistry() {
+    this.#bindings.setTelemetryLogSink(
+      (event2) => {
+        const fields = { otelEvent: event2.name };
+        const message = event2.message || event2.name;
+        if (event2.level === "error") {
+          logger2().error(fields, message);
+        } else {
+          logger2().warn(fields, message);
+        }
+      }
+    );
     return asRegistryHandle(new this.#bindings.CoreRegistry());
   }
   registerActor(registry, name, factory) {
@@ -1573,6 +1599,9 @@ var NapiCoreRuntime = class {
   async shutdownRegistry(registry) {
     await asNativeRegistry(registry).shutdown();
   }
+  async shutdownTelemetry() {
+    await this.#bindings.shutdownTelemetry();
+  }
   async registryActorStopThresholdMs(registry) {
     return await asNativeRegistry(registry).actorStopThresholdMs() ?? void 0;
   }
@@ -1726,6 +1755,57 @@ var NapiCoreRuntime = class {
   actorId(ctx) {
     return asNativeActorContext(ctx).actorId();
   }
+  runWithActorInvocationContext(ctx, run) {
+    return this.#runAs(asNativeActorContext(ctx), run);
+  }
+  #runAs(nativeCtx, run) {
+    var _a;
+    const span = (_a = nativeCtx.invocationTraceContext()) == null ? void 0 : _a.span;
+    return this.#invocationContext.run(
+      nativeCtx,
+      () => runWithActorInvocationSpan(span, run)
+    );
+  }
+  actorInvocationTraceContext(ctx) {
+    return this.#actorContextForOperation(ctx).invocationTraceContext() ?? void 0;
+  }
+  startCallSpan(ctx, actorName, actionName) {
+    const call = this.#actorContextForOperation(ctx).startCallSpan(
+      actorName,
+      actionName
+    );
+    if (!call) return void 0;
+    return {
+      span: call.spanContext() ?? void 0,
+      finish: (error) => call.finish(error)
+    };
+  }
+  async startWorkflowSpan(ctx) {
+    const span = await asNativeActorContext(ctx).startWorkflowSpan();
+    const spanCtx = span.ctx();
+    return {
+      run: (body) => this.#runAs(spanCtx, body),
+      startStep: (name, attempt) => {
+        const step = spanCtx.startWorkflowStepSpan(name, attempt);
+        if (!step) return void 0;
+        const stepCtx = step.ctx();
+        return {
+          run: (body) => this.#runAs(stepCtx, body),
+          finish: (outcome, error) => step.finish(
+            outcome,
+            error === void 0 ? void 0 : encodeErrorForBridge(error)
+          )
+        };
+      },
+      finish: (outcome) => span.finish(outcome)
+    };
+  }
+  runOutsideActorInvocationContext(run) {
+    return this.#invocationContext.exit(run);
+  }
+  currentInvocationScope() {
+    return this.#invocationContext.getStore();
+  }
   actorName(ctx) {
     return asNativeActorContext(ctx).name();
   }
@@ -1757,7 +1837,7 @@ var NapiCoreRuntime = class {
     asNativeActorContext(ctx).broadcast(name, toNapiBuffer(args));
   }
   actorWaitUntil(ctx, promise) {
-    asNativeActorContext(ctx).waitUntil(promise);
+    this.#actorContextForOperation(ctx).waitUntil(promise);
   }
   async actorWaitForTrackedShutdownWork(ctx) {
     return await asNativeActorContext(ctx).waitForTrackedShutdownWork();
@@ -1858,7 +1938,7 @@ var NapiCoreRuntime = class {
     await asNativeSqlTransaction(transaction).rollback();
   }
   async actorBeginStateTransaction(ctx, timeoutMs) {
-    return await asNativeActorContext(ctx).beginStateTransaction(
+    return await this.#actorContextForOperation(ctx).beginStateTransaction(
       timeoutMs
     );
   }
@@ -1896,23 +1976,21 @@ var NapiCoreRuntime = class {
   }
   async actorSqlClose(ctx) {
     const nativeCtx = asNativeActorContext(ctx);
-    const database = this.#sql.get(nativeCtx);
-    if (!database) {
+    if (!this.#sqlOpened.delete(nativeCtx)) {
       return;
     }
-    this.#sql.delete(nativeCtx);
-    await database.close();
+    await nativeCtx.sql().close();
   }
   async actorRuntimeSocketProvision(ctx) {
     return await asNativeActorContext(ctx).provisionActorRuntimeSocket();
   }
   async actorQueueSend(ctx, name, body) {
     return toNapiQueueMessage(
-      await asNativeActorContext(ctx).queue().send(name, toNapiBuffer(body))
+      await this.#actorContextForOperation(ctx).queue().send(name, toNapiBuffer(body))
     );
   }
   async actorQueueNextBatch(ctx, options, signal) {
-    const messages = await asNativeActorContext(ctx).queue().nextBatch(
+    const messages = await this.#actorContextForOperation(ctx).queue().nextBatch(
       options,
       signal ? asNativeCancellationToken(signal) : signal
     );
@@ -1920,7 +1998,7 @@ var NapiCoreRuntime = class {
   }
   async actorQueueWaitForNames(ctx, names, options, signal) {
     return toNapiQueueMessage(
-      await asNativeActorContext(ctx).queue().waitForNames(
+      await this.#actorContextForOperation(ctx).queue().waitForNames(
         names,
         options,
         signal ? asNativeCancellationToken(signal) : signal
@@ -1942,7 +2020,7 @@ var NapiCoreRuntime = class {
     );
   }
   async actorQueueEnqueueAndWait(ctx, name, body, options, signal) {
-    return await asNativeActorContext(ctx).queue().enqueueAndWait(
+    return await this.#actorContextForOperation(ctx).queue().enqueueAndWait(
       name,
       toNapiBuffer(body),
       options,
@@ -1962,22 +2040,22 @@ var NapiCoreRuntime = class {
     await asNativeActorContext(ctx).queue().reset();
   }
   async actorScheduleAfter(ctx, durationMs, actionName, args) {
-    return await asNativeActorContext(ctx).schedule().after(durationMs, actionName, toNapiBuffer(args));
+    return await this.#actorContextForOperation(ctx).schedule().after(durationMs, actionName, toNapiBuffer(args));
   }
   async actorScheduleAt(ctx, timestampMs, actionName, args) {
-    return await asNativeActorContext(ctx).schedule().at(timestampMs, actionName, toNapiBuffer(args));
+    return await this.#actorContextForOperation(ctx).schedule().at(timestampMs, actionName, toNapiBuffer(args));
   }
   async actorScheduleCancel(ctx, id) {
-    return await asNativeActorContext(ctx).schedule().cancel(id);
+    return await this.#actorContextForOperation(ctx).schedule().cancel(id);
   }
   async actorScheduleGet(ctx, id) {
-    return await asNativeActorContext(ctx).schedule().get(id) ?? void 0;
+    return await this.#actorContextForOperation(ctx).schedule().get(id) ?? void 0;
   }
   async actorScheduleList(ctx) {
-    return await asNativeActorContext(ctx).schedule().list();
+    return await this.#actorContextForOperation(ctx).schedule().list();
   }
   async actorCronSet(ctx, name, expression, timezone, actionName, args, maxHistory) {
-    await asNativeActorContext(ctx).schedule().cronSet(
+    await this.#actorContextForOperation(ctx).schedule().cronSet(
       name,
       expression,
       timezone,
@@ -1987,7 +2065,7 @@ var NapiCoreRuntime = class {
     );
   }
   async actorCronEvery(ctx, name, intervalMs, actionName, args, maxHistory) {
-    await asNativeActorContext(ctx).schedule().cronEvery(
+    await this.#actorContextForOperation(ctx).schedule().cronEvery(
       name,
       intervalMs,
       actionName,
@@ -1996,16 +2074,16 @@ var NapiCoreRuntime = class {
     );
   }
   async actorCronGet(ctx, name) {
-    return await asNativeActorContext(ctx).schedule().cronGet(name) ?? void 0;
+    return await this.#actorContextForOperation(ctx).schedule().cronGet(name) ?? void 0;
   }
   async actorCronList(ctx) {
-    return await asNativeActorContext(ctx).schedule().cronList();
+    return await this.#actorContextForOperation(ctx).schedule().cronList();
   }
   async actorCronDelete(ctx, name) {
-    return await asNativeActorContext(ctx).schedule().cronDelete(name);
+    return await this.#actorContextForOperation(ctx).schedule().cronDelete(name);
   }
   async actorCronHistory(ctx, name, limit) {
-    return await asNativeActorContext(ctx).schedule().cronHistory(name, limit);
+    return await this.#actorContextForOperation(ctx).schedule().cronHistory(name, limit);
   }
   connId(conn) {
     return asNativeConn(conn).id();
@@ -2039,10 +2117,13 @@ var NapiCoreRuntime = class {
   }
 };
 async function loadNapiRuntime() {
-  const bindings = await import(["@rivetkit", "rivetkit-napi"].join("/"));
+  const [{ AsyncLocalStorage }, bindings] = await Promise.all([
+    import("async_hooks"),
+    import(["@rivetkit", "rivetkit-napi"].join("/"))
+  ]);
   return {
     bindings,
-    runtime: new NapiCoreRuntime(bindings)
+    runtime: new NapiCoreRuntime(bindings, new AsyncLocalStorage())
   };
 }
 
@@ -2742,6 +2823,29 @@ var WasmCoreRuntime = class {
   actorId(ctx) {
     return callHandle(asWasmActorContext(ctx), "actorId");
   }
+  runWithActorInvocationContext(_ctx, run) {
+    return run();
+  }
+  actorInvocationTraceContext(_ctx) {
+    return void 0;
+  }
+  startCallSpan(_ctx, _actorName, _actionName) {
+    return void 0;
+  }
+  async startWorkflowSpan(_ctx) {
+    return {
+      run: (body) => body(),
+      startStep: () => void 0,
+      finish: async () => {
+      }
+    };
+  }
+  runOutsideActorInvocationContext(run) {
+    return run();
+  }
+  currentInvocationScope() {
+    return void 0;
+  }
   actorName(ctx) {
     return callHandle(asWasmActorContext(ctx), "name");
   }
@@ -5055,6 +5159,7 @@ var NativeConnectionMap = class {
   }
   [Symbol.toStringTag] = "NativeConnectionMap";
 };
+var NO_INVOCATION_SCOPE = {};
 var ActorContextHandleAdapter = class {
   #runtime;
   #ctx;
@@ -5069,6 +5174,7 @@ var ActorContextHandleAdapter = class {
   #db;
   #dispatchCancelToken;
   #kv;
+  #logByInvocationScope = /* @__PURE__ */ new WeakMap();
   #queue;
   #request;
   #schedule;
@@ -5211,7 +5317,11 @@ var ActorContextHandleAdapter = class {
           await callNative(
             () => this.#runtime.actorSetRunWakeAt(this.#ctx, timestamp)
           );
-        }
+        },
+        startWorkflowSpan: () => callNative(
+          () => this.#runtime.startWorkflowSpan(this.#ctx)
+        ),
+        runOutsideWorkflowSpan: (body) => this.#runtime.runOutsideActorInvocationContext(body)
       };
     }
     return this.#run;
@@ -5256,8 +5366,29 @@ var ActorContextHandleAdapter = class {
     }
     return this.#connMap;
   }
+  #invocationTraceContext() {
+    return callNativeSync(
+      () => this.#runtime.actorInvocationTraceContext(this.#ctx)
+    );
+  }
+  /** Cached per invocation, because the `run` context outlives workflow runs and steps. */
   get log() {
-    return logger2();
+    const scope = this.#runtime.currentInvocationScope() ?? NO_INVOCATION_SCOPE;
+    const cached = this.#logByInvocationScope.get(scope);
+    if (cached) return cached;
+    const invocation = this.#invocationTraceContext();
+    const log = logger2().child({
+      actorId: this.actorId,
+      actorName: this.name,
+      actorKey: this.key,
+      ...(invocation == null ? void 0 : invocation.rayId) && { rayId: invocation.rayId },
+      ...(invocation == null ? void 0 : invocation.span) && {
+        traceId: invocation.span.traceId,
+        spanId: invocation.span.spanId
+      }
+    });
+    this.#logByInvocationScope.set(scope, log);
+    return log;
   }
   get abortSignal() {
     if (!this.#abortSignal) {
@@ -6045,11 +6176,19 @@ function buildNativeFactory(runtime, registryConfig, definition) {
     events: config.events,
     queues: config.queues
   };
-  const createClient = () => createClientWithDriver(
+  const createClient = (ctx) => 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 = getRunInspectorKind(config.run) === "workflow" ? new RunHandlerCoordinator(config.run) : void 0;
@@ -6072,7 +6211,7 @@ function buildNativeFactory(runtime, registryConfig, definition) {
   const makeActorCtx = (ctx, request, cancelToken) => new ActorContextHandleAdapter(
     runtime,
     ctx,
-    createClient,
+    () => createClient(ctx),
     schemaConfig,
     databaseProvider,
     request,
@@ -6085,7 +6224,7 @@ function buildNativeFactory(runtime, registryConfig, definition) {
     runtime,
     ctx,
     conn,
-    createClient,
+    () => createClient(ctx),
     schemaConfig,
     databaseProvider,
     request,
@@ -6798,111 +6937,125 @@ function buildNativeFactory(runtime, registryConfig, definition) {
       async (error, payload) => {
         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;
-          }
-          const requestAbortController = new AbortController();
-          const handlerRequest = buildNativeHttpRequest({
-            ...request,
-            abortController: requestAbortController
-          });
-          const rawConnParams = handlerRequest.headers.get(HEADER_CONN_PARAMS);
-          let requestCtx;
-          let conn;
-          let removeRequestAbortListener;
-          let cleanupDeferredToBody = false;
-          let cleanedUp = false;
-          const cleanupRequest = async () => {
-            if (cleanedUp) return;
-            cleanedUp = true;
-            removeRequestAbortListener == null ? void 0 : removeRequestAbortListener();
-            try {
-              await (requestCtx == null ? void 0 : requestCtx.dispose());
-            } finally {
-              if (conn) {
-                await runtime.connDisconnect(conn);
-              }
-            }
-          };
-          try {
-            const connParams = validateConnParams(
-              schemaConfig.connParamsSchema,
-              rawConnParams ? JSON.parse(rawConnParams) : void 0
-            );
-            conn = await callNative(
-              () => runtime.actorConnectConn(
+          return await runtime.runWithActorInvocationContext(
+            ctx,
+            async () => {
+              const inspectorResponse = await maybeHandleNativeInspectorRequest(
                 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) => {
-                logger2().error({
-                  msg: "failed to clean up native streaming http request",
-                  error: cleanupError
-                });
+              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;
+              }
+              const requestAbortController = new AbortController();
+              const handlerRequest = buildNativeHttpRequest({
+                ...request,
+                abortController: requestAbortController
               });
-            }
-            return conversion.response;
-          } finally {
-            try {
-              await cancelNativeHttpRequestBody(
-                request.bodyStream
-              );
-            } finally {
-              if (!cleanupDeferredToBody) {
-                await cleanupRequest();
+              const rawConnParams = handlerRequest.headers.get(HEADER_CONN_PARAMS);
+              let requestCtx;
+              let conn;
+              let removeRequestAbortListener;
+              let cleanupDeferredToBody = false;
+              let cleanedUp = false;
+              const cleanupRequest = async () => {
+                if (cleanedUp) return;
+                cleanedUp = true;
+                removeRequestAbortListener == null ? void 0 : removeRequestAbortListener();
+                try {
+                  await (requestCtx == null ? void 0 : requestCtx.dispose());
+                } finally {
+                  if (conn) {
+                    await runtime.connDisconnect(conn);
+                  }
+                }
+              };
+              try {
+                const connParams = validateConnParams(
+                  schemaConfig.connParamsSchema,
+                  rawConnParams ? JSON.parse(rawConnParams) : void 0
+                );
+                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) => {
+                    logger2().error({
+                      msg: "failed to clean up native streaming http request",
+                      error: cleanupError
+                    });
+                  });
+                }
+                return conversion.response;
+              } finally {
+                try {
+                  await cancelNativeHttpRequestBody(
+                    request.bodyStream
+                  );
+                } finally {
+                  if (!cleanupDeferredToBody) {
+                    await cleanupRequest();
+                  }
+                }
               }
             }
-          }
+          );
         } catch (error2) {
           logger2().error({
             msg: "native onRequest failed",
@@ -6990,21 +7143,27 @@ function buildNativeFactory(runtime, registryConfig, definition) {
           async (error, payload) => {
             const { ctx, conn, args, scheduledFire, cancelToken } = unwrapTsfnPayload(error, payload);
             const actorCtx = conn != null ? makeConnCtx(ctx, conn, void 0, cancelToken) : makeActorCtx(ctx, void 0, cancelToken);
-            try {
-              return encodeValue(
-                await handler(
-                  actorCtx,
-                  ...validateActionArgs(
-                    schemaConfig.actionInputSchemas,
-                    name,
-                    decodeArgs(args)
-                  ),
-                  ...scheduledFire ? [scheduledFire] : []
-                )
-              );
-            } finally {
-              await actorCtx.dispose();
-            }
+            const runAction = async () => {
+              try {
+                return encodeValue(
+                  await handler(
+                    actorCtx,
+                    ...validateActionArgs(
+                      schemaConfig.actionInputSchemas,
+                      name,
+                      decodeArgs(args)
+                    ),
+                    ...scheduledFire ? [scheduledFire] : []
+                  )
+                );
+              } finally {
+                await actorCtx.dispose();
+              }
+            };
+            return await runtime.runWithActorInvocationContext(
+              ctx,
+              runAction
+            );
           }
         )
       ])
@@ -7026,7 +7185,7 @@ function buildNativeFactory(runtime, registryConfig, definition) {
           runtime,
           ctx,
           conn,
-          createClient,
+          () => createClient(ctx),
           schemaConfig,
           databaseProvider,
           jsRequest,
@@ -7035,43 +7194,48 @@ function buildNativeFactory(runtime, registryConfig, definition) {
           cancelToken,
           run !== void 0
         );
-        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 === void 0 || timeoutMs === null ? void 0 : Number(timeoutMs)
-                }
+              if (!schemaConfig.queues || !hasSchemaConfigKey(schemaConfig.queues, name)) {
+                return { status: "completed" };
+              }
+              const canPublish = getQueueCanPublish(
+                schemaConfig.queues,
+                name
               );
-              return {
-                status: "completed",
-                response: response === void 0 ? void 0 : encodeValue(response)
-              };
-            } catch (error2) {
-              if (error2.group === "queue" && error2.code === "timed_out") {
-                return { status: "timedOut" };
+              if (canPublish && !await canPublish(actorCtx)) {
+                throw forbiddenError();
               }
-              throw error2;
+              const decodedBody = decodeValue(body);
+              if (wait) {
+                try {
+                  const response = await actorCtx.queue.enqueueAndWait(
+                    name,
+                    decodedBody,
+                    {
+                      timeout: timeoutMs === void 0 || timeoutMs === null ? void 0 : Number(timeoutMs)
+                    }
+                  );
+                  return {
+                    status: "completed",
+                    response: response === void 0 ? void 0 : encodeValue(response)
+                  };
+                } catch (error2) {
+                  if (error2.group === "queue" && error2.code === "timed_out") {
+                    return { status: "timedOut" };
+                  }
+                  throw error2;
+                }
+              }
+              await actorCtx.queue.send(name, decodedBody);
+              return { status: "completed" };
+            } finally {
+              await actorCtx.dispose();
             }
           }
-          await actorCtx.queue.send(name, decodedBody);
-          return { status: "completed" };
-        } finally {
-          await actorCtx.dispose();
-        }
+        );
       }
     ),
     serializeState: wrapNativeCallback(
@@ -7721,6 +7885,7 @@ var Registry = class {
     var _a;
     const modeAPromise = this.#runtimeServeConfiguredPromise;
     const modeBPromise = this.#runtimeServerlessPromise;
+    const runtimes = /* @__PURE__ */ new Set();
     const gracePeriodMs = ((_a = config.shutdown) == null ? void 0 : _a.gracePeriodMs) ?? await this.#actorStopThresholdMs(modeAPromise ?? modeBPromise) ?? 30 * 60 * 1e3;
     const drain = async () => {
       const registries = [];
@@ -7729,6 +7894,7 @@ var Registry = class {
           (async () => {
             try {
               const { runtime, registry } = await modeAPromise;
+              runtimes.add(runtime);
               await runtime.shutdownRegistry(registry);
             } catch (err) {
               logger2().warn(
@@ -7744,6 +7910,7 @@ var Registry = class {
           (async () => {
             try {
               const { runtime, registry } = await modeBPromise;
+              runtimes.add(runtime);
               await runtime.shutdownRegistry(registry);
             } catch (err) {
               logger2().warn(
@@ -7772,6 +7939,12 @@ var Registry = class {
         }
       )
     ]);
+    await Promise.all(
+      [...runtimes].map((runtime) => {
+        var _a2;
+        return (_a2 = runtime.shutdownTelemetry) == null ? void 0 : _a2.call(runtime);
+      })
+    );
   }
   async #actorStopThresholdMs(configuredRegistryPromise) {
     var _a;