feat: Propogating the otel context#1033
Open
copybara-service[bot] wants to merge 1 commit intomainfrom
Open
Conversation
3ba4fa8 to
4709af8
Compare
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
4709af8 to
f74f5d2
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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
Scopemanagement (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.RunnerRefactoring:"invocation"span torunAsyncandrunLivecalls.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 (likeafterRunCallbackandonEventCallback) execute within the captured context.BaseLlmFlowTest,FunctionsTest,PluginManagerTest, andRunnerTestthat specifically verify context propagation usingContextKeyandSchedulers.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.BUILDfiles: Updated dependencies for OpenTelemetry APIs and SDK testing.