From fc1abd912475ae25878521b7cfa44638e5c6bf9f Mon Sep 17 00:00:00 2001 From: Google Team Member Date: Tue, 15 Sep 2026 12:39:49 -0700 Subject: [PATCH] fix: respect OpenTelemetry LIFO scoping in reactive streams Refactor ADK telemetry to avoid holding thread-local OpenTelemetry Scope instances open across asynchronous reactive stream lifecycles. PiperOrigin-RevId: 981979719 --- .../java/com/google/adk/agents/BaseAgent.java | 49 ++--- .../google/adk/flows/llmflows/Functions.java | 76 +++++--- .../google/adk/telemetry/Instrumentation.java | 24 ++- .../com/google/adk/telemetry/Tracing.java | 115 ++++++------ .../adk/telemetry/ContextPropagationTest.java | 170 ++++++++++++++++++ .../adk/telemetry/InstrumentationTest.java | 35 ++++ 6 files changed, 355 insertions(+), 114 deletions(-) diff --git a/core/src/main/java/com/google/adk/agents/BaseAgent.java b/core/src/main/java/com/google/adk/agents/BaseAgent.java index 0a3c550a4..eeb04bdf1 100644 --- a/core/src/main/java/com/google/adk/agents/BaseAgent.java +++ b/core/src/main/java/com/google/adk/agents/BaseAgent.java @@ -26,6 +26,7 @@ import com.google.adk.plugins.Plugin; import com.google.adk.telemetry.Instrumentation; import com.google.adk.telemetry.Instrumentation.AgentInvocation; +import com.google.adk.telemetry.Tracing; import com.google.adk.utils.AgentEnums.AgentOrigin; import com.google.common.collect.ImmutableList; import com.google.errorprone.annotations.CanIgnoreReturnValue; @@ -331,31 +332,39 @@ private Flowable run( }, agentInvocation -> { InvocationContext invocationContext = agentInvocation.getCtx(); + Context otelContext = agentInvocation.context().otelContext(); Flowable mainAndAfterEvents = Flowable.defer(() -> runImplementation.apply(invocationContext)) + .compose(Tracing.withContext(otelContext)) .concatWith( Flowable.defer( - () -> - callCallback( - afterCallbacksToFunctions( - invocationContext.pluginManager(), afterAgentCallback), - invocationContext) - .toFlowable())); - - return callCallback( - beforeCallbacksToFunctions( - invocationContext.pluginManager(), beforeAgentCallback), - invocationContext) - .flatMapPublisher( - beforeEvent -> { - if (invocationContext.endInvocation()) { - return Flowable.just(beforeEvent); - } - return Flowable.just(beforeEvent).concatWith(mainAndAfterEvents); - }) - .switchIfEmpty(mainAndAfterEvents) + () -> + callCallback( + afterCallbacksToFunctions( + invocationContext.pluginManager(), + afterAgentCallback), + invocationContext) + .toFlowable()) + .compose(Tracing.withContext(otelContext))); + + return Flowable.defer( + () -> + callCallback( + beforeCallbacksToFunctions( + invocationContext.pluginManager(), beforeAgentCallback), + invocationContext) + .compose(Tracing.withContext(otelContext)) + .flatMapPublisher( + beforeEvent -> { + if (invocationContext.endInvocation()) { + return Flowable.just(beforeEvent); + } + return Flowable.just(beforeEvent).concatWith(mainAndAfterEvents); + }) + .switchIfEmpty(mainAndAfterEvents)) .doOnNext(agentInvocation::addEvent) - .doOnError(agentInvocation::setError); + .doOnError(agentInvocation::setError) + .compose(Tracing.withContext(otelContext)); }, AgentInvocation::close); } diff --git a/core/src/main/java/com/google/adk/flows/llmflows/Functions.java b/core/src/main/java/com/google/adk/flows/llmflows/Functions.java index 2b5c07435..3a1d45c2f 100644 --- a/core/src/main/java/com/google/adk/flows/llmflows/Functions.java +++ b/core/src/main/java/com/google/adk/flows/llmflows/Functions.java @@ -304,25 +304,10 @@ private static Function> getFunctionCallMapper( Map functionArgs = functionCall.args().map(HashMap::new).orElse(new HashMap<>()); - Maybe> maybeFunctionResult = - maybeInvokeBeforeToolCall(invocationContext, tool, functionArgs, toolContext) - .switchIfEmpty( - Maybe.defer( - () -> - isLive - ? processFunctionLive( - invocationContext, - tool, - toolContext, - functionCall, - functionArgs) - : callTool(tool, functionArgs, toolContext)) - .compose(Tracing.withContext(parentContext))); - return postProcessFunctionResult( - maybeFunctionResult, invocationContext, tool, + functionCall, functionArgs, toolContext, isLive, @@ -487,9 +472,9 @@ static boolean hasPendingLongRunningCall(List events) { } private static Maybe postProcessFunctionResult( - Maybe> maybeFunctionResult, InvocationContext invocationContext, BaseTool tool, + FunctionCall functionCall, Map functionArgs, ToolContext toolContext, boolean isLive, @@ -498,11 +483,38 @@ private static Maybe postProcessFunctionResult( () -> 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 -> { + Context toolOtelContext = toolExecution.context().otelContext(); + Maybe> maybeFunctionResult = + Maybe.defer( + () -> + maybeInvokeBeforeToolCall( + invocationContext, tool, functionArgs, toolContext)) + .compose(Tracing.withContext(toolOtelContext)) + .switchIfEmpty( + Maybe.defer( + () -> + isLive + ? processFunctionLive( + invocationContext, + tool, + toolContext, + functionCall, + functionArgs) + : callTool(tool, functionArgs, toolContext)) + .compose(Tracing.withContext(toolOtelContext))); + return processFunctionResult( + maybeFunctionResult, + invocationContext, + tool, + functionArgs, + toolContext, + isLive, + toolOtelContext) + .compose(Tracing.withContext(toolOtelContext)) + .doOnSuccess(event -> toolExecution.context().setFunctionResponseEvent(event)) + .doOnError(toolExecution::setError); + }, ToolExecution::close); } @@ -512,14 +524,19 @@ private static Maybe processFunctionResult( BaseTool tool, Map functionArgs, ToolContext toolContext, - boolean isLive) { + boolean isLive, + Context toolOtelContext) { return maybeFunctionResult .map(Optional::of) .defaultIfEmpty(Optional.empty()) .onErrorResumeNext( t -> { Maybe> errorCallbackResult = - handleOnToolErrorCallback(invocationContext, tool, functionArgs, toolContext, t); + Maybe.defer( + () -> + handleOnToolErrorCallback( + invocationContext, tool, functionArgs, toolContext, t)) + .compose(Tracing.withContext(toolOtelContext)); Maybe>> mappedResult; if (isLive) { // In live mode, handle null results from the error callback gracefully. @@ -535,8 +552,15 @@ private static Maybe processFunctionResult( optionalInitialResult -> { Map initialFunctionResult = optionalInitialResult.orElse(null); - return maybeInvokeAfterToolCall( - invocationContext, tool, functionArgs, toolContext, initialFunctionResult) + return Maybe.defer( + () -> + maybeInvokeAfterToolCall( + invocationContext, + tool, + functionArgs, + toolContext, + initialFunctionResult)) + .compose(Tracing.withContext(toolOtelContext)) .map(Optional::of) .defaultIfEmpty(Optional.ofNullable(initialFunctionResult)) .flatMapMaybe( diff --git a/core/src/main/java/com/google/adk/telemetry/Instrumentation.java b/core/src/main/java/com/google/adk/telemetry/Instrumentation.java index 620bb0f02..71be2ba00 100644 --- a/core/src/main/java/com/google/adk/telemetry/Instrumentation.java +++ b/core/src/main/java/com/google/adk/telemetry/Instrumentation.java @@ -23,7 +23,6 @@ import io.opentelemetry.api.trace.Span; import io.opentelemetry.api.trace.StatusCode; import io.opentelemetry.context.Context; -import io.opentelemetry.context.Scope; import java.time.Duration; import java.util.ArrayList; import java.util.Collections; @@ -91,9 +90,6 @@ public abstract static class ClosableTelemetryScope implements AutoCloseable { /** The OpenTelemetry span associated with this scope. */ protected final Span span; - /** The OpenTelemetry scope associated with this span. */ - protected final Scope scope; - /** The telemetry context for this scope. */ protected final TelemetryContext telemetryContext; @@ -104,16 +100,15 @@ public abstract static class ClosableTelemetryScope implements AutoCloseable { protected final AtomicBoolean closed = new AtomicBoolean(false); /** - * Constructs a new {@code ClosableTelemetryScope} with the given span. + * Constructs a new {@code ClosableTelemetryScope} with the given span and parent context. * * @param span The OpenTelemetry span to manage. + * @param parentContext The OpenTelemetry parent context. */ - @SuppressWarnings("MustBeClosedChecker") - ClosableTelemetryScope(Span span) { + ClosableTelemetryScope(Span span, Context parentContext) { this.startTimeNanos = System.nanoTime(); this.span = span; - this.scope = span.makeCurrent(); - this.telemetryContext = new TelemetryContext(Context.current()); + this.telemetryContext = new TelemetryContext(parentContext.with(span)); } /** @@ -136,7 +131,7 @@ public void setError(Throwable caughtError) { span.setStatus(StatusCode.ERROR, caughtError.getMessage()); } - /** Closes the scope and ends the underlying span, recording any applicable metrics. */ + /** Ends the underlying span and records any applicable metrics. */ @Override public final void close() { if (closed.getAndSet(true)) { @@ -144,6 +139,7 @@ public final void close() { } try { beforeSpanEnd(); + } finally { span.end(); Duration elapsed = Duration.ofNanos(System.nanoTime() - startTimeNanos); try { @@ -151,8 +147,6 @@ public final void close() { } catch (RuntimeException e) { handleMetricsError(e); } - } finally { - scope.close(); } } @@ -184,7 +178,8 @@ public AgentInvocation(InvocationContext ctx, BaseAgent agent, Context parentCon Tracing.getTracer() .spanBuilder("invoke_agent " + agent.name()) .setParent(parentContext) - .startSpan()); + .startSpan(), + parentContext); this.agent = agent; this.ctx = ctx; Tracing.traceAgentInvocation(span, agent.name(), agent.description(), ctx); @@ -254,7 +249,8 @@ public ToolExecution( Tracing.getTracer() .spanBuilder("execute_tool " + tool.name()) .setParent(parentContext) - .startSpan()); + .startSpan(), + parentContext); this.tool = tool; this.agent = agent; this.functionArgs = functionArgs; diff --git a/core/src/main/java/com/google/adk/telemetry/Tracing.java b/core/src/main/java/com/google/adk/telemetry/Tracing.java index 226d7011c..f284f574f 100644 --- a/core/src/main/java/com/google/adk/telemetry/Tracing.java +++ b/core/src/main/java/com/google/adk/telemetry/Tracing.java @@ -427,37 +427,23 @@ public static Tracer getTracer() { } /** - * Executes a Flowable with an OpenTelemetry Scope active for its entire lifecycle. - * - *

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). - * - *

Why not try-with-resources? 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). - * - *

The scope is properly closed via doFinally when the stream terminates, ensuring no resource - * leaks regardless of completion mode (success, error, or cancellation). + * Executes a {@link Flowable} supplier within {@code spanContext} and propagates that context + * across subscription and stream emissions via {@link #withContext(Context)}. Ends {@code span} + * when the stream terminates or is cancelled. * * @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 The type of items emitted by the Flowable - * @return Flowable with OpenTelemetry scope lifecycle management + * @return Flowable with OpenTelemetry context propagation and span lifecycle management */ - @SuppressWarnings("MustBeClosedChecker") // Scope lifecycle managed by RxJava doFinally public static Flowable traceFlowable( Context spanContext, Span span, Supplier> flowableSupplier) { - Scope scope = spanContext.makeCurrent(); - return flowableSupplier - .get() - .doFinally( - () -> { - scope.close(); - span.end(); - }); + final Flowable upstream; + try (Scope scope = spanContext.makeCurrent()) { + upstream = flowableSupplier.get(); + } + return upstream.compose(withContext(spanContext)).doFinally(span::end); } /** @@ -541,19 +527,16 @@ private Context getParentContext() { private final class TracingLifecycle { private Span span; - private Scope scope; + private Context spanContext; - @SuppressWarnings("MustBeClosedChecker") void start() { - span = tracer.spanBuilder(spanName).setParent(getParentContext()).startSpan(); + Context parentContext = getParentContext(); + span = tracer.spanBuilder(spanName).setParent(parentContext).startSpan(); spanConfigurers.forEach(c -> c.accept(span)); - scope = span.makeCurrent(); + spanContext = parentContext.with(span); } void end() { - if (scope != null) { - scope.close(); - } if (span != null) { span.end(); } @@ -572,7 +555,7 @@ public Publisher apply(Flowable upstream) { () -> { TracingLifecycle lifecycle = new TracingLifecycle(); lifecycle.start(); - Flowable pipeline = upstream; + Flowable pipeline = upstream.compose(withContext(lifecycle.spanContext)); if (onSuccessConsumer != null) { pipeline = pipeline.doOnNext(t -> onSuccessConsumer.accept(lifecycle.span, t)); } @@ -592,7 +575,7 @@ public SingleSource apply(Single upstream) { () -> { TracingLifecycle lifecycle = new TracingLifecycle(); lifecycle.start(); - Single pipeline = upstream; + Single pipeline = upstream.compose(withContext(lifecycle.spanContext)); if (onSuccessConsumer != null) { pipeline = pipeline.doOnSuccess(t -> onSuccessConsumer.accept(lifecycle.span, t)); } @@ -612,7 +595,7 @@ public MaybeSource apply(Maybe upstream) { () -> { TracingLifecycle lifecycle = new TracingLifecycle(); lifecycle.start(); - Maybe pipeline = upstream; + Maybe pipeline = upstream.compose(withContext(lifecycle.spanContext)); if (onSuccessConsumer != null) { pipeline = pipeline.doOnSuccess(t -> onSuccessConsumer.accept(lifecycle.span, t)); } @@ -632,7 +615,7 @@ public CompletableSource apply(Completable upstream) { () -> { TracingLifecycle lifecycle = new TracingLifecycle(); lifecycle.start(); - return upstream.doFinally(lifecycle::end); + return upstream.compose(withContext(lifecycle.spanContext)).doFinally(lifecycle::end); }); } } @@ -673,7 +656,14 @@ private ContextTransformer(Context context) { */ @Override public Publisher apply(Flowable upstream) { - return upstream.lift(subscriber -> TracingObserver.wrap(context, subscriber)); + return new Flowable() { + @Override + protected void subscribeActual(Subscriber subscriber) { + try (Scope scope = context.makeCurrent()) { + upstream.subscribe(TracingObserver.wrap(context, subscriber)); + } + } + }; } /** @@ -684,7 +674,14 @@ public Publisher apply(Flowable upstream) { */ @Override public SingleSource apply(Single upstream) { - return upstream.lift(observer -> TracingObserver.wrap(context, observer)); + return new Single() { + @Override + protected void subscribeActual(SingleObserver observer) { + try (Scope scope = context.makeCurrent()) { + upstream.subscribe(TracingObserver.wrap(context, observer)); + } + } + }; } /** @@ -695,7 +692,14 @@ public SingleSource apply(Single upstream) { */ @Override public MaybeSource apply(Maybe upstream) { - return upstream.lift(observer -> TracingObserver.wrap(context, observer)); + return new Maybe() { + @Override + protected void subscribeActual(MaybeObserver observer) { + try (Scope scope = context.makeCurrent()) { + upstream.subscribe(TracingObserver.wrap(context, observer)); + } + } + }; } /** @@ -706,36 +710,39 @@ public MaybeSource apply(Maybe upstream) { */ @Override public CompletableSource apply(Completable upstream) { - return upstream.lift(observer -> TracingObserver.wrap(context, observer)); + return new Completable() { + @Override + protected void subscribeActual(CompletableObserver observer) { + try (Scope scope = context.makeCurrent()) { + upstream.subscribe(TracingObserver.wrap(context, observer)); + } + } + }; } } /** - * An observer that wraps another observer and ensures that the OpenTelemetry context is active - * during all callback methods. - * - *

This implementation only wraps the data-flow callbacks (`onNext`, `onSuccess`, etc.). The - * `Subscription.request/cancel` and `Disposable.dispose` calls are not wrapped in the context. If - * the upstream logic depends on the context during these signals, they might lose trace - * information. Given this is a manual `withContext` utility, this might be an acceptable - * trade-off for simplicity/performance, but worth keeping in mind. + * Observer wrapper that activates an OpenTelemetry {@link Context} during downstream callbacks + * ({@code onSubscribe}, {@code onNext}, {@code onSuccess}, {@code onError}, {@code onComplete}). + * Upstream flow-control signals ({@code request}, {@code cancel}, {@code dispose}) are not + * wrapped. * * @param The type of the items emitted by the stream. */ private static final class TracingObserver implements Subscriber, SingleObserver, MaybeObserver, CompletableObserver { private final Context context; - private final Subscriber subscriber; - private final SingleObserver singleObserver; - private final MaybeObserver maybeObserver; - private final CompletableObserver completableObserver; + private final @Nullable Subscriber subscriber; + private final @Nullable SingleObserver singleObserver; + private final @Nullable MaybeObserver maybeObserver; + private final @Nullable CompletableObserver completableObserver; private TracingObserver( Context context, - Subscriber subscriber, - SingleObserver singleObserver, - MaybeObserver maybeObserver, - CompletableObserver completableObserver) { + @Nullable Subscriber subscriber, + @Nullable SingleObserver singleObserver, + @Nullable MaybeObserver maybeObserver, + @Nullable CompletableObserver completableObserver) { this.context = context; this.subscriber = subscriber; this.singleObserver = singleObserver; diff --git a/core/src/test/java/com/google/adk/telemetry/ContextPropagationTest.java b/core/src/test/java/com/google/adk/telemetry/ContextPropagationTest.java index 33810f081..e641cab18 100644 --- a/core/src/test/java/com/google/adk/telemetry/ContextPropagationTest.java +++ b/core/src/test/java/com/google/adk/telemetry/ContextPropagationTest.java @@ -57,6 +57,7 @@ import io.reactivex.rxjava3.core.Flowable; import io.reactivex.rxjava3.core.Maybe; import io.reactivex.rxjava3.core.Single; +import io.reactivex.rxjava3.processors.PublishProcessor; import io.reactivex.rxjava3.schedulers.Schedulers; import java.util.Comparator; import java.util.List; @@ -806,6 +807,175 @@ public void testNestedAgentTraceHierarchy() throws InterruptedException { assertParent(agentBSpan, agentBCallLlm); } + @Test + public void + traceFlowable_andTraceTransformer_doNotLeakContextOnCallingThreadDuringAsyncExecution() { + Span callerSpan = tracer.spanBuilder("caller_rpc").startSpan(); + PublishProcessor asyncStream1 = PublishProcessor.create(); + PublishProcessor asyncStream2 = PublishProcessor.create(); + + try (Scope callerScope = callerSpan.makeCurrent()) { + Span flowableSpan = + tracer.spanBuilder("async_flowable").setParent(Context.current()).startSpan(); + Flowable tracedFlowable = + Tracing.traceFlowable( + Context.current().with(flowableSpan), flowableSpan, () -> asyncStream1); + Flowable tracedTransformer = + asyncStream2.compose(Tracing.trace("async_transformer")); + + // Subscribe while streams are still pending (not completed) + var sub1 = tracedFlowable.test(); + var sub2 = tracedTransformer.test(); + + // Calling thread's active span must remain caller_rpc, not async_flowable or + // async_transformer + assertEquals( + callerSpan.getSpanContext().getSpanId(), Span.current().getSpanContext().getSpanId()); + + asyncStream1.onNext(1); + asyncStream1.onComplete(); + asyncStream2.onNext(2); + asyncStream2.onComplete(); + + sub1.assertComplete(); + sub2.assertComplete(); + + assertEquals( + callerSpan.getSpanContext().getSpanId(), Span.current().getSpanContext().getSpanId()); + } finally { + callerSpan.end(); + } + } + + @Test + public void withContext_propagatesContextToDeferredUpstreamAcrossAsyncSubscription() + throws InterruptedException { + ContextKey testKey = ContextKey.named("async-defer-key"); + Context testContext = Context.root().with(testKey, "expected-value"); + + AtomicReference flowableObserved = new AtomicReference<>(); + AtomicReference singleObserved = new AtomicReference<>(); + AtomicReference maybeObserved = new AtomicReference<>(); + AtomicReference completableObserved = new AtomicReference<>(); + + Flowable.defer( + () -> { + flowableObserved.set(Context.current().get(testKey)); + return Flowable.just(1); + }) + .compose(Tracing.withContext(testContext)) + .subscribeOn(Schedulers.computation()) + .test() + .await() + .assertComplete(); + + Single.defer( + () -> { + singleObserved.set(Context.current().get(testKey)); + return Single.just(1); + }) + .compose(Tracing.withContext(testContext)) + .subscribeOn(Schedulers.computation()) + .test() + .await() + .assertComplete(); + + Maybe.defer( + () -> { + maybeObserved.set(Context.current().get(testKey)); + return Maybe.just(1); + }) + .compose(Tracing.withContext(testContext)) + .subscribeOn(Schedulers.computation()) + .test() + .await() + .assertComplete(); + + Completable.defer( + () -> { + completableObserved.set(Context.current().get(testKey)); + return Completable.complete(); + }) + .compose(Tracing.withContext(testContext)) + .subscribeOn(Schedulers.computation()) + .test() + .await() + .assertComplete(); + + assertEquals("expected-value", flowableObserved.get()); + assertEquals("expected-value", singleObserved.get()); + assertEquals("expected-value", maybeObserved.get()); + assertEquals("expected-value", completableObserved.get()); + } + + @Test + public void agentAndToolCallbacks_preserveSpanContextAcrossAsyncBoundaries() + throws InterruptedException { + AtomicReference beforeAgentSpanId = new AtomicReference<>(); + AtomicReference afterAgentSpanId = new AtomicReference<>(); + AtomicReference beforeToolSpanId = new AtomicReference<>(); + AtomicReference afterToolSpanId = new AtomicReference<>(); + + BaseTool asyncTool = + new SearchFlightsTool() { + @Override + public Single> runAsync( + Map args, ToolContext context) { + return Single.>fromCallable(() -> ImmutableMap.of("result", args)) + .subscribeOn(Schedulers.computation()); + } + }; + + TestLlm testLlm = + TestUtils.createTestLlm( + TestUtils.createLlmResponse( + Content.builder() + .role("model") + .parts( + Part.fromFunctionCall( + "search_flights", ImmutableMap.of("destination", "NYC"))) + .build()), + TestUtils.createLlmResponse(Content.fromParts(Part.fromText("flight found")))); + + LlmAgent callbackAgent = + LlmAgent.builder() + .name("callback_agent") + .description("agent with callbacks") + .model(testLlm) + .tools(ImmutableList.of(asyncTool)) + .beforeAgentCallbackSync( + ctx -> { + beforeAgentSpanId.set(Span.current().getSpanContext().getSpanId()); + return Optional.empty(); + }) + .afterAgentCallbackSync( + ctx -> { + afterAgentSpanId.set(Span.current().getSpanContext().getSpanId()); + return Optional.empty(); + }) + .beforeToolCallbackSync( + (invCtx, tool, input, toolCtx) -> { + beforeToolSpanId.set(Span.current().getSpanContext().getSpanId()); + return Optional.empty(); + }) + .afterToolCallbackSync( + (invCtx, tool, input, toolCtx, response) -> { + afterToolSpanId.set(Span.current().getSpanContext().getSpanId()); + return Optional.empty(); + }) + .build(); + + runAgent(callbackAgent); + + SpanData invokeAgentSpan = findSpanByName("invoke_agent callback_agent"); + SpanData executeToolSpan = findSpanByName("execute_tool search_flights"); + + assertEquals(invokeAgentSpan.getSpanContext().getSpanId(), beforeAgentSpanId.get()); + assertEquals(invokeAgentSpan.getSpanContext().getSpanId(), afterAgentSpanId.get()); + assertEquals(executeToolSpan.getSpanContext().getSpanId(), beforeToolSpanId.get()); + assertEquals(executeToolSpan.getSpanContext().getSpanId(), afterToolSpanId.get()); + } + private void runAgent(BaseAgent agent) throws InterruptedException { Runner runner = Runner.builder().agent(agent).appName("test_app").sessionService(sessionService).build(); diff --git a/core/src/test/java/com/google/adk/telemetry/InstrumentationTest.java b/core/src/test/java/com/google/adk/telemetry/InstrumentationTest.java index 3c1ec3269..a28572df4 100644 --- a/core/src/test/java/com/google/adk/telemetry/InstrumentationTest.java +++ b/core/src/test/java/com/google/adk/telemetry/InstrumentationTest.java @@ -17,6 +17,7 @@ package com.google.adk.telemetry; import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertThrows; import com.google.adk.agents.BaseAgent; import com.google.adk.agents.InvocationContext; @@ -32,6 +33,7 @@ import io.opentelemetry.api.common.AttributeKey; import io.opentelemetry.api.common.Attributes; import io.opentelemetry.api.metrics.Meter; +import io.opentelemetry.api.trace.Span; import io.opentelemetry.api.trace.Tracer; import io.opentelemetry.context.Context; import io.opentelemetry.sdk.metrics.data.HistogramPointData; @@ -40,8 +42,11 @@ import io.opentelemetry.sdk.trace.data.SpanData; import io.reactivex.rxjava3.core.Flowable; import io.reactivex.rxjava3.core.Single; +import java.time.Duration; import java.util.List; import java.util.Map; +import java.util.concurrent.atomic.AtomicBoolean; +import org.jspecify.annotations.Nullable; import org.junit.After; import org.junit.Before; import org.junit.Rule; @@ -193,6 +198,36 @@ public void recordToolExecution_success() { .isEqualTo("my-tool"); } + @Test + public void close_whenBeforeSpanEndThrows_stillEndsSpanAndRecordsMetrics() { + AtomicBoolean metricsRecorded = new AtomicBoolean(false); + Span testSpan = Tracing.getTracer().spanBuilder("test_failing_before_span_end").startSpan(); + Instrumentation.ClosableTelemetryScope scope = + new Instrumentation.ClosableTelemetryScope(testSpan, Context.current()) { + @Override + protected void beforeSpanEnd() { + throw new IllegalStateException("simulated serialization failure in beforeSpanEnd"); + } + + @Override + protected void recordMetrics(Duration elapsed, @Nullable Throwable error) { + metricsRecorded.set(true); + } + + @Override + protected void handleMetricsError(RuntimeException e) {} + }; + + IllegalStateException thrown = assertThrows(IllegalStateException.class, scope::close); + assertThat(thrown).hasMessageThat().contains("simulated serialization failure"); + + List spans = openTelemetryRule.getSpans(); + assertThat(spans).hasSize(1); + assertThat(spans.get(0).getName()).isEqualTo("test_failing_before_span_end"); + assertThat(spans.get(0).hasEnded()).isTrue(); + assertThat(metricsRecorded.get()).isTrue(); + } + private MetricData findMetricByName(String name) { return openTelemetryRule.getMetrics().stream() .filter(m -> m.getName().equals(name))