feat: refactor OpenTelemetry (OTel) instrumentation within the ADK core, moving from manual span management to structured helper classes

### Key Changes

*   **Structured Instrumentation:** Replaces manual `Tracing.trace` calls and explicit `Scope` management with `Flowable.using` and `Maybe.using` patterns. It introduces helper classes like `AgentInvocation` and `ToolExecution` to encapsulate telemetry logic.
*   **Metrics Integration:** Adds support for tracking new metrics during agent execution:
    *   `gen_ai.agent.invocation.duration`
    *   `gen_ai.agent.request.size` / `gen_ai.agent.response.size`
    *   `gen_ai.agent.workflow.steps`
*   **Reactive API Improvements:** Leverages `Tracing.withContext()` and `doOnNext`/`doOnError` hooks within `AgentInvocation` and `ToolExecution` to automatically capture events and errors without polluting the core logic.
*   **Trace Hierarchy Refinement:** Updates how spans are nested. For example, in `ContextPropagationTest`, the child agent span is now correctly parented to the specific LLM call span that triggered it, rather than the parent agent span.
*   **Testing:** Significantly enhances `BaseAgentTest` to verify metric collection using `InMemoryMetricReader`. Updates various tests to match the new span naming convention (e.g., removing brackets from tool execution spans).

### Impact
This refactor simplifies the core agent and flow logic by removing boilerplate telemetry code, making the instrumentation more robust and easier to maintain while expanding the observability of the ADK through new metrics.

PiperOrigin-RevId: 918116236

Author: google-genai-bot

core/src/main/java/com/google/adk/agents/BaseAgent.java

@@ -24,13 +24,12 @@
 import com.google.adk.agents.Callbacks.BeforeAgentCallback;
 import com.google.adk.events.Event;
 import com.google.adk.plugins.Plugin;
-import com.google.adk.telemetry.Tracing;
+import com.google.adk.telemetry.Instrumentation.AgentInvocation;
 import com.google.adk.utils.AgentEnums.AgentOrigin;
 import com.google.common.collect.ImmutableList;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import com.google.errorprone.annotations.DoNotCall;
 import com.google.genai.types.Content;
-import io.opentelemetry.context.Context;
 import io.reactivex.rxjava3.core.Completable;
 import io.reactivex.rxjava3.core.Flowable;
 import io.reactivex.rxjava3.core.Maybe;
