feat(server-utils): Restore caller context for callback tracing channels

packages/server-utils/src/integrations/tracing-channel/lru-memoizer.ts

@@ -1,14 +1,15 @@
 import * as diagnosticsChannel from 'node:diagnostics_channel';
 import type { IntegrationFn } from '@sentry/core';
-import { debug, defineIntegration, getCurrentScope, withScope } from '@sentry/core';
+import { debug, defineIntegration, waitForTracingChannelBinding } from '@sentry/core';
 import { DEBUG_BUILD } from '../../debug-build';
 import { CHANNELS } from '../../orchestrion/channels';
+import { bindTracingChannelToSpan } from '../../tracing-channel';
 
 // Same name as the OTel integration by design — when enabled, the OTel
 // 'LruMemoizer' integration is omitted from the default set.
 const INTEGRATION_NAME = 'LruMemoizer' as const;
 
-interface LruMemoizerChannelContext {
+interface LruMemoizerLoadContext {
   arguments: unknown[];
 }
 
@@ -22,35 +23,13 @@ const _lruMemoizerChannelIntegration = (() => {
       }
 
       DEBUG_BUILD && debug.log(`[orchestrion:lru-memoizer] subscribing to channel "${CHANNELS.LRU_MEMOIZER_LOAD}"`);
-      const lruMemoizerCh = diagnosticsChannel.tracingChannel(CHANNELS.LRU_MEMOIZER_LOAD);
 
-      lruMemoizerCh.subscribe({
-        start(rawCtx) {
-          const ctx = rawCtx as LruMemoizerChannelContext;
-          if (ctx.arguments.length === 0) {
-            return;
-          }
-
-          // Capture the scope while we're still synchronously inside the memoized call.
-          // lru-memoizer queues the callback and fires it later via setImmediate, where the
-          // active scope no longer reflects the caller's context.
-          const scope = getCurrentScope();
-          const cbIdx = ctx.arguments.length - 1;
-          const orchestrionWrappedCb = ctx.arguments[cbIdx];
-
-          if (typeof orchestrionWrappedCb !== 'function') {
-            return;
-          }
-
-          const wrapped = orchestrionWrappedCb as (...a: unknown[]) => unknown;
-          ctx.arguments[cbIdx] = function (this: unknown, ...args: unknown[]): unknown {
-            return withScope(scope, () => wrapped.apply(this, args));
-          };
-        },
-        end() {},
-        asyncStart() {},
-        asyncEnd() {},
-        error() {},
+      waitForTracingChannelBinding(() => {
+        bindTracingChannelToSpan(
+          diagnosticsChannel.tracingChannel<LruMemoizerLoadContext>(CHANNELS.LRU_MEMOIZER_LOAD),
+          // We only want the helper's caller-context restore for the callback lru-memoizer fires from a detached `setImmediate`.
+          () => undefined,
+        );
       });
     },
   };

packages/server-utils/src/tracing-channel.ts

@@ -6,7 +6,15 @@ import { DEBUG_BUILD } from './debug-build';
 import { ERROR_TYPE } from '@sentry/conventions/attributes';
 
 export type TracingChannelPayloadWithSpan<TData extends object> = TData & {
+  /**
+   * The current active span for the traced call.
+   */
   _sentrySpan?: Span;
+
+  /**
+   * The context's active store value, used to restore the context for asyncStart continuations for callback-based tracing.
+   */
+  _sentryCallerStore?: unknown;
 };
 
 /*
@@ -158,21 +166,33 @@ function bindSpanToChannelStore<TData extends object>(
   // 3. Read: inside the op, Sentry's scope machinery calls getScopes() → asyncStorage.getStore() on that same ALS, so getCurrentScope/getIsolationScope/getActiveSpan resolve to the scope carrying our span.
   // 4. Nest: any child span started in the traced op parents to that active span.
   channel.start.bindStore(asyncLocalStorage, (data: TracingChannelPayloadWithSpan<TData>) => {
+    // Stash the caller's store before we swap in the span store, so `asyncStart` can restore it for
+    // callback-style channels (see `_sentryCallerStore`).
+    data._sentryCallerStore = asyncLocalStorage.getStore();
+
     const span = getSpan(data);
     if (!span) {
       // Leave the active context untouched so nested operations keep parenting to the enclosing span.
-      return asyncLocalStorage.getStore() as TData;
+      return data._sentryCallerStore as TData;
     }
     data._sentrySpan = span;
 
     return binding.getStoreWithActiveSpan(span) as TData;
   });
 
+  // Restore the caller's context for the async continuation. Only callback-style channels `runStores`
+  // `asyncStart` (so the callback runs inside this store). promise channels `publish` it, leaving this
+  // inert, their continuation already inherits the caller's context natively.
+  channel.asyncStart.bindStore(asyncLocalStorage, (data: TracingChannelPayloadWithSpan<TData>) => {
+    return data._sentryCallerStore as TData;
+  });
+
   return {
     channel,
     unbind: () => {
-      // Removes the store
+      // Removes the stores
       channel.start.unbindStore(asyncLocalStorage);
+      channel.asyncStart.unbindStore(asyncLocalStorage);
     },
   };
 }

packages/server-utils/test/tracing-channel.test.ts

@@ -175,6 +175,81 @@ describe('bindTracingChannelToSpan', () => {
     expect(childParentSpanId).toBe(parent.spanContext().spanId);
   });
 
+  it('restores the caller context in a callback dispatched from a detached context (asyncStart rebind)', async () => {
+    installTestAsyncContextStrategy();
+    initTestClient();
+
+    let channelSpanId: string | undefined;
+    const { channel } = bindTracingChannelToSpan(
+      tracingChannel<{ operation: string }>('test:asyncStart:caller-context'),
+      () => {
+        const span = startInactiveSpan({ name: 'channel-span' });
+        channelSpanId = span.spanContext().spanId;
+        return span;
+      },
+    );
+
+    let enclosingSpanId: string | undefined;
+    let childParentSpanId: string | undefined;
+
+    await new Promise<void>(done => {
+      startSpan({ forceTransaction: true, name: 'enclosing-span' }, enclosing => {
+        enclosingSpanId = enclosing.spanContext().spanId;
+        channel.traceCallback(
+          (cb: (err: Error | null, result?: string) => void) => {
+            // Fire the callback after the enclosing scope has exited, so it runs in a detached
+            // async context — the asyncStart rebind is the only thing that can restore the caller's.
+            setTimeout(() => cb(null, 'ok'), 1);
+          },
+          0,
+          { operation: 'read' },
+          undefined,
+          () => {
+            startSpan({ name: 'child-span' }, child => {
+              childParentSpanId = spanToJSON(child).parent_span_id;
+            });
+            done();
+          },
+        );
+      });
+    });
+
+    // A span started inside the callback parents to the caller (the enclosing span), not to the
+    // channel span — matching how a promise's `.then` continuation behaves.
+    expect(childParentSpanId).toBe(enclosingSpanId);
+    expect(childParentSpanId).not.toBe(channelSpanId);
+  });
+
+  it('does not leak an unrelated active store into the callback when the caller had none', () => {
+    installTestAsyncContextStrategy();
+    initTestClient();
+
+    const { channel } = bindTracingChannelToSpan(tracingChannel<{ operation: string }>('test:asyncStart:no-leak'), () =>
+      startInactiveSpan({ name: 'channel-span' }),
+    );
+
+    // Caller issues the op with no active context, so the caller store is captured as `undefined`.
+    const ctx = { operation: 'read' };
+    channel.start.runStores(ctx, () => undefined);
+
+    let otherRequestSpanId: string | undefined;
+    let childParentSpanId: string | undefined;
+
+    // The callback fires later, dispatched from *another* request's active context.
+    startSpan({ forceTransaction: true, name: 'other-request' }, other => {
+      otherRequestSpanId = other.spanContext().spanId;
+      channel.asyncStart.runStores(ctx, () => {
+        startSpan({ name: 'child-span' }, child => {
+          childParentSpanId = spanToJSON(child).parent_span_id;
+        });
+      });
+    });
+
+    // The caller had no context, so the callback must restore to none — not adopt the other request's.
+    expect(childParentSpanId).toBeUndefined();
+    expect(childParentSpanId).not.toBe(otherRequestSpanId);
+  });
+
   describe('auto lifecycle ending strategy', () => {
     // Returns a channel whose span we can observe, plus spies for `span.end` and `captureException`.
     function setup(name: string): {