Skip to content

feat: Propogating the otel context#1033

Open
copybara-service[bot] wants to merge 1 commit intomainfrom
test_881463869
Open

feat: Propogating the otel context#1033
copybara-service[bot] wants to merge 1 commit intomainfrom
test_881463869

Conversation

@copybara-service
Copy link

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.

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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant