tomekpanek
diff --git a/‎core/pom.xml‎
Lines changed: 5 additions & 0 deletions b/‎core/pom.xml‎
Lines changed: 5 additions & 0 deletions
diff --git a/‎core/src/main/java/com/google/adk/Telemetry.java‎
Lines changed: 46 additions & 1 deletion b/‎core/src/main/java/com/google/adk/Telemetry.java‎
Lines changed: 46 additions & 1 deletion
diff --git a/‎core/src/main/java/com/google/adk/agents/BaseAgent.java‎
Lines changed: 50 additions & 40 deletions b/‎core/src/main/java/com/google/adk/agents/BaseAgent.java‎
Lines changed: 50 additions & 40 deletions
diff --git a/‎core/src/main/java/com/google/adk/events/EventActions.java‎
Lines changed: 30 additions & 0 deletions b/‎core/src/main/java/com/google/adk/events/EventActions.java‎
Lines changed: 30 additions & 0 deletions
@@ -184,6 +184,11 @@
       <groupId>io.opentelemetry</groupId>
       <artifactId>opentelemetry-sdk-trace</artifactId>
     </dependency>
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-sdk-testing</artifactId>
+      <scope>test</scope>
+    </dependency>
   </dependencies>
   <build>
     <resources>
 
@@ -30,10 +30,14 @@
 import io.opentelemetry.api.GlobalOpenTelemetry;
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.Tracer;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
+import io.reactivex.rxjava3.core.Flowable;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.function.Supplier;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -46,10 +50,17 @@
 public class Telemetry {
 
   private static final Logger log = LoggerFactory.getLogger(Telemetry.class);
-  private static final Tracer tracer = GlobalOpenTelemetry.getTracer("gcp.vertex.agent");
+
+  @SuppressWarnings("NonFinalStaticField")
+  private static Tracer tracer = GlobalOpenTelemetry.getTracer("gcp.vertex.agent");
 
   private Telemetry() {}
 
+  /** Sets the OpenTelemetry instance to be used for tracing. This is for testing purposes only. */
+  public static void setTracerForTesting(Tracer tracer) {
+    Telemetry.tracer = tracer;
+  }
+
   /**
    * Traces tool call arguments.
    *
@@ -219,4 +230,38 @@ public static void traceSendData(
   public static Tracer getTracer() {
     return tracer;
   }
+
+  /**
+   * Executes a Flowable with an OpenTelemetry Scope active for its entire lifecycle.
+   *
+   * <p>This helper manages the OpenTelemetry Scope lifecycle for RxJava Flowables to ensure proper
+   * context propagation across async boundaries. The scope remains active from when the Flowable is
+   * returned through all operators until stream completion (onComplete, onError, or cancel).
+   *
+   * <p><b>Why not try-with-resources?</b> RxJava Flowables execute lazily - operators run at
+   * subscription time, not at chain construction time. Using try-with-resources would close the
+   * scope before the Flowable subscribes, causing Context.current() to return ROOT in nested
+   * operations and breaking parent-child span relationships (fragmenting traces).
+   *
+   * <p>The scope is properly closed via doFinally when the stream terminates, ensuring no resource
+   * leaks regardless of completion mode (success, error, or cancellation).
+   *
+   * @param spanContext The context containing the span to activate
+   * @param span The span to end when the stream completes
+   * @param flowableSupplier Supplier that creates the Flowable to execute with active scope
+   * @param <T> The type of items emitted by the Flowable
+   * @return Flowable with OpenTelemetry scope lifecycle management
+   */
+  @SuppressWarnings("MustBeClosedChecker") // Scope lifecycle managed by RxJava doFinally
+  public static <T> Flowable<T> traceFlowable(
+      Context spanContext, Span span, Supplier<Flowable<T>> flowableSupplier) {
+    Scope scope = spanContext.makeCurrent();
+    return flowableSupplier
+        .get()
+        .doFinally(
+            () -> {
+              scope.close();
+              span.end();
+            });
+  }
 }
@@ -28,7 +28,7 @@
 import com.google.genai.types.Content;
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.Tracer;
-import io.opentelemetry.context.Scope;
+import io.opentelemetry.context.Context;
 import io.reactivex.rxjava3.core.Flowable;
 import io.reactivex.rxjava3.core.Maybe;
 import io.reactivex.rxjava3.core.Single;