@@ -322,11 +321,10 @@ public Flowable<Event> runAsync(InvocationContext parentContext) {
   private Flowable<Event> run(
       InvocationContext parentContext,
       Function<InvocationContext, Flowable<Event>> runImplementation) {
-    Context parentSpanContext = Context.current();
-    return Flowable.defer(
-        () -> {
-          InvocationContext invocationContext = createInvocationContext(parentContext);
-
+    return Flowable.using(
+        () -> new AgentInvocation(createInvocationContext(parentContext), this),
+        agentInvocation -> {
+          InvocationContext invocationContext = agentInvocation.getCtx();
           Flowable<Event> mainAndAfterEvents =
               Flowable.defer(() -> runImplementation.apply(invocationContext))
                   .concatWith(
@@ -350,14 +348,10 @@ private Flowable<Event> run(
                     return Flowable.just(beforeEvent).concatWith(mainAndAfterEvents);
                   })
               .switchIfEmpty(mainAndAfterEvents)
-              .compose(
-                  Tracing.<Event>trace("invoke_agent " + name())
-                      .setParent(parentSpanContext)
-                      .configure(
-                          span ->
-                              Tracing.traceAgentInvocation(
-                                  span, name(), description(), invocationContext)));
-        });
+              .doOnNext(agentInvocation::addEvent)
+              .doOnError(agentInvocation::setError);
+        },
+        AgentInvocation::close);
   }
 
   /**

core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java

@@ -435,12 +435,10 @@ private Flowable<Event> runOneStep(Context spanContext, InvocationContext contex
                                                   "Agent not found: " + agentToTransfer)));
                                     }
                                     return postProcessedEvents.concatWith(
-                                        Flowable.defer(
-                                            () -> {
-                                              try (Scope s = spanContext.makeCurrent()) {
-                                                return nextAgent.get().runAsync(context);
-                                              }
-                                            }));
+                                        nextAgent
+                                            .get()
+                                            .runAsync(context)
+                                            .compose(Tracing.withContext(spanContext)));
                                   }
                                   return postProcessedEvents;
                                 });
@@ -622,12 +620,10 @@ public void onError(Throwable e) {
                                     "Agent not found: " + event.actions().transferToAgent().get());
                               }
                               Flowable<Event> nextAgentEvents =
-                                  Flowable.defer(
-                                      () -> {
-                                        try (Scope s = spanContext.makeCurrent()) {
-                                          return nextAgent.get().runLive(invocationContext);
-                                        }
-                                      });
+                                  nextAgent
+                                      .get()
+                                      .runLive(invocationContext)
+                                      .compose(Tracing.withContext(spanContext));
                               events = Flowable.concat(events, nextAgentEvents);
                             }
                             return events;

core/src/main/java/com/google/adk/flows/llmflows/Functions.java

@@ -29,6 +29,8 @@
 import com.google.adk.events.Event;
 import com.google.adk.events.EventActions;
 import com.google.adk.events.ToolConfirmation;
+import com.google.adk.telemetry.Instrumentation;
+import com.google.adk.telemetry.Instrumentation.ToolExecution;
 import com.google.adk.telemetry.Tracing;
 import com.google.adk.tools.BaseTool;
 import com.google.adk.tools.FunctionTool;
@@ -417,57 +419,50 @@ private static Maybe<Event> postProcessFunctionResult(
       ToolContext toolContext,
       boolean isLive,
       Context parentContext) {
+    return Maybe.using(
+        () ->
+            Instrumentation.recordToolExecution(
+                tool, invocationContext.agent(), functionArgs, parentContext),
+        toolExecution ->
+            processFunctionResult(
+                    maybeFunctionResult, invocationContext, tool, functionArgs, toolContext, isLive)
+                .doOnSuccess(event -> toolExecution.context().setFunctionResponseEvent(event))
+                .doOnError(toolExecution::setError),
+        ToolExecution::close);
+  }
+
+  private static Maybe<Event> processFunctionResult(
+      Maybe<Map<String, Object>> maybeFunctionResult,
+      InvocationContext invocationContext,
+      BaseTool tool,
+      Map<String, Object> functionArgs,
+      ToolContext toolContext,
+      boolean isLive) {
     return maybeFunctionResult
         .map(Optional::of)
         .defaultIfEmpty(Optional.empty())
         .onErrorResumeNext(
-            t -> {
-              Maybe<Map<String, Object>> errorCallbackResult =
-                  handleOnToolErrorCallback(invocationContext, tool, functionArgs, toolContext, t);
-              Maybe<Optional<Map<String, Object>>> mappedResult;
-              if (isLive) {
-                // In live mode, handle null results from the error callback gracefully.
-                mappedResult = errorCallbackResult.map(Optional::ofNullable);
-              } else {
-                // In non-live mode, a null result from the error callback will cause an NPE
-                // when wrapped with Optional.of(), potentially matching prior behavior.
-                mappedResult = errorCallbackResult.map(Optional::of);
-              }
-              return mappedResult.switchIfEmpty(Single.error(t));
-            })
+            t ->
+                handleOnToolErrorCallback(invocationContext, tool, functionArgs, toolContext, t)
+                    .map(res -> isLive ? Optional.ofNullable(res) : Optional.of(res))
+                    .switchIfEmpty(Single.error(t)))
         .flatMapMaybe(
             optionalInitialResult -> {
               Map<String, Object> initialFunctionResult = optionalInitialResult.orElse(null);
-
               return maybeInvokeAfterToolCall(
                       invocationContext, tool, functionArgs, toolContext, initialFunctionResult)
                   .map(Optional::of)
-                  .defaultIfEmpty(Optional.ofNullable(initialFunctionResult))
-                  .flatMapMaybe(
-                      finalOptionalResult -> {
-                        Map<String, Object> finalFunctionResult = finalOptionalResult.orElse(null);
-                        if (tool.longRunning() && finalFunctionResult == null) {
-                          return Maybe.empty();
-                        }
-                        Event event =
-                            buildResponseEvent(
-                                tool, finalFunctionResult, toolContext, invocationContext);
-                        return Maybe.just(event);
-                      });
+                  .defaultIfEmpty(Optional.ofNullable(initialFunctionResult));
             })
-        .compose(
-            Tracing.<Event>trace("execute_tool [" + tool.name() + "]")
-                .setParent(parentContext)
-                .onSuccess(
-                    (span, event) ->
-                        Tracing.traceToolExecution(
-                            span,
-                            tool.name(),
-                            tool.description(),
-                            tool.getClass().getSimpleName(),
-                            functionArgs,
-                            event,
-                            null)));
+        .flatMapMaybe(
+            finalOptionalResult -> {
+              if (tool.longRunning() && finalOptionalResult.isEmpty()) {
+                return Maybe.empty();
+              }
+              return Maybe.just(
+                  buildResponseEvent(
+                      tool, finalOptionalResult.orElse(null), toolContext, invocationContext));
+            });
   }
 
   private static Optional<Event> mergeParallelFunctionResponseEvents(

core/src/main/java/com/google/adk/telemetry/Instrumentation.java

@@ -160,8 +160,13 @@ public static final class ToolExecution extends ClosableTelemetryScope {
     private final BaseAgent agent;
     private final Map<String, Object> functionArgs;
 
-    public ToolExecution(BaseTool tool, BaseAgent agent, Map<String, Object> functionArgs) {
-      super(Tracing.getTracer().spanBuilder("execute_tool " + tool.name()).startSpan());
+    public ToolExecution(
+        BaseTool tool, BaseAgent agent, Map<String, Object> functionArgs, Context parentContext) {
+      super(
+          Tracing.getTracer()
+              .spanBuilder("execute_tool " + tool.name())
+              .setParent(parentContext)
+              .startSpan());
       this.tool = tool;
       this.agent = agent;
       this.functionArgs = functionArgs;
@@ -202,6 +207,11 @@ public static AgentInvocation recordAgentInvocation(InvocationContext ctx, BaseA
   /** Creates a ToolExecution context to record tool execution telemetry. */
   public static ToolExecution recordToolExecution(
       BaseTool tool, BaseAgent agent, Map<String, Object> functionArgs) {
-    return new ToolExecution(tool, agent, functionArgs);
+    return recordToolExecution(tool, agent, functionArgs, Context.current());
+  }
+
+  public static ToolExecution recordToolExecution(
+      BaseTool tool, BaseAgent agent, Map<String, Object> functionArgs, Context parentContext) {
+    return new ToolExecution(tool, agent, functionArgs, parentContext);
   }
 }

