Skip to content

test: Adding a test for telemetry in RunnerTest #540

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
copybara-service merged 1 commit into from
Oct 23, 2025
Merged

test: Adding a test for telemetry in RunnerTest #540

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 8 additions & 1 deletion core/src/main/java/com/google/adk/Telemetry.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
*
Expand Down
30 changes: 0 additions & 30 deletions core/src/main/java/com/google/adk/events/EventActions.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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. */
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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);
}

Expand All @@ -174,7 +160,6 @@ public int hashCode() {
transferToAgent,
escalate,
requestedAuthConfigs,
requestedToolConfirmations,
endInvocation);
}

Expand All @@ -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() {}
Expand All @@ -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();
}

Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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();
}
Expand All @@ -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;
}
Expand Down
71 changes: 0 additions & 71 deletions core/src/main/java/com/google/adk/tools/ToolConfirmation.java
View file
Open in desktop

This file was deleted.

55 changes: 3 additions & 52 deletions core/src/main/java/com/google/adk/tools/ToolContext.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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() {
Expand All @@ -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.
Expand All @@ -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) {
Expand All @@ -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;
Expand All @@ -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);
}
}
}
52 changes: 52 additions & 0 deletions core/src/test/java/com/google/adk/runner/RunnerTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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");
Expand All @@ -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();
Expand Down Expand Up @@ -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 =
Expand Down Expand Up @@ -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();
}
}
Loading