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.
|
Hi @r4inee, Thank you for taking time and submitting this PR. Currently this PR is under review by our team, we will keep you posted if any additional information is required. thank you. |
|
@tilgalas, Could you please review this. |
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