@@ -205,39 +205,45 @@ public Flowable<Event> runAsync(InvocationContext parentContext) {
     Tracer tracer = Telemetry.getTracer();
     return Flowable.defer(
         () -> {
-          Span span = tracer.spanBuilder("agent_run [" + name() + "]").startSpan();
-          try (Scope scope = span.makeCurrent()) {
-            InvocationContext invocationContext = createInvocationContext(parentContext);
-
-            Flowable<Event> executionFlowable =
-                callCallback(
-                        beforeCallbacksToFunctions(
-                            invocationContext.pluginManager(),
-                            beforeAgentCallback.orElse(ImmutableList.of())),
-                        invocationContext)
-                    .flatMapPublisher(
-                        beforeEventOpt -> {
-                          if (invocationContext.endInvocation()) {
-                            return Flowable.fromOptional(beforeEventOpt);
-                          }
-
-                          Flowable<Event> beforeEvents = Flowable.fromOptional(beforeEventOpt);
-                          Flowable<Event> mainEvents =
-                              Flowable.defer(() -> runAsyncImpl(invocationContext));
-                          Flowable<Event> afterEvents =
-                              Flowable.defer(
-                                  () ->
-                                      callCallback(
-                                              afterCallbacksToFunctions(
-                                                  invocationContext.pluginManager(),
-                                                  afterAgentCallback.orElse(ImmutableList.of())),
-                                              invocationContext)
-                                          .flatMapPublisher(Flowable::fromOptional));
-
-                          return Flowable.concat(beforeEvents, mainEvents, afterEvents);
-                        });
-            return executionFlowable.doFinally(span::end);
-          }
+          Span span =
+              tracer
+                  .spanBuilder("agent_run [" + name() + "]")
+                  .setParent(Context.current())
+                  .startSpan();
+          Context spanContext = Context.current().with(span);
+
+          InvocationContext invocationContext = createInvocationContext(parentContext);
+
+          return Telemetry.traceFlowable(
+              spanContext,
+              span,
+              () ->
+                  callCallback(
+                          beforeCallbacksToFunctions(
+                              invocationContext.pluginManager(),
+                              beforeAgentCallback.orElse(ImmutableList.of())),
+                          invocationContext)
+                      .flatMapPublisher(
+                          beforeEventOpt -> {
+                            if (invocationContext.endInvocation()) {
+                              return Flowable.fromOptional(beforeEventOpt);
+                            }
+
+                            Flowable<Event> beforeEvents = Flowable.fromOptional(beforeEventOpt);
+                            Flowable<Event> mainEvents =
+                                Flowable.defer(() -> runAsyncImpl(invocationContext));
+                            Flowable<Event> afterEvents =
+                                Flowable.defer(
+                                    () ->
+                                        callCallback(
+                                                afterCallbacksToFunctions(
+                                                    invocationContext.pluginManager(),
+                                                    afterAgentCallback.orElse(ImmutableList.of())),
+                                                invocationContext)
+                                            .flatMapPublisher(Flowable::fromOptional));
+
+                            return Flowable.concat(beforeEvents, mainEvents, afterEvents);
+                          }));
         });
   }
 
@@ -340,12 +346,16 @@ public Flowable<Event> runLive(InvocationContext parentContext) {
     Tracer tracer = Telemetry.getTracer();
     return Flowable.defer(
         () -> {
-          Span span = tracer.spanBuilder("agent_run [" + name() + "]").startSpan();
-          try (Scope scope = span.makeCurrent()) {
-            InvocationContext invocationContext = createInvocationContext(parentContext);
-            Flowable<Event> executionFlowable = runLiveImpl(invocationContext);
-            return executionFlowable.doFinally(span::end);
-          }
+          Span span =
+              tracer
+                  .spanBuilder("agent_run [" + name() + "]")
+                  .setParent(Context.current())
+                  .startSpan();
+          Context spanContext = Context.current().with(span);
+
+          InvocationContext invocationContext = createInvocationContext(parentContext);
+
+          return Telemetry.traceFlowable(spanContext, span, () -> runLiveImpl(invocationContext));
         });
   }
 
 
@@ -17,6 +17,7 @@
 
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
+import com.google.adk.tools.ToolConfirmation;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import com.google.genai.types.Part;
 import java.util.Objects;
@@ -37,6 +38,8 @@ public class EventActions {
   private Optional<Boolean> escalate = Optional.empty();
   private ConcurrentMap<String, ConcurrentMap<String, Object>> requestedAuthConfigs =
       new ConcurrentHashMap<>();
+  private ConcurrentMap<String, ToolConfirmation> requestedToolConfirmations =
+      new ConcurrentHashMap<>();
   private Optional<Boolean> endInvocation = Optional.empty();
 
   /** Default constructor for Jackson. */
@@ -113,6 +116,16 @@ public void setRequestedAuthConfigs(
     this.requestedAuthConfigs = requestedAuthConfigs;
   }
 
+  @JsonProperty("requestedToolConfirmations")
+  public ConcurrentMap<String, ToolConfirmation> requestedToolConfirmations() {
+    return requestedToolConfirmations;
+  }
+
+  public void setRequestedToolConfirmations(
+      ConcurrentMap<String, ToolConfirmation> requestedToolConfirmations) {
+    this.requestedToolConfirmations = requestedToolConfirmations;
+  }
+
   @JsonProperty("endInvocation")
   public Optional<Boolean> endInvocation() {
     return endInvocation;
@@ -148,6 +161,7 @@ public boolean equals(Object o) {
         && Objects.equals(transferToAgent, that.transferToAgent)
         && Objects.equals(escalate, that.escalate)
         && Objects.equals(requestedAuthConfigs, that.requestedAuthConfigs)
+        && Objects.equals(requestedToolConfirmations, that.requestedToolConfirmations)
         && Objects.equals(endInvocation, that.endInvocation);
   }
 
@@ -160,6 +174,7 @@ public int hashCode() {
         transferToAgent,
         escalate,
         requestedAuthConfigs,
+        requestedToolConfirmations,
         endInvocation);
   }
 
@@ -172,6 +187,8 @@ public static class Builder {
     private Optional<Boolean> escalate = Optional.empty();
     private ConcurrentMap<String, ConcurrentMap<String, Object>> requestedAuthConfigs =
         new ConcurrentHashMap<>();
+    private ConcurrentMap<String, ToolConfirmation> requestedToolConfirmations =
+        new ConcurrentHashMap<>();
     private Optional<Boolean> endInvocation = Optional.empty();
 
     public Builder() {}
@@ -183,6 +200,8 @@ private Builder(EventActions eventActions) {
       this.transferToAgent = eventActions.transferToAgent();
       this.escalate = eventActions.escalate();
       this.requestedAuthConfigs = new ConcurrentHashMap<>(eventActions.requestedAuthConfigs());
+      this.requestedToolConfirmations =
+          new ConcurrentHashMap<>(eventActions.requestedToolConfirmations());
       this.endInvocation = eventActions.endInvocation();
     }
 
@@ -229,6 +248,13 @@ public Builder requestedAuthConfigs(
       return this;
     }
 
+    @CanIgnoreReturnValue
+    @JsonProperty("requestedToolConfirmations")
+    public Builder requestedToolConfirmations(ConcurrentMap<String, ToolConfirmation> value) {
+      this.requestedToolConfirmations = value;
+      return this;
+    }
+
     @CanIgnoreReturnValue
     @JsonProperty("endInvocation")
     public Builder endInvocation(boolean endInvocation) {
@@ -256,6 +282,9 @@ public Builder merge(EventActions other) {
       if (other.requestedAuthConfigs() != null) {
         this.requestedAuthConfigs.putAll(other.requestedAuthConfigs());
       }
+      if (other.requestedToolConfirmations() != null) {
+        this.requestedToolConfirmations.putAll(other.requestedToolConfirmations());
+      }
       if (other.endInvocation().isPresent()) {
         this.endInvocation = other.endInvocation();
       }
@@ -270,6 +299,7 @@ public EventActions build() {
       eventActions.setTransferToAgent(this.transferToAgent);
       eventActions.setEscalate(this.escalate);
       eventActions.setRequestedAuthConfigs(this.requestedAuthConfigs);
+      eventActions.setRequestedToolConfirmations(this.requestedToolConfirmations);
       eventActions.setEndInvocation(this.endInvocation);
       return eventActions;
     }