diff --git a/core/package.json b/core/package.json index 4dde99c49..2a5c455c4 100644 --- a/core/package.json +++ b/core/package.json @@ -72,6 +72,7 @@ }, "devDependencies": { "@mikro-orm/sqlite": "^6.6.6", + "@opentelemetry/context-async-hooks": "^2.1.0", "@types/adm-zip": "^0.5.8", "@types/express": "^4.17.25", "@types/lodash-es": "^4.17.12", diff --git a/core/src/telemetry/tracing.ts b/core/src/telemetry/tracing.ts index 4d97440bc..56abc84bd 100644 --- a/core/src/telemetry/tracing.ts +++ b/core/src/telemetry/tracing.ts @@ -35,6 +35,13 @@ const GEN_AI_TOOL_DESCRIPTION = 'gen_ai.tool.description'; const GEN_AI_TOOL_NAME = 'gen_ai.tool.name'; const GEN_AI_TOOL_TYPE = 'gen_ai.tool.type'; +const ADK_WORKFLOW_NAME = 'adk.workflow.name'; +const ADK_NODE_PATH = 'adk.node.path'; +const ADK_NODE_RUN_ID = 'adk.node.run_id'; +const ADK_NODE_ATTEMPT = 'adk.node.attempt'; +const ADK_NODE_STATUS = 'adk.node.status'; +const ADK_NODE_INTERRUPT_COUNT = 'adk.node.interrupt_count'; + export const tracer = trace.getTracer('gcp.vertex.agent', version); /** @@ -90,6 +97,55 @@ export function traceAgentInvocation({ }); } +export interface TraceWorkflowInvocationParams { + workflowName: string; + nodePath: string; +} + +export function traceWorkflowInvocation({ + workflowName, + nodePath, +}: TraceWorkflowInvocationParams): void { + const span = trace.getActiveSpan(); + if (!span) return; + + span.setAttributes({ + [GEN_AI_OPERATION_NAME]: 'invoke_workflow', + [ADK_WORKFLOW_NAME]: workflowName, + [ADK_NODE_PATH]: nodePath, + }); +} + +export type NodeExecutionStatus = 'completed' | 'waiting' | 'failed'; + +export interface TraceNodeExecutionParams { + nodePath: string; + runId: string; + attempt: number; + status: NodeExecutionStatus; + interruptCount: number; +} + +export function traceNodeExecution({ + nodePath, + runId, + attempt, + status, + interruptCount, +}: TraceNodeExecutionParams): void { + const span = trace.getActiveSpan(); + if (!span) return; + + span.setAttributes({ + [GEN_AI_OPERATION_NAME]: 'execute_node', + [ADK_NODE_PATH]: nodePath, + [ADK_NODE_RUN_ID]: runId, + [ADK_NODE_ATTEMPT]: attempt, + [ADK_NODE_STATUS]: status, + [ADK_NODE_INTERRUPT_COUNT]: interruptCount, + }); +} + export interface TraceToolCallParams { tool: BaseTool; args: Record; diff --git a/core/src/workflow/node_runner.ts b/core/src/workflow/node_runner.ts index 77c18ed57..956b89f7c 100644 --- a/core/src/workflow/node_runner.ts +++ b/core/src/workflow/node_runner.ts @@ -4,8 +4,11 @@ * SPDX-License-Identifier: Apache-2.0 */ +import {context, type Span, SpanStatusCode, trace} from '@opentelemetry/api'; import {InvocationContext} from '../agents/invocation_context.js'; import {Event} from '../events/event.js'; +import {traceNodeExecution, tracer} from '../telemetry/tracing.js'; +import {formatError} from '../utils/error_utils.js'; import {BaseNode} from './base_node.js'; import {createSubBranch} from './branch_path.js'; import { @@ -84,20 +87,56 @@ export interface ExecuteChildNodeParams { * duplicate observable events across retries should emit only after their * fallible work has succeeded. */ -export async function executeChildNode({ - parent, - node, - input, - options = {}, - abortSignal, - nodeState: callerNodeState, -}: ExecuteChildNodeParams): Promise { +export function executeChildNode( + params: ExecuteChildNodeParams, +): Promise { + const {parent, node, options = {}} = params; const nodeName = options.nodeName ?? node.name; - const runId = options.runId ?? nodeName; const nodePath = options.overrideNodePath ?? (parent.nodePath ? `${parent.nodePath}.${nodeName}` : nodeName); + // The span is started here, synchronously, rather than inside `runChildNode`: + // its parent must be whatever span was active when the workflow SCHEDULED + // this node. Nodes are raced concurrently in `Workflow.runLoop`, so a parent + // captured any later would nest concurrent siblings inside whichever task + // happened to resolve first. + const span = tracer.startSpan(`execute_node ${nodeName}`); + + // Deliberately not `async`, and the callback is deliberately not `async` + // either: `context.with` hands the inner promise straight back, so the child + // settles on exactly the same microtask it did before tracing existed. Async + // wrappers here would each add promise-adoption ticks, which is enough to + // change how concurrently scheduled nodes interleave — an observability + // change must not move execution around. + return context.with(trace.setSpan(context.active(), span), () => + runChildNode({params, nodeName, nodePath, span}), + ); +} + +interface RunChildNodeParams { + params: ExecuteChildNodeParams; + nodeName: string; + nodePath: string; + span: Span; +} + +/** The body of {@link executeChildNode}, running under its `execute_node` span. */ +async function runChildNode({ + params: { + parent, + node, + input, + options = {}, + abortSignal, + nodeState: callerNodeState, + }, + nodeName, + nodePath, + span, +}: RunChildNodeParams): Promise { + const runId = options.runId ?? nodeName; + let branch = parent.branch; if (options.overrideBranch !== undefined) { branch = options.overrideBranch; @@ -150,67 +189,158 @@ export async function executeChildNode({ }); const pluginManager = child.invocationContext.pluginManager; - if (pluginManager?.hasPlugins) { - const skipOutput = await pluginManager.runBeforeNodeCallback({ - node, - nodeContext: child, - input, - }); - if (skipOutput !== undefined) { - child.output = skipOutput; - if (options.useAsOutput) { - parent.output = child.output; - parent.route = child.route; + + try { + if (pluginManager?.hasPlugins) { + const skipOutput = await pluginManager.runBeforeNodeCallback({ + node, + nodeContext: child, + input, + }); + if (skipOutput !== undefined) { + child.output = skipOutput; + // A skipped node still fills its slot in the trace, so record it as + // completed rather than leaving an attribute-less span behind. + traceNodeExecution({ + nodePath, + runId, + attempt: nodeState.attemptCount, + status: 'completed', + interruptCount: child.interruptIds.length, + }); + if (options.useAsOutput) { + parent.output = child.output; + parent.route = child.route; + } + return child; } - return child; } - } - let succeeded = false; - while (!succeeded) { - resetState(child); - child.attemptCount = nodeState.attemptCount; - try { - await runOnce({node, child, input, nodeName, branch, isolationScope}); - succeeded = true; - } catch (err) { - // Cancellation is terminal: an aborted invocation (or a sibling failure - // that cancelled this node) is never retried. - if (isInvocationAbortedError(err)) { - throw err; - } - // Check retry eligibility with the attempt that just failed, compute its - // backoff delay, THEN advance the counter (matches Python semantics). - const retryConfig = node.preparedRetryConfig; - if ( - !retryConfig || - !shouldRetryNode({error: err, retryConfig, nodeState}) - ) { - throw err; + let succeeded = false; + while (!succeeded) { + resetState(child); + child.attemptCount = nodeState.attemptCount; + try { + await runAttempt({ + node, + child, + input, + nodeName, + branch, + isolationScope, + nodePath, + runId, + attempt: nodeState.attemptCount, + }); + succeeded = true; + } catch (err) { + // Cancellation is terminal: an aborted invocation (or a sibling + // failure that cancelled this node) is never retried. + if (isInvocationAbortedError(err)) { + throw err; + } + // Check retry eligibility with the attempt that just failed, compute + // its backoff delay, THEN advance the counter (matches Python + // semantics). + const retryConfig = node.preparedRetryConfig; + if ( + !retryConfig || + !shouldRetryNode({error: err, retryConfig, nodeState}) + ) { + throw err; + } + const delaySeconds = getRetryDelaySeconds({retryConfig, nodeState}); + nodeState.attemptCount += 1; + await delay(delaySeconds * 1000, effectiveAbortSignal); } - const delaySeconds = getRetryDelaySeconds({retryConfig, nodeState}); - nodeState.attemptCount += 1; - await delay(delaySeconds * 1000, effectiveAbortSignal); } - } - if (pluginManager?.hasPlugins) { - const replacedOutput = await pluginManager.runAfterNodeCallback({ - node, - nodeContext: child, - output: child.output, + traceNodeExecution({ + nodePath, + runId, + attempt: nodeState.attemptCount, + status: child.interruptIds.length > 0 ? 'waiting' : 'completed', + interruptCount: child.interruptIds.length, }); - if (replacedOutput !== undefined) { - child.output = replacedOutput; + + if (pluginManager?.hasPlugins) { + const replacedOutput = await pluginManager.runAfterNodeCallback({ + node, + nodeContext: child, + output: child.output, + }); + if (replacedOutput !== undefined) { + child.output = replacedOutput; + } + } + + if (options.useAsOutput) { + parent.output = child.output; + parent.route = child.route; } - } - if (options.useAsOutput) { - parent.output = child.output; - parent.route = child.route; + return child; + } catch (err) { + traceNodeExecution({ + nodePath, + runId, + attempt: nodeState.attemptCount, + status: 'failed', + interruptCount: child.interruptIds.length, + }); + span.setStatus({code: SpanStatusCode.ERROR, message: formatError(err)}); + throw err; + } finally { + span.end(); } +} + +interface RunAttemptParams extends RunOnceParams { + nodePath: string; + runId: string; + attempt: number; +} - return child; +/** + * Not `async`: a node without a retry config must reach `runOnce` and settle on + * exactly the microtask it would have without tracing (see `executeChildNode`). + */ +function runAttempt(params: RunAttemptParams): Promise { + const {node, nodePath, runId, attempt} = params; + if (!node.preparedRetryConfig) { + return runOnce(params); + } + return tracer.startActiveSpan( + `execute_node_attempt ${params.nodeName}`, + async (span) => { + try { + await runOnce(params); + traceNodeExecution({ + nodePath, + runId, + attempt, + status: + params.child.interruptIds.length > 0 ? 'waiting' : 'completed', + interruptCount: params.child.interruptIds.length, + }); + } catch (err) { + traceNodeExecution({ + nodePath, + runId, + attempt, + status: 'failed', + interruptCount: params.child.interruptIds.length, + }); + span.setStatus({ + code: SpanStatusCode.ERROR, + message: formatError(err), + }); + throw err; + } finally { + span.end(); + } + }, + ); } /** diff --git a/core/src/workflow/workflow.ts b/core/src/workflow/workflow.ts index 2ef72558d..94c71977a 100644 --- a/core/src/workflow/workflow.ts +++ b/core/src/workflow/workflow.ts @@ -4,7 +4,9 @@ * SPDX-License-Identifier: Apache-2.0 */ +import {context, trace} from '@opentelemetry/api'; import {Event} from '../events/event.js'; +import {tracer, traceWorkflowInvocation} from '../telemetry/tracing.js'; import {experimental} from '../utils/experimental.js'; import {BaseNode, BaseNodeConfig} from './base_node.js'; import {commonPrefixOf} from './branch_path.js'; @@ -173,10 +175,21 @@ export class Workflow extends BaseNode { abort.controller.signal, ); + const span = tracer.startSpan(`invoke_workflow ${this.name}`); try { - await this.orchestrate(ctx, nodeInput, dynamicState, abort.controller); + // Sync callback returning the promise, not `async () => await …`: the + // wrapper must not insert microtask hops around the orchestration loop + // (see the note in `executeChildNode`). + await context.with(trace.setSpan(context.active(), span), () => { + traceWorkflowInvocation({ + workflowName: this.name, + nodePath: ctx.nodePath, + }); + return this.orchestrate(ctx, nodeInput, dynamicState, abort.controller); + }); } finally { abort.dispose(); + span.end(); } } diff --git a/core/test/workflow/telemetry_test.ts b/core/test/workflow/telemetry_test.ts new file mode 100644 index 000000000..5ba2ec2ed --- /dev/null +++ b/core/test/workflow/telemetry_test.ts @@ -0,0 +1,203 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {context, trace} from '@opentelemetry/api'; +import {AsyncLocalStorageContextManager} from '@opentelemetry/context-async-hooks'; +import { + BasicTracerProvider, + InMemorySpanExporter, + SimpleSpanProcessor, + type ReadableSpan, +} from '@opentelemetry/sdk-trace-base'; +import {afterAll, beforeAll, beforeEach, describe, expect, it} from 'vitest'; +import {Workflow} from '../../src/workflow/workflow.js'; +import {driveWorkflow, FnNode} from './test_helpers.js'; + +const exporter = new InMemorySpanExporter(); +const contextManager = new AsyncLocalStorageContextManager(); +const provider = new BasicTracerProvider({ + spanProcessors: [new SimpleSpanProcessor(exporter)], +}); + +beforeAll(() => { + context.setGlobalContextManager(contextManager.enable()); + trace.setGlobalTracerProvider(provider); +}); + +afterAll(async () => { + await provider.shutdown(); + trace.disable(); + context.disable(); + contextManager.disable(); +}); + +beforeEach(() => { + exporter.reset(); +}); + +function onlySpan(name: string): ReadableSpan { + const matches = exporter.getFinishedSpans().filter((s) => s.name === name); + expect(matches.map((s) => s.name)).toEqual([name]); + return matches[0]; +} + +function spansNamed(name: string): ReadableSpan[] { + return exporter.getFinishedSpans().filter((s) => s.name === name); +} + +function expectChildOf(child: ReadableSpan, parent: ReadableSpan): void { + expect(child.spanContext().traceId).toBe(parent.spanContext().traceId); + expect(child.parentSpanContext?.spanId).toBe(parent.spanContext().spanId); +} + +describe('workflow telemetry — span tree', () => { + it('nests every node span under the workflow span in a sequential graph', async () => { + const a = new FnNode('step_a', (_c, input) => `${input}->A`); + const b = new FnNode('step_b', (_c, input) => `${input}->B`); + const c = new FnNode('step_c', (_c, input) => `${input}->C`); + const wf = new Workflow({name: 'seq', edges: [['START', a, b, c]]}); + + expect((await driveWorkflow(wf, 'INIT')).output).toBe('INIT->A->B->C'); + + const wfSpan = onlySpan('invoke_workflow seq'); + const nodeSpans = ['step_a', 'step_b', 'step_c'].map((name) => + onlySpan(`execute_node ${name}`), + ); + for (const nodeSpan of nodeSpans) { + expectChildOf(nodeSpan, wfSpan); + } + + expectChildOf(wfSpan, onlySpan('execute_node seq')); + + expect(wfSpan.attributes).toMatchObject({ + 'gen_ai.operation.name': 'invoke_workflow', + 'adk.workflow.name': 'seq', + 'adk.node.path': 'seq', + }); + expect(nodeSpans[0].attributes).toMatchObject({ + 'gen_ai.operation.name': 'execute_node', + 'adk.node.path': 'seq.step_a', + 'adk.node.attempt': 1, + 'adk.node.status': 'completed', + 'adk.node.interrupt_count': 0, + }); + }); + + it('keeps concurrently scheduled nodes as siblings, not nested', async () => { + let startLeft!: () => void; + let startRight!: () => void; + const leftStarted = new Promise((r) => (startLeft = r)); + const rightStarted = new Promise((r) => (startRight = r)); + + // Each node runs a dynamic child while its sibling is also in flight, so + // the inner spans below can only land on the right parent if each node's + // context binding survives being interleaved with the other's. + const left = new FnNode('left', async (ctx) => { + startLeft(); + await rightStarted; + await ctx.runNode(new FnNode('inner_left', () => 'l'), 'x'); + return 'L'; + }); + const right = new FnNode('right', async (ctx) => { + startRight(); + await leftStarted; + await ctx.runNode(new FnNode('inner_right', () => 'r'), 'x'); + return 'R'; + }); + const join = new FnNode('join', (_c, input) => input); + const wf = new Workflow({ + name: 'fan', + edges: [['START', [left, right], join]], + }); + + await driveWorkflow(wf, 'x'); + + const wfSpan = onlySpan('invoke_workflow fan'); + const leftSpan = onlySpan('execute_node left'); + const rightSpan = onlySpan('execute_node right'); + + // Overlap is guaranteed structurally, not asserted on timestamps: neither + // node can return until the other has started, so the run only completes + // if both were in flight at once. OTel HrTime is too coarse to compare + // sub-millisecond span intervals reliably. + expectChildOf(leftSpan, wfSpan); + expectChildOf(rightSpan, wfSpan); + expect(leftSpan.parentSpanContext?.spanId).not.toBe( + rightSpan.spanContext().spanId, + ); + expect(rightSpan.parentSpanContext?.spanId).not.toBe( + leftSpan.spanContext().spanId, + ); + + // The point of binding the context per node: work started inside one node + // nests under that node, never under whichever sibling happens to be + // running alongside it. + expectChildOf(onlySpan('execute_node inner_left'), leftSpan); + expectChildOf(onlySpan('execute_node inner_right'), rightSpan); + }); + + it('emits one attempt span per try for a retried node', async () => { + let attempts = 0; + const flaky = new FnNode( + 'flaky', + () => { + attempts++; + if (attempts < 2) { + throw new Error('transient'); + } + return 'ok'; + }, + {retryConfig: {maxAttempts: 2, initialDelay: 0.001, jitter: 0}}, + ); + const wf = new Workflow({name: 'retry_wf', edges: [['START', flaky]]}); + + expect((await driveWorkflow(wf, 'x')).output).toBe('ok'); + expect(attempts).toBe(2); + + const nodeSpan = onlySpan('execute_node flaky'); + const attemptSpans = spansNamed('execute_node_attempt flaky'); + expect(attemptSpans).toHaveLength(2); + for (const attemptSpan of attemptSpans) { + expectChildOf(attemptSpan, nodeSpan); + } + + expect(attemptSpans[0].attributes).toMatchObject({ + 'adk.node.path': 'retry_wf.flaky', + 'adk.node.attempt': 1, + 'adk.node.status': 'failed', + }); + expect(attemptSpans[1].attributes).toMatchObject({ + 'adk.node.attempt': 2, + 'adk.node.status': 'completed', + }); + expect(nodeSpan.attributes).toMatchObject({ + 'adk.node.attempt': 2, + 'adk.node.status': 'completed', + }); + }); + + it('does not emit attempt spans for a node without a retry config', async () => { + const plain = new FnNode('plain', () => 'ok'); + const wf = new Workflow({name: 'plain_wf', edges: [['START', plain]]}); + + await driveWorkflow(wf, 'x'); + + expect(spansNamed('execute_node_attempt plain')).toHaveLength(0); + }); + + it('marks a failed node span as failed', async () => { + const boom = new FnNode('boom', () => { + throw new Error('kaboom'); + }); + const wf = new Workflow({name: 'fail_wf', edges: [['START', boom]]}); + + await expect(driveWorkflow(wf, 'x')).rejects.toThrow('kaboom'); + + expect(onlySpan('execute_node boom').attributes).toMatchObject({ + 'adk.node.status': 'failed', + }); + }); +}); diff --git a/package-lock.json b/package-lock.json index 5f83f3847..baab5ae92 100644 --- a/package-lock.json +++ b/package-lock.json @@ -76,6 +76,7 @@ }, "devDependencies": { "@mikro-orm/sqlite": "^6.6.6", + "@opentelemetry/context-async-hooks": "^2.1.0", "@types/adm-zip": "^0.5.8", "@types/express": "^4.17.25", "@types/lodash-es": "^4.17.12",