test: Adding a test for telemetry in RunnerTest

PiperOrigin-RevId: 823192673

Author: google-genai-bot

core/src/main/java/com/google/adk/Telemetry.java

@@ -50,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.
    *

core/src/main/java/com/google/adk/events/EventActions.java

@@ -17,7 +17,6 @@
 
 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;
@@ -38,8 +37,6 @@ 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. */
@@ -116,16 +113,6 @@ 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;
@@ -161,7 +148,6 @@ 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);
   }
 
@@ -174,7 +160,6 @@ public int hashCode() {
         transferToAgent,
         escalate,
         requestedAuthConfigs,
-        requestedToolConfirmations,
         endInvocation);
   }
 
@@ -187,8 +172,6 @@ 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() {}
@@ -200,8 +183,6 @@ 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();
     }
 
@@ -248,13 +229,6 @@ 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) {
@@ -282,9 +256,6 @@ 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();
       }
@@ -299,7 +270,6 @@ 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;
     }

core/src/main/java/com/google/adk/tools/ToolConfirmation.java

@@ -23,21 +23,17 @@
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import io.reactivex.rxjava3.core.Single;
 import java.util.Optional;
-import javax.annotation.Nullable;
 
 /** ToolContext object provides a structured context for executing tools or functions. */
 public class ToolContext extends CallbackContext {
   private Optional<String> functionCallId = Optional.empty();
-  private Optional<ToolConfirmation> toolConfirmation = Optional.empty();
 
   private ToolContext(
       InvocationContext invocationContext,
       EventActions eventActions,
-      Optional<String> functionCallId,
-      Optional<ToolConfirmation> toolConfirmation) {
+      Optional<String> functionCallId) {
     super(invocationContext, eventActions);
     this.functionCallId = functionCallId;
-    this.toolConfirmation = toolConfirmation;
   }
 
   public EventActions actions() {
@@ -56,14 +52,6 @@ public void functionCallId(String functionCallId) {
     this.functionCallId = Optional.ofNullable(functionCallId);
   }
 
-  public Optional<ToolConfirmation> toolConfirmation() {
-    return toolConfirmation;
-  }
-
-  public void toolConfirmation(ToolConfirmation toolConfirmation) {
-    this.toolConfirmation = Optional.ofNullable(toolConfirmation);
-  }
-
   @SuppressWarnings("unused")
   private void requestCredential() {
     // TODO: b/414678311 - Implement credential request logic. Make this public.
@@ -76,35 +64,6 @@ private void getAuthResponse() {
     throw new UnsupportedOperationException("Auth response retrieval not implemented yet.");
   }
 
-  /**
-   * Requests confirmation for the given function call.
-   *
-   * @param hint A hint to the user on how to confirm the tool call.
-   * @param payload The payload used to confirm the tool call.
-   */
-  public void requestConfirmation(@Nullable String hint, @Nullable Object payload) {
-    if (functionCallId.isEmpty()) {
-      throw new IllegalStateException("function_call_id is not set.");
-    }
-    this.eventActions
-        .requestedToolConfirmations()
-        .put(functionCallId.get(), ToolConfirmation.builder().hint(hint).payload(payload).build());
-  }
-
-  /**
-   * Requests confirmation for the given function call.
-   *
-   * @param hint A hint to the user on how to confirm the tool call.
-   */
-  public void requestConfirmation(@Nullable String hint) {
-    requestConfirmation(hint, null);
-  }
-
-  /** Requests confirmation for the given function call. */
-  public void requestConfirmation() {
-    requestConfirmation(null, null);
-  }
-
   /** Searches the memory of the current user. */
   public Single<SearchMemoryResponse> searchMemory(String query) {
     if (invocationContext.memoryService() == null) {
@@ -123,16 +82,14 @@ public static Builder builder(InvocationContext invocationContext) {
   public Builder toBuilder() {
     return new Builder(invocationContext)
         .actions(eventActions)
-        .functionCallId(functionCallId.orElse(null))
-        .toolConfirmation(toolConfirmation.orElse(null));
+        .functionCallId(functionCallId.orElse(null));
   }
 
   /** Builder for {@link ToolContext}. */
   public static final class Builder {
     private final InvocationContext invocationContext;
     private EventActions eventActions = EventActions.builder().build(); // Default empty actions
     private Optional<String> functionCallId = Optional.empty();
-    private Optional<ToolConfirmation> toolConfirmation = Optional.empty();
 
     private Builder(InvocationContext invocationContext) {
       this.invocationContext = invocationContext;
@@ -150,14 +107,8 @@ public Builder functionCallId(String functionCallId) {
       return this;
     }
 
-    @CanIgnoreReturnValue
-    public Builder toolConfirmation(ToolConfirmation toolConfirmation) {
-      this.toolConfirmation = Optional.ofNullable(toolConfirmation);
-      return this;
-    }
-
     public ToolContext build() {
-      return new ToolContext(invocationContext, eventActions, functionCallId, toolConfirmation);
+      return new ToolContext(invocationContext, eventActions, functionCallId);
     }
   }
 }

core/src/main/java/com/google/adk/tools/ToolContext.java

@@ -28,7 +28,9 @@
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import com.google.adk.Telemetry;
 import com.google.adk.agents.InvocationContext;
+import com.google.adk.agents.LiveRequestQueue;
 import com.google.adk.agents.LlmAgent;
 import com.google.adk.agents.RunConfig;
 import com.google.adk.events.Event;
@@ -44,19 +46,27 @@
 import com.google.genai.types.Content;
 import com.google.genai.types.FunctionCall;
 import com.google.genai.types.Part;
+import io.opentelemetry.api.trace.Tracer;
+import io.opentelemetry.sdk.testing.junit4.OpenTelemetryRule;
+import io.opentelemetry.sdk.trace.data.SpanData;
 import io.reactivex.rxjava3.core.Completable;
 import io.reactivex.rxjava3.core.Flowable;
 import io.reactivex.rxjava3.core.Maybe;
 import java.util.List;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.concurrent.ConcurrentHashMap;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 import org.mockito.ArgumentCaptor;
 
 @RunWith(JUnit4.class)
 public final class RunnerTest {
+  @Rule public final OpenTelemetryRule openTelemetryRule = OpenTelemetryRule.create();
 
   private final BasePlugin plugin = mockPlugin("test");
   private final Content pluginContent = createContent("from plugin");
@@ -65,6 +75,7 @@ public final class RunnerTest {
   private final Runner runner = new InMemoryRunner(agent, "test", ImmutableList.of(plugin));
   private final Session session =
       runner.sessionService().createSession("test", "user").blockingGet();
+  private Tracer originalTracer;
 
   private final FailingEchoTool failingEchoTool = new FailingEchoTool();
   private final EchoTool echoTool = new EchoTool();
@@ -94,6 +105,17 @@ private BasePlugin mockPlugin(String name) {
     return plugin;
   }
 
+  @Before
+  public void setUp() {
+    this.originalTracer = Telemetry.getTracer();
+    Telemetry.setTracerForTesting(openTelemetryRule.getOpenTelemetry().getTracer("RunnerTest"));
+  }
+
+  @After
+  public void tearDown() {
+    Telemetry.setTracerForTesting(originalTracer);
+  }
+
   @Test
   public void pluginDoesNothing() {
     var events =
@@ -606,4 +628,34 @@ public void beforeRunCallback_withStateDelta_seesMergedState() {
   private Content createContent(String text) {
     return Content.builder().parts(Part.builder().text(text).build()).build();
   }
+
+  @Test
+  public void runAsync_createsInvocationSpan() {
+    var unused =
+        runner.runAsync("user", session.id(), createContent("test message")).toList().blockingGet();
+
+    List<SpanData> spans = openTelemetryRule.getSpans();
+    assertThat(spans).isNotEmpty();
+
+    Optional<SpanData> invocationSpan =
+        spans.stream().filter(span -> Objects.equals(span.getName(), "invocation")).findFirst();
+
+    assertThat(invocationSpan).isPresent();
+    assertThat(invocationSpan.get().hasEnded()).isTrue();
+  }
+
+  @Test
+  public void runLive_createsInvocationSpan() {
+    LiveRequestQueue liveRequestQueue = new LiveRequestQueue();
+    var unused = runner.runLive(session, liveRequestQueue, RunConfig.builder().build()).test();
+
+    List<SpanData> spans = openTelemetryRule.getSpans();
+    assertThat(spans).isNotEmpty();
+
+    Optional<SpanData> invocationSpan =
+        spans.stream().filter(span -> Objects.equals(span.getName(), "invocation")).findFirst();
+
+    assertThat(invocationSpan).isPresent();
+    assertThat(invocationSpan.get().hasEnded()).isTrue();
+  }
 }