feat: Propogating the otel context

This change ensures that the OpenTelemetry context is correctly propagated across asynchronous boundaries throughout the ADK, primarily within RxJava streams.

### Key Changes

*   **Context Propagation:** Replaces manual `Scope` management (which often fails in reactive code) with `.compose(Tracing.withContext(context))`. This ensures the OTel context is preserved when work moves between different threads or schedulers.
*   **`Runner` Refactoring:**
    *   Adds a top-level `"invocation"` span to `runAsync` and `runLive` calls.
    *   Captures the context at entry points and propagates it through the internal execution flow (`runAsyncImpl`, `runLiveImpl`, `runAgentWithFreshSession`).
*   **`BaseLlmFlow` & `Functions`:** Updates preprocessing, postprocessing, and tool execution logic to maintain context. This ensures that spans created within tools or processors are correctly parented.
*   **`PluginManager`:** Ensures that plugin callbacks (like `afterRunCallback` and `onEventCallback`) execute within the captured context.
*   **Testing:** Adds several unit tests across `BaseLlmFlowTest`, `FunctionsTest`, `PluginManagerTest`, and `RunnerTest` that specifically verify context propagation using `ContextKey` and `Schedulers.computation()`.

### Files Modified
*   **`BaseLlmFlow.java`**, **`Functions.java`**, **`PluginManager.java`**, **`Runner.java`**: Core logic updates for context propagation.
*   **`LlmAgentTest.java`**, **`BaseLlmFlowTest.java`**, **`FunctionsTest.java`**, **`PluginManagerTest.java`**, **`RunnerTest.java`**: New tests for OTel integration.
*   **`BUILD` files**: Updated dependencies for OpenTelemetry APIs and SDK testing.

PiperOrigin-RevId: 881463869

Author: google-genai-bot

core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java

@@ -91,8 +91,9 @@ public BaseLlmFlow(
    * RequestProcessor} transforming the provided {@code llmRequestRef} in-place, and emits the
    * events generated by them.
    */
-  protected Flowable<Event> preprocess(
+  private Flowable<Event> preprocess(
       InvocationContext context, AtomicReference<LlmRequest> llmRequestRef) {
+    Context currentContext = Context.current();
     LlmAgent agent = (LlmAgent) context.agent();
 
     RequestProcessor toolsProcessor =
@@ -114,9 +115,11 @@ protected Flowable<Event> preprocess(
         .concatMap(
             processor ->
                 Single.defer(() -> processor.processRequest(context, llmRequestRef.get()))
+                    .compose(Tracing.withContext(currentContext))
                     .doOnSuccess(result -> llmRequestRef.set(result.updatedRequest()))
                     .flattenAsFlowable(
-                        result -> result.events() != null ? result.events() : ImmutableList.of()));
+                        result -> result.events() != null ? result.events() : ImmutableList.of()))
+        .compose(Tracing.withContext(currentContext));
   }
 
   /**
@@ -146,13 +149,13 @@ protected Flowable<Event> postprocess(
     }
     Context parentContext = Context.current();
 
-    return currentLlmResponse.flatMapPublisher(
-        updatedResponse -> {
-          try (Scope scope = parentContext.makeCurrent()) {
-            return buildPostprocessingEvents(
-                updatedResponse, eventIterables, context, baseEventForLlmResponse, llmRequest);
-          }
-        });
+    return currentLlmResponse
+        .compose(Tracing.withContext(parentContext))
+        .flatMapPublisher(
+            updatedResponse ->
+                buildPostprocessingEvents(
+                    updatedResponse, eventIterables, context, baseEventForLlmResponse, llmRequest))
+        .compose(Tracing.withContext(parentContext));
   }
 
   /**
@@ -222,6 +225,7 @@ private Flowable<LlmResponse> callLlm(
    */
   private Maybe<LlmResponse> handleBeforeModelCallback(
       InvocationContext context, LlmRequest.Builder llmRequestBuilder, Event modelResponseEvent) {
+    Context currentContext = Context.current();
     Event callbackEvent = modelResponseEvent.toBuilder().build();
     CallbackContext callbackContext =
         new CallbackContext(context, callbackEvent.actions(), callbackEvent.id());
@@ -240,7 +244,11 @@ private Maybe<LlmResponse> handleBeforeModelCallback(
         Maybe.defer(
             () ->
                 Flowable.fromIterable(callbacks)
-                    .concatMapMaybe(callback -> callback.call(callbackContext, llmRequestBuilder))
+                    .concatMapMaybe(
+                        callback ->
+                            callback
+                                .call(callbackContext, llmRequestBuilder)
+                                .compose(Tracing.withContext(currentContext)))
                     .firstElement());
 
     return pluginResult.switchIfEmpty(callbackResult);
@@ -257,6 +265,7 @@ private Maybe<LlmResponse> handleOnModelErrorCallback(
       LlmRequest.Builder llmRequestBuilder,
       Event modelResponseEvent,
       Throwable throwable) {
+    Context currentContext = Context.current();
     Event callbackEvent = modelResponseEvent.toBuilder().build();
     CallbackContext callbackContext =
         new CallbackContext(context, callbackEvent.actions(), callbackEvent.id());
@@ -277,7 +286,11 @@ private Maybe<LlmResponse> handleOnModelErrorCallback(
             () -> {
               LlmRequest llmRequest = llmRequestBuilder.build();
               return Flowable.fromIterable(callbacks)
-                  .concatMapMaybe(callback -> callback.call(callbackContext, llmRequest, ex))
+                  .concatMapMaybe(
+                      callback ->
+                          callback
+                              .call(callbackContext, llmRequest, ex)
+                              .compose(Tracing.withContext(currentContext)))
                   .firstElement();
             });
 
@@ -292,6 +305,7 @@ private Maybe<LlmResponse> handleOnModelErrorCallback(
    */
   private Single<LlmResponse> handleAfterModelCallback(
       InvocationContext context, LlmResponse llmResponse, Event modelResponseEvent) {
+    Context currentContext = Context.current();
     Event callbackEvent = modelResponseEvent.toBuilder().build();
     CallbackContext callbackContext =
         new CallbackContext(context, callbackEvent.actions(), callbackEvent.id());
@@ -310,7 +324,11 @@ private Single<LlmResponse> handleAfterModelCallback(
         Maybe.defer(
             () ->
                 Flowable.fromIterable(callbacks)
-                    .concatMapMaybe(callback -> callback.call(callbackContext, llmResponse))
+                    .concatMapMaybe(
+                        callback ->
+                            callback
+                                .call(callbackContext, llmResponse)
+                                .compose(Tracing.withContext(currentContext)))
                     .firstElement());
 
     return pluginResult.switchIfEmpty(callbackResult).defaultIfEmpty(llmResponse);

core/src/main/java/com/google/adk/flows/llmflows/Functions.java

@@ -42,7 +42,6 @@
 import com.google.genai.types.Part;
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.context.Context;
-import io.opentelemetry.context.Scope;
 import io.reactivex.rxjava3.core.Flowable;
 import io.reactivex.rxjava3.core.Maybe;
 import io.reactivex.rxjava3.core.Observable;
@@ -163,7 +162,9 @@ public static Maybe<Event> handleFunctionCalls(
     }
     return functionResponseEventsObservable
         .toList()
-        .flatMapMaybe(
+        .toMaybe()
+        .compose(Tracing.withContext(parentContext))
+        .flatMap(
             events -> {
               if (events.isEmpty()) {
                 return Maybe.empty();
@@ -226,7 +227,9 @@ public static Maybe<Event> handleFunctionCallsLive(
 
     return responseEventsObservable
         .toList()
-        .flatMapMaybe(
+        .toMaybe()
+        .compose(Tracing.withContext(parentContext))
+        .flatMap(
             events -> {
               if (events.isEmpty()) {
                 return Maybe.empty();
@@ -243,47 +246,45 @@ private static Function<FunctionCall, Maybe<Event>> getFunctionCallMapper(
       Context parentContext) {
     return functionCall ->
         Maybe.defer(
-            () -> {
-              try (Scope scope = parentContext.makeCurrent()) {
-                BaseTool tool = tools.get(functionCall.name().get());
-                ToolContext toolContext =
-                    ToolContext.builder(invocationContext)
-                        .functionCallId(functionCall.id().orElse(""))
-                        .toolConfirmation(
-                            functionCall.id().map(toolConfirmations::get).orElse(null))
-                        .build();
-
-                Map<String, Object> functionArgs =
-                    functionCall.args().map(HashMap::new).orElse(new HashMap<>());
-
-                Maybe<Map<String, Object>> maybeFunctionResult =
-                    maybeInvokeBeforeToolCall(invocationContext, tool, functionArgs, toolContext)
-                        .switchIfEmpty(
-                            Maybe.defer(
-                                () -> {
-                                  try (Scope innerScope = parentContext.makeCurrent()) {
-                                    return isLive
-                                        ? processFunctionLive(
-                                            invocationContext,
-                                            tool,
-                                            toolContext,
-                                            functionCall,
-                                            functionArgs,
-                                            parentContext)
-                                        : callTool(tool, functionArgs, toolContext, parentContext);
-                                  }
-                                }));
-
-                return postProcessFunctionResult(
-                    maybeFunctionResult,
-                    invocationContext,
-                    tool,
-                    functionArgs,
-                    toolContext,
-                    isLive,
-                    parentContext);
-              }
-            });
+                () -> {
+                  BaseTool tool = tools.get(functionCall.name().get());
+                  ToolContext toolContext =
+                      ToolContext.builder(invocationContext)
+                          .functionCallId(functionCall.id().orElse(""))
+                          .toolConfirmation(
+                              functionCall.id().map(toolConfirmations::get).orElse(null))
+                          .build();
+
+                  Map<String, Object> functionArgs =
+                      functionCall.args().map(HashMap::new).orElse(new HashMap<>());
+
+                  Maybe<Map<String, Object>> maybeFunctionResult =
+                      maybeInvokeBeforeToolCall(invocationContext, tool, functionArgs, toolContext)
+                          .switchIfEmpty(
+                              Maybe.defer(
+                                      () ->
+                                          isLive
+                                              ? processFunctionLive(
+                                                  invocationContext,
+                                                  tool,
+                                                  toolContext,
+                                                  functionCall,
+                                                  functionArgs,
+                                                  parentContext)
+                                              : callTool(
+                                                  tool, functionArgs, toolContext, parentContext))
+                                  .compose(Tracing.withContext(parentContext)));
+
+                  return postProcessFunctionResult(
+                      maybeFunctionResult,
+                      invocationContext,
+                      tool,
+                      functionArgs,
+                      toolContext,
+                      isLive,
+                      parentContext);
+                })
+            .compose(Tracing.withContext(parentContext));
   }
 
   /**
@@ -410,34 +411,27 @@ private static Maybe<Event> postProcessFunctionResult(
             })
         .flatMapMaybe(
             optionalInitialResult -> {
-              try (Scope scope = parentContext.makeCurrent()) {
-                Map<String, Object> initialFunctionResult = optionalInitialResult.orElse(null);
-
-                return maybeInvokeAfterToolCall(
-                        invocationContext, tool, functionArgs, toolContext, initialFunctionResult)
-                    .map(Optional::of)
-                    .defaultIfEmpty(Optional.ofNullable(initialFunctionResult))
-                    .flatMapMaybe(
-                        finalOptionalResult -> {
-                          Map<String, Object> finalFunctionResult =
-                              finalOptionalResult.orElse(null);
-                          if (tool.longRunning() && finalFunctionResult == null) {
-                            return Maybe.empty();
-                          }
-                          return Maybe.fromCallable(
-                                  () ->
-                                      buildResponseEvent(
-                                          tool,
-                                          finalFunctionResult,
-                                          toolContext,
-                                          invocationContext))
-                              .compose(
-                                  Tracing.<Event>trace("tool_response [" + tool.name() + "]")
-                                      .setParent(parentContext))
-                              .doOnSuccess(event -> Tracing.traceToolResponse(event.id(), event));
-                        });
-              }
-            });
+              Map<String, Object> initialFunctionResult = optionalInitialResult.orElse(null);
+
+              return maybeInvokeAfterToolCall(
+                      invocationContext, tool, functionArgs, toolContext, initialFunctionResult)
+                  .map(Optional::of)
+                  .defaultIfEmpty(Optional.ofNullable(initialFunctionResult))
+                  .flatMapMaybe(
+                      finalOptionalResult -> {
+                        Map<String, Object> finalFunctionResult = finalOptionalResult.orElse(null);
+                        if (tool.longRunning() && finalFunctionResult == null) {
+                          return Maybe.empty();
+                        }
+                        Event event =
+                            buildResponseEvent(
+                                tool, finalFunctionResult, toolContext, invocationContext);
+                        Tracing.traceToolResponse(event.id(), event);
+                        return Maybe.just(event);
+                      });
+            })
+        .compose(
+            Tracing.<Event>trace("tool_response [" + tool.name() + "]").setParent(parentContext));
   }
 
   private static Optional<Event> mergeParallelFunctionResponseEvents(

core/src/main/java/com/google/adk/plugins/PluginManager.java

@@ -21,11 +21,13 @@
 import com.google.adk.events.Event;
 import com.google.adk.models.LlmRequest;
 import com.google.adk.models.LlmResponse;
+import com.google.adk.telemetry.Tracing;
 import com.google.adk.tools.BaseTool;
 import com.google.adk.tools.ToolContext;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableList;
 import com.google.genai.types.Content;
+import io.opentelemetry.context.Context;
 import io.reactivex.rxjava3.core.Completable;
 import io.reactivex.rxjava3.core.Flowable;
 import io.reactivex.rxjava3.core.Maybe;
@@ -126,6 +128,7 @@ public Maybe<Content> beforeRunCallback(InvocationContext invocationContext) {
 
   @Override
   public Completable afterRunCallback(InvocationContext invocationContext) {
+    Context capturedContext = Context.current();
     return Flowable.fromIterable(plugins)
         .concatMapCompletable(
             plugin ->
@@ -136,20 +139,22 @@ public Completable afterRunCallback(InvocationContext invocationContext) {
                             logger.error(
                                 "[{}] Error during callback 'afterRunCallback'",
                                 plugin.getName(),
-                                e)));
+                                e))
+                    .compose(Tracing.withContext(capturedContext)));
   }
 
   @Override
   public Completable close() {
+    Context capturedContext = Context.current();
     return Flowable.fromIterable(plugins)
         .concatMapCompletableDelayError(
             plugin ->
                 plugin
                     .close()
                     .doOnError(
                         e ->
-                            logger.error(
-                                "[{}] Error during callback 'close'", plugin.getName(), e)));
+                            logger.error("[{}] Error during callback 'close'", plugin.getName(), e))
+                    .compose(Tracing.withContext(capturedContext)));
   }
 
   @Override
@@ -227,7 +232,7 @@ public Maybe<Map<String, Object>> onToolErrorCallback(
    */
   private <T> Maybe<T> runMaybeCallbacks(
       Function<Plugin, Maybe<T>> callbackExecutor, String callbackName) {
-
+    Context capturedContext = Context.current();
     return Flowable.fromIterable(this.plugins)
         .concatMapMaybe(
             plugin ->
@@ -246,7 +251,8 @@ private <T> Maybe<T> runMaybeCallbacks(
                                 "[{}] Error during callback '{}'",
                                 plugin.getName(),
                                 callbackName,
-                                e)))
+                                e))
+                    .compose(Tracing.<T>withContext(capturedContext)))
         .firstElement();
   }
 }