core/src/test/java/com/google/adk/agents/BaseAgentTest.java

@@ -22,26 +22,42 @@
 import com.google.adk.agents.Callbacks.AfterAgentCallback;
 import com.google.adk.agents.Callbacks.BeforeAgentCallback;
 import com.google.adk.events.Event;
+import com.google.adk.telemetry.Metrics;
 import com.google.adk.testing.TestBaseAgent;
 import com.google.adk.testing.TestCallback;
 import com.google.adk.testing.TestUtils;
 import com.google.common.collect.ImmutableList;
 import com.google.genai.types.Content;
 import com.google.genai.types.Part;
+import io.opentelemetry.api.GlobalOpenTelemetry;
+import io.opentelemetry.api.common.AttributeKey;
+import io.opentelemetry.api.metrics.Meter;
+import io.opentelemetry.sdk.OpenTelemetrySdk;
+import io.opentelemetry.sdk.metrics.SdkMeterProvider;
+import io.opentelemetry.sdk.metrics.data.HistogramPointData;
+import io.opentelemetry.sdk.metrics.data.MetricData;
+import io.opentelemetry.sdk.testing.exporter.InMemoryMetricReader;
+import io.opentelemetry.sdk.testing.time.TestClock;
+import io.opentelemetry.sdk.trace.SdkTracerProvider;
 import io.reactivex.rxjava3.core.Completable;
 import io.reactivex.rxjava3.core.Maybe;
 import java.util.List;
 import java.util.concurrent.atomic.AtomicBoolean;
+import org.junit.After;
+import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 
 @RunWith(JUnit4.class)
 public final class BaseAgentTest {
-
   private static final String TEST_AGENT_NAME = "testAgent";
   private static final String TEST_AGENT_DESCRIPTION = "A test agent";
 
+  private InMemoryMetricReader inMemoryMetricReader;
+  private TestClock testClock;
+  private Meter originalMeter;
+
   private static class ClosableTestAgent extends TestBaseAgent {
     final AtomicBoolean closed = new AtomicBoolean(false);
 
@@ -56,6 +72,35 @@ public Completable close() {
     }
   }
 
+  @Before
+  public void setUp() {
+    GlobalOpenTelemetry.resetForTest();
+    testClock = TestClock.create();
+    inMemoryMetricReader = InMemoryMetricReader.create();
+    SdkMeterProvider sdkMeterProvider =
+        SdkMeterProvider.builder()
+            .registerMetricReader(inMemoryMetricReader)
+            .setClock(testClock)
+            .build();
+
+    OpenTelemetrySdk openTelemetrySdk =
+        OpenTelemetrySdk.builder()
+            .setTracerProvider(SdkTracerProvider.builder().build())
+            .setMeterProvider(sdkMeterProvider)
+            .build();
+
+    GlobalOpenTelemetry.set(openTelemetrySdk);
+    originalMeter = GlobalOpenTelemetry.getMeter("gcp.vertex.agent");
+    Metrics.setMeterForTesting(openTelemetrySdk.getMeter("gcp.vertex.agent"));
+  }
+
+  @After
+  public void tearDown() {
+    if (originalMeter != null) {
+      Metrics.setMeterForTesting(originalMeter);
+    }
+  }
+
   @Test
   public void constructor_setsNameAndDescription() {
     String name = "testName";
@@ -173,6 +218,36 @@ public void runAsync_noCallbacks_invokesRunAsyncImpl() {
     assertThat(results).hasSize(1);
     assertThat(results.get(0).content()).hasValue(runAsyncImplContent);
     assertThat(runAsyncImpl.wasCalled()).isTrue();
+    MetricData durationMetric = findMetricByName("gen_ai.agent.invocation.duration");
+    assertThat(durationMetric.getUnit()).isEqualTo("ms");
+    HistogramPointData durationPoint =
+        (HistogramPointData) durationMetric.getHistogramData().getPoints().iterator().next();
+    assertThat(durationPoint.getAttributes().get(AttributeKey.stringKey("gen_ai.agent.name")))
+        .isEqualTo("testAgent");
+
+    MetricData reqSizeMetric = findMetricByName("gen_ai.agent.request.size");
+    assertThat(reqSizeMetric.getUnit()).isEqualTo("By");
+    HistogramPointData reqSizePoint =
+        (HistogramPointData) reqSizeMetric.getHistogramData().getPoints().iterator().next();
+    assertThat(reqSizePoint.getSum()).isEqualTo(12.0);
+    assertThat(reqSizePoint.getAttributes().get(AttributeKey.stringKey("gen_ai.agent.name")))
+        .isEqualTo("testAgent");
+
+    MetricData respSizeMetric = findMetricByName("gen_ai.agent.response.size");
+    assertThat(respSizeMetric.getUnit()).isEqualTo("By");
+    HistogramPointData respSizePoint =
+        (HistogramPointData) respSizeMetric.getHistogramData().getPoints().iterator().next();
+    assertThat(respSizePoint.getSum()).isEqualTo(11.0);
+    assertThat(respSizePoint.getAttributes().get(AttributeKey.stringKey("gen_ai.agent.name")))
+        .isEqualTo("testAgent");
+
+    MetricData workflowStepsMetric = findMetricByName("gen_ai.agent.workflow.steps");
+    assertThat(workflowStepsMetric.getUnit()).isEqualTo("1");
+    HistogramPointData workflowStepsPoint =
+        (HistogramPointData) workflowStepsMetric.getHistogramData().getPoints().iterator().next();
+    assertThat(workflowStepsPoint.getSum()).isEqualTo(1.0);
+    assertThat(workflowStepsPoint.getAttributes().get(AttributeKey.stringKey("gen_ai.agent.name")))
+        .isEqualTo("testAgent");
   }
 
   @Test
@@ -627,4 +702,11 @@ public void close_twoLevelsSubAgents_closesAllSubAgents() {
     assertThat(subAgent.closed.get()).isTrue();
     assertThat(subSubAgent.closed.get()).isTrue();
   }
+
+  private MetricData findMetricByName(String name) {
+    return inMemoryMetricReader.collectAllMetrics().stream()
+        .filter(m -> m.getName().equals(name))
+        .findFirst()
+        .orElseThrow(() -> new AssertionError("Metric not found: " + name));
+  }
 }

core/src/test/java/com/google/adk/agents/LlmAgentTest.java

@@ -494,7 +494,7 @@ public void runAsync_withTools_createsToolSpans() throws InterruptedException {
     List<SpanData> spans = openTelemetryRule.getSpans();
     SpanData agentSpan = findSpanByName(spans, "invoke_agent test agent");
     List<SpanData> llmSpans = findSpansByName(spans, "call_llm");
-    List<SpanData> toolSpans = findSpansByName(spans, "execute_tool [echo_tool]");
+    List<SpanData> toolSpans = findSpansByName(spans, "execute_tool echo_tool");
 
     assertThat(llmSpans).hasSize(2);
     assertThat(toolSpans).hasSize(1);

core/src/test/java/com/google/adk/runner/RunnerTest.java

@@ -1330,7 +1330,7 @@ public void runAsync_createsToolSpansWithCorrectParent() {
     List<SpanData> spans = openTelemetryRule.getSpans();
     List<SpanData> llmSpans = spans.stream().filter(s -> s.getName().equals("call_llm")).toList();
     List<SpanData> toolSpans =
-        spans.stream().filter(s -> s.getName().equals("execute_tool [echo_tool]")).toList();
+        spans.stream().filter(s -> s.getName().equals("execute_tool echo_tool")).toList();
 
     assertThat(llmSpans).hasSize(2);
     assertThat(toolSpans).hasSize(1);
@@ -1365,7 +1365,7 @@ public void runLive_createsToolSpansWithCorrectParent() throws Exception {
     List<SpanData> spans = openTelemetryRule.getSpans();
     List<SpanData> llmSpans = spans.stream().filter(s -> s.getName().equals("call_llm")).toList();
     List<SpanData> toolSpans =
-        spans.stream().filter(s -> s.getName().equals("execute_tool [echo_tool]")).toList();
+        spans.stream().filter(s -> s.getName().equals("execute_tool echo_tool")).toList();
 
     // In runLive, there is one call_llm span for the execution
     assertThat(llmSpans).hasSize(1);