fix: Fix RxJava tracing context propagation#1253
Conversation
334fc81 to
3de1916
Compare
There was a problem hiding this comment.
Tracing.TracerProvider — Start spans before subscribing to upstream streams, instead of in doOnSubscribe.
This ensures deferred RxJava sources see the tracing span as Span.current().
There was a problem hiding this comment.
Runner#runAsyncImpl — Capture Context.current() inside Flowable.defer(...) so the runner captures execution-time context, not stale assembly-time context.
| () -> | ||
| Instrumentation.recordAgentInvocation( | ||
| createInvocationContext(parentContext), this, otelContext), | ||
| () -> { |
There was a problem hiding this comment.
BaseAgent#run — Capture the parent context inside deferred execution so invoke_agent is correctly parented under the active invocation span.
There was a problem hiding this comment.
BaseLlmFlow#callLlm — Create the call_llm span before beforeModelCallback, so beforeModelCallback, model execution, afterModelCallback, and model error callbacks share the same call_llm span.
| () -> | ||
| receiveFlow.takeWhile( | ||
| event -> !event.actions().endInvocation().orElse(false))) | ||
| .compose(Tracing.withContext(spanContext)); |
There was a problem hiding this comment.
BaseLlmFlow#callLlm — Re-scope emissions back to the parent context after the LLM segment, preventing follow-up tool/agent/LLM work from inheriting the previous call_llm context.
Link to Issue or Description of Change
1. Link to an existing issue (if applicable):
Problem:
ADK Java plugin callbacks are not consistently invoked with the semantically correct OpenTelemetry span as Span.current(). This affects plugin authors who enrich ADK-created spans using:
Solution:
Tracing.TracerProvider— Start spans before subscribing to upstream streams, instead of indoOnSubscribe. This ensures deferred RxJava sources see the tracing span asSpan.current().Runner#runAsyncImpl— CaptureContext.current()insideFlowable.defer(...)so the runner captures execution-time context, not stale assembly-time context.BaseAgent#run— Capture the parent context inside deferred execution soinvoke_agentis correctly parented under the activeinvocationspan.BaseLlmFlow#callLlm— Create thecall_llmspan beforebeforeModelCallback, sobeforeModelCallback, model execution,afterModelCallback, and model error callbacks share the samecall_llmspan.BaseLlmFlow#callLlm— Re-scope emissions back to the parent context after the LLM segment, preventing follow-up tool/agent/LLM work from inheriting the previouscall_llmcontext.Testing Plan
ContextPropagationTest— Add coverage provingTracing.trace(...)starts spans before deferred upstream code runs.ContextPropagationTest#testAgentWithToolCallTraceHierarchyandContextPropagationTest#runnerRunAsync_propagatesContextwhich checks context propgation and hierarchy. AddedtestModelCallbacksObserveCallLlmSpanto check that the callbacks sees the correct current span.Unit Tests:
Manual End-to-End (E2E) Tests:
Hierarchy remains correct

With context propgation

Checklist