diff --git a/dev-packages/node-integration-tests/suites/tracing/flue/scenario-separate-requests.mjs b/dev-packages/node-integration-tests/suites/tracing/flue/scenario-separate-requests.mjs new file mode 100644 index 000000000000..13a54ffc9d25 --- /dev/null +++ b/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); diff --git a/dev-packages/node-integration-tests/suites/tracing/flue/test.ts b/dev-packages/node-integration-tests/suites/tracing/flue/test.ts index 5d3b6fe6cc0e..b750ee653e4f 100644 --- a/dev-packages/node-integration-tests/suites/tracing/flue/test.ts +++ b/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, + ); }); diff --git a/packages/server-utils/src/ai/flue/index.ts b/packages/server-utils/src/ai/flue/index.ts index 44c9ade4667c..d4856f95f2ac 100644 --- a/packages/server-utils/src/ai/flue/index.ts +++ b/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(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(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(); }, diff --git a/packages/server-utils/test/ai/lib/tracing/flue.test.ts b/packages/server-utils/test/ai/lib/tracing/flue.test.ts index bb3712d6ac32..e39ed3fd3bfb 100644 --- a/packages/server-utils/test/ai/lib/tracing/flue.test.ts +++ b/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 { + 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 () => {