feat: Propagating 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

@@ -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/flows/llmflows/Functions.java

@@ -29,6 +29,7 @@
 import com.google.adk.events.Event;
 import com.google.adk.events.ToolConfirmation;
 import com.google.adk.models.LlmRequest;
+import com.google.adk.telemetry.Tracing;
 import com.google.adk.tools.BaseTool;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
@@ -37,6 +38,7 @@
 import com.google.genai.types.FunctionCall;
 import com.google.genai.types.FunctionResponse;
 import com.google.genai.types.Part;
+import io.opentelemetry.context.Context;
 import io.reactivex.rxjava3.core.Maybe;
 import io.reactivex.rxjava3.core.Single;
 import java.util.Collection;
@@ -216,10 +218,13 @@ private Maybe<Event> assembleEvent(
                     .build())
             .build();
 
-    return toolsMapSingle.flatMapMaybe(
-        toolsMap ->
-            Functions.handleFunctionCalls(
-                invocationContext, functionCallEvent, toolsMap, toolConfirmations));
+    Context parentContext = Context.current();
+    return toolsMapSingle
+        .flatMapMaybe(
+            toolsMap ->
+                Functions.handleFunctionCalls(
+                    invocationContext, functionCallEvent, toolsMap, toolConfirmations))
+        .compose(Tracing.withContext(parentContext));
   }
 
   private static Optional<Map.Entry<String, ToolConfirmation>> maybeCreateToolConfirmationEntry(

core/src/main/java/com/google/adk/flows/llmflows/RequestConfirmationLlmRequestProcessor.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,11 +139,13 @@ 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 ->
@@ -149,7 +154,8 @@ public Completable close() {
                     .doOnError(
                         e ->
                             logger.error(
-                                "[{}] Error during callback 'close'", plugin.getName(), e)));
+                                "[{}] Error during callback 'close'", plugin.getName(), e)))
+        .compose(Tracing.withContext(capturedContext));
   }
 
   @Override
@@ -227,7 +233,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 ->
@@ -247,6 +253,7 @@ private <T> Maybe<T> runMaybeCallbacks(
                                 plugin.getName(),
                                 callbackName,
                                 e)))
-        .firstElement();
+        .firstElement()
+        .compose(Tracing.withContext(capturedContext));
   }
 }