Skip to content

fix: Fix RxJava tracing context propagation#1253

Open
r4inee wants to merge 1 commit into
google:mainfrom
r4inee:ssu/fix-rxjava-tracing-context
Open

fix: Fix RxJava tracing context propagation#1253
r4inee wants to merge 1 commit into
google:mainfrom
r4inee:ssu/fix-rxjava-tracing-context

Conversation

@r4inee

@r4inee r4inee commented Jun 9, 2026

Copy link
Copy Markdown
Contributor

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:

Span.current().setAttribute(...);

Solution:

  • 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().

  • Runner#runAsyncImpl — Capture Context.current() inside Flowable.defer(...) so the runner captures execution-time context, not stale assembly-time context.

  • BaseAgent#run — Capture the parent context inside deferred execution so invoke_agent is correctly parented under the active invocation span.

  • BaseLlmFlow#callLlm — Create the call_llm span before beforeModelCallback, so beforeModelCallback, model execution, afterModelCallback, and model error callbacks share the same call_llm span.

  • 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.

Testing Plan

  • ContextPropagationTest — Add coverage proving Tracing.trace(...) starts spans before deferred upstream code runs.
  • Existing coverage - ContextPropagationTest#testAgentWithToolCallTraceHierarchy and ContextPropagationTest#runnerRunAsync_propagatesContext which checks context propgation and hierarchy. Added testModelCallbacksObserveCallLlmSpan to check that the callbacks sees the correct current span.

Unit Tests:

  • I have added or updated unit tests for my change.
  • All unit tests pass locally.

Manual End-to-End (E2E) Tests:

Hierarchy remains correct
Screenshot 2026-06-09 at 10 41 40 pm

With context propgation
Screenshot 2026-06-09 at 10 42 06 pm

Checklist

  • I have read the CONTRIBUTING.md document.
  • My pull request contains a single commit.
  • I have performed a self-review of my own code.
  • I have commented my code, particularly in hard-to-understand areas.
  • I have added tests that prove my fix is effective or that my feature works.
  • New and existing unit tests pass locally with my changes.
  • I have manually tested my changes end-to-end.
  • Any dependent changes have been merged and published in downstream modules.

@r4inee r4inee force-pushed the ssu/fix-rxjava-tracing-context branch from 334fc81 to 3de1916 Compare June 9, 2026 12:54
@r4inee

r4inee commented Jun 9, 2026

Copy link
Copy Markdown
Contributor Author

cc @krwc @mazas-google

@r4inee r4inee changed the title Fix RxJava tracing context propagation fix: Fix RxJava tracing context propagation Jun 9, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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().

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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),
() -> {

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

BaseAgent#run — Capture the parent context inside deferred execution so invoke_agent is correctly parented under the active invocation span.

@r4inee r4inee Jun 9, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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));

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

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.

Telemetry: Plugins observe incorrect current span

1 participant