fix(server-utils): Nest Flue runs on Node under the request that dispatched them

dev-packages/node-integration-tests/suites/tracing/flue/scenario-separate-requests.mjs

@@ -0,0 +1,37 @@
+import * as Sentry from '@sentry/node';
+import { __flueBindAgentModule, init, instrument, useModel } from '@flue/runtime';
+import { start } from '@flue/runtime/node';
+import { fauxAssistantMessage, fauxProvider } from '@earendil-works/pi-ai/providers/faux';
+
+instrument(Sentry.createFlueInstrumentation());
+
+const faux = fauxProvider({
+  provider: 'faux',
+  models: [{ id: 'faux-model', cost: { input: 1, output: 2, cacheRead: 0, cacheWrite: 0 } }],
+});
+faux.setResponses([fauxAssistantMessage('Hello, a.'), fauxAssistantMessage('Hello, b.')]);
+
+function Hello() {
+  useModel('faux/faux-model');
+  return 'You are a helpful assistant.';
+}
+__flueBindAgentModule(Hello, { identity: 'Hello' });
+
+const flue = await start({ agents: [Hello], providers: [faux.provider] });
+
+// Each request gets its own isolation scope and trace, as the http integration gives a real one. The
+// first dispatch starts Flue's claim loop, which keeps request-a's context for every later run.
+for (const name of ['request-a', 'request-b']) {
+  await Sentry.withIsolationScope(() =>
+    Sentry.startNewTrace(() =>
+      Sentry.startSpan({ name, op: 'http.server' }, async () => {
+        const agent = init(Hello, { id: name });
+        const receipt = await agent.dispatch('Say hello.');
+        await agent.read(receipt);
+      }),
+    ),
+  );
+}
+
+await flue[Symbol.asyncDispose]?.();
+await Sentry.flush(2000);

dev-packages/node-integration-tests/suites/tracing/flue/test.ts

@@ -1,3 +1,4 @@
+import type { SerializedStreamedSpanContainer } from '@sentry/core';
 import {
   GEN_AI_AGENT_NAME,
   GEN_AI_CONVERSATION_ID,
@@ -175,4 +176,36 @@ describe.runIf(supports({ min: 22 }))('Flue integration', () => {
     },
     FLUE_DEPENDENCIES,
   );
+
+  createEsmAndCjsTests(
+    __dirname,
+    'scenario-separate-requests.mjs',
+    'instrument.mjs',
+    (createRunner, test, mode) => {
+      if (mode === 'cjs') {
+        return;
+      }
+
+      test('puts each dispatch in the trace of the request that admitted it', async () => {
+        const expectAgentUnder = (requestName: string) => ({
+          span: (container: SerializedStreamedSpanContainer) => {
+            const request = container.items.find(span => span.name === requestName)!;
+            expect(request.is_segment).toBe(true);
+
+            const agents = container.items.filter(span => span.name === 'invoke_agent Hello');
+            expect(agents).toHaveLength(1);
+            expect(agents[0]!.parent_span_id).toBe(request.span_id);
+          },
+        });
+
+        await createRunner()
+          .unordered()
+          .expect(expectAgentUnder('request-a'))
+          .expect(expectAgentUnder('request-b'))
+          .start()
+          .completed();
+      });
+    },
+    FLUE_DEPENDENCIES,
+  );
 });

packages/server-utils/src/ai/flue/index.ts

@@ -1,5 +1,15 @@
-import type { Span } from '@sentry/core';
-import { continueTrace, getActiveSpan, LRUMap, startSpan, withActiveSpan } from '@sentry/core';
+import type { Scope, Span } from '@sentry/core';
+import {
+  continueTrace,
+  getActiveSpan,
+  getDefaultIsolationScope,
+  getIsolationScope,
+  GLOBAL_OBJ,
+  LRUMap,
+  startNewTrace,
+  startSpan,
+  withActiveSpan,
+} from '@sentry/core';
 import {
   GEN_AI_AGENT_NAME,
   GEN_AI_CONVERSATION_ID,
@@ -9,6 +19,7 @@ import {
 import type { GenAiOptions } from '../core/utils';
 import { getGenAiSpanOp, resolveAIRecordingOptions } from '../core/utils';
 import { skipPiAiProviderIntegrations } from '../pi-ai/providers';
+import { withCleanScopes } from '../../utils/withCleanScopes';
 import { FLUE_INSTRUMENTATION_KEY, FLUE_OPERATION, FLUE_ORIGIN, MAX_TRACKED_FLUE_SPANS } from './constants';
 import type { SpanTracker } from './utils';
 import {
@@ -23,6 +34,39 @@ import type { FlueInstrumentation } from './types';
 
 export type FlueOptions = GenAiOptions;
 
+interface Admission {
+  isolationScope: Scope;
+  span: Span | undefined;
+}
+
+/**
+ * Flue's Node coordinator runs every submission from one long-lived claim loop, started by the
+ * first admission and keeping that call's async context for the life of the process. Left alone,
+ * every agent run lands in the trace of whichever request admitted the first one
+ * (https://github.com/withastro/flue/issues/813).
+ */
+function runInAdmissionContext<T>(admission: Admission | undefined, next: () => T): T {
+  if (isCloudflareWorkers()) {
+    return next();
+  }
+  if (!admission) {
+    // Admitted by another process, replayed by idempotency key, or evicted: there is nothing to
+    // restore, but the loop's context still belongs to an unrelated request.
+    return withCleanScopes(getDefaultIsolationScope().clone(), () => startNewTrace(next));
+  }
+  return withCleanScopes(admission.isolationScope, () =>
+    admission.span ? withActiveSpan(admission.span, next) : startNewTrace(next),
+  );
+}
+
+/**
+ * Cloudflare runs each turn from a Durable Object alarm, which already starts a trace of its own, so
+ * there is no leaked context to undo there.
+ */
+function isCloudflareWorkers(): boolean {
+  return GLOBAL_OBJ.navigator?.userAgent === 'Cloudflare-Workers';
+}
+
 /**
  * Build the object to hand to `instrument()` from `@flue/runtime`.
  *
@@ -56,6 +100,8 @@ export function createFlueInstrumentation(options: FlueOptions = {}): FlueInstru
   // Capped, unlike the above: these are keyed off ids that only a matching end observation removes.
   const turnSpans: SpanTracker = new LRUMap(MAX_TRACKED_FLUE_SPANS);
   const toolSpans: SpanTracker = new LRUMap(MAX_TRACKED_FLUE_SPANS);
+  // Capped: a submission admitted here but run by another process never settles here.
+  const admissions = new LRUMap<string, Admission>(MAX_TRACKED_FLUE_SPANS);
 
   return {
     key: FLUE_INSTRUMENTATION_KEY,
@@ -80,9 +126,13 @@ export function createFlueInstrumentation(options: FlueOptions = {}): FlueInstru
 
       // A submission opens with a wrapper operation whose id *is* the submission id; the run it
       // wraps gets its own operation id, and that is the one the observations reference. Spanning
-      // the wrapper too would double-count every agent invocation.
+      // the wrapper too would double-count every agent invocation, but it is the outermost point of
+      // every attempt, so it is where the admitting context goes back in.
       const operationId = operation.operationId;
-      if (!operationId || operationId === ctx.submissionId || agentSpans.has(operationId)) {
+      if (operationId && operationId === ctx.submissionId) {
+        return runInAdmissionContext(admissions.get(operationId), next);
+      }
+      if (!operationId || agentSpans.has(operationId)) {
         return next();
       }
 
@@ -148,6 +198,20 @@ export function createFlueInstrumentation(options: FlueOptions = {}): FlueInstru
       }
 
       switch (observation.type) {
+        // Emitted synchronously from inside the admitting call, so the active context is its own.
+        case 'submission_queued':
+          if (observation.submissionId && !isCloudflareWorkers()) {
+            admissions.set(observation.submissionId, {
+              isolationScope: getIsolationScope(),
+              span: getActiveSpan(),
+            });
+          }
+          return;
+        case 'submission_settled':
+          if (observation.submissionId) {
+            admissions.remove(observation.submissionId);
+          }
+          return;
         case 'turn_start':
           startTurnSpan(observation, turnSpans);
           return;
@@ -179,6 +243,7 @@ export function createFlueInstrumentation(options: FlueOptions = {}): FlueInstru
         span.end();
       }
       agentSpans.clear();
+      admissions.clear();
       turnSpans.clear();
       toolSpans.clear();
     },

packages/server-utils/test/ai/lib/tracing/flue.test.ts

@@ -1,16 +1,22 @@
-import { afterEach, beforeEach, describe, expect, it } from 'vitest';
+import { afterEach, beforeEach, describe, expect, it, onTestFinished, vi } from 'vitest';
 import type { Event, Span } from '@sentry/core';
 import {
   _INTERNAL_clearAiProviderSkips,
   _INTERNAL_shouldSkipAiProviderWrapping,
   _INTERNAL_skipAiProviderWrapping,
+  getIsolationScope,
   getMainCarrier,
+  Scope,
+  setAsyncContextStrategy,
   setCurrentClient,
   spanToStaticSpanJSON,
   startInactiveSpan,
   startSpan,
+  withActiveSpan,
+  withIsolationScope,
 } from '@sentry/core';
 import { ANTHROPIC_AI_INTEGRATION_NAME } from '../../../../src/ai/anthropic-ai/constants';
+import { setAsyncLocalStorageAsyncContextStrategy } from '../../../../src/async-context';
 import { createFlueInstrumentation } from '../../../../src/ai/flue';
 import { MAX_TRACKED_FLUE_SPANS } from '../../../../src/ai/flue/constants';
 import { GOOGLE_GENAI_INTEGRATION_NAME } from '../../../../src/ai/google-genai/constants';
@@ -223,6 +229,114 @@ describe('createFlueInstrumentation', () => {
     });
   });
 
+  // On Node, Flue runs every submission from one claim loop that keeps the async context of the
+  // call that started it, so the run has to be put back where it was admitted.
+  describe('admission context', () => {
+    const SUBMISSION_OP = { type: 'agent', operationId: 'sub_1' };
+
+    function admit(): void {
+      instrumentation.observe({ type: 'submission_queued', submissionId: 'sub_1' }, {});
+    }
+
+    /** Run the submission the way the claim loop does: wrapper first, then the spanned operation. */
+    function runSubmission(): Promise<unknown> {
+      return instrumentation.interceptor(SUBMISSION_OP, AGENT_CTX, () =>
+        instrumentation.interceptor(AGENT_OP, AGENT_CTX, async () => undefined),
+      );
+    }
+
+    it('parents the agent span to the span that admitted the submission, not the one the loop carries', async () => {
+      const admitting = startInactiveSpan({ name: 'POST /agents/hello/b' });
+      withActiveSpan(admitting, admit);
+
+      await startSpan({ name: 'POST /agents/hello/a' }, () => runSubmission());
+
+      const agent = findSpan('invoke_agent Hello');
+      expect(agent?.trace_id).toBe(spanToStaticSpanJSON(admitting).trace_id);
+      expect(agent?.parent_span_id).toBe(admitting.spanContext().spanId);
+    });
+
+    it('restores the isolation scope the submission was admitted in', async () => {
+      // The default stack strategy ignores the scope it is handed.
+      setAsyncLocalStorageAsyncContextStrategy();
+      onTestFinished(() => setAsyncContextStrategy(undefined));
+
+      const admittingScope = new Scope();
+      withIsolationScope(admittingScope, admit);
+
+      let runScope: Scope | undefined;
+      await instrumentation.interceptor(SUBMISSION_OP, AGENT_CTX, async () => {
+        runScope = getIsolationScope();
+      });
+
+      expect(runScope).toBe(admittingScope);
+    });
+
+    it('starts a new trace when nothing was active at admission', async () => {
+      admit();
+
+      let loopTraceId: string | undefined;
+      await startSpan({ name: 'POST /agents/hello/a' }, span => {
+        loopTraceId = spanToStaticSpanJSON(span).trace_id;
+        return runSubmission();
+      });
+
+      const agent = findSpan('invoke_agent Hello');
+      expect(agent?.trace_id).toMatch(/^[0-9a-f]{32}$/);
+      expect(agent?.trace_id).not.toBe(loopTraceId);
+    });
+
+    // Each Cloudflare turn runs from its own Durable Object alarm, which already starts a fresh trace.
+    it('leaves the context alone on Cloudflare Workers', async () => {
+      vi.stubGlobal('navigator', { userAgent: 'Cloudflare-Workers' });
+      onTestFinished(() => {
+        vi.unstubAllGlobals();
+      });
+      withActiveSpan(startInactiveSpan({ name: 'POST /agents/hello/b' }), admit);
+
+      let alarmSpanId: string | undefined;
+      await startSpan({ name: 'alarm' }, span => {
+        alarmSpanId = span.spanContext().spanId;
+        return runSubmission();
+      });
+
+      expect(findSpan('invoke_agent Hello')?.parent_span_id).toBe(alarmSpanId);
+    });
+
+    // A submission admitted by another process, replayed by idempotency key, or already settled has
+    // no record here, but the loop still carries the context of whichever request started it.
+    it('drops the loop context when nothing was recorded for the submission', async () => {
+      withActiveSpan(startInactiveSpan({ name: 'POST /agents/hello/b' }), admit);
+      instrumentation.observe({ type: 'submission_settled', submissionId: 'sub_1' }, {});
+
+      let loopTraceId: string | undefined;
+      await startSpan({ name: 'POST /agents/hello/a' }, span => {
+        loopTraceId = span.spanContext().traceId;
+        return runSubmission();
+      });
+
+      const agent = findSpan('invoke_agent Hello');
+      expect(agent?.parent_span_id).toBeUndefined();
+      expect(agent?.trace_id).toMatch(/^[0-9a-f]{32}$/);
+      expect(agent?.trace_id).not.toBe(loopTraceId);
+    });
+
+    it('continues the replayed traceparent when nothing was recorded for the submission', async () => {
+      const traceId = 'cccccccccccccccccccccccccccccccc';
+      await startSpan({ name: 'POST /agents/hello/a' }, () =>
+        instrumentation.interceptor(SUBMISSION_OP, AGENT_CTX, () =>
+          instrumentation.interceptor(
+            AGENT_OP,
+            { ...AGENT_CTX, traceCarrier: { traceparent: `00-${traceId}-dddddddddddddddd-01` } },
+            async () => undefined,
+          ),
+        ),
+      );
+
+      expect(findSpan('invoke_agent Hello')?.trace_id).toBe(traceId);
+    });
+  });
+
   // The spanned operation carries no `agentName` — only the submission wrapper does, and that opens
   // no span. The name arrives on the observations instead.
   it('names the agent span from the observations', async () => {