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..5a54694ce 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)); } /** @@ -142,17 +137,13 @@ public final void close() { if (closed.getAndSet(true)) { return; } + beforeSpanEnd(); + span.end(); + Duration elapsed = Duration.ofNanos(System.nanoTime() - startTimeNanos); try { - beforeSpanEnd(); - span.end(); - Duration elapsed = Duration.ofNanos(System.nanoTime() - startTimeNanos); - try { - recordMetrics(elapsed, caughtError); - } catch (RuntimeException e) { - handleMetricsError(e); - } - } finally { - scope.close(); + recordMetrics(elapsed, caughtError); + } catch (RuntimeException e) { + handleMetricsError(e); } } @@ -184,7 +175,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 +246,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;