Trace context lost in RxJava-based async flow (Flowable)

[dongsu90]

Describe the bug
In ADK-Java, when using Flowable for asynchronous execution in runner.java and BaseAgent.java, the OpenTelemetry trace context is not properly propagated. As a result, the trace tree in the OpenTelemetry backend does not show the expected parent-child span relationships. This breaks the trace continuity and makes debugging and performance analysis difficult.

To Reproduce
Steps to reproduce the behavior:

1. Clone the ADK-Java repository and set up a sample agent.
2. Start the AdkWebServer.
3. Observe the trace output in AdkWeb UI.
4. Notice that spans created inside flatMap do not appear as children of the initial span.

Expected behavior
Spans created inside the Flowable chain (e.g., in BaseAgent.runAsync) should maintain the parent-child relationship with the span created in runner.runAsync. The trace tree should reflect the full execution flow, preserving context across asynchronous boundaries.

Screenshots
the trace tree showing disconnected spans.

Desktop (please complete the following information):

• OS: ubuntu-24
• Java 17
• adk-java v0.2.0

Additional context
This issue appears to be related to RxJava's handling of thread-local context. OpenTelemetry context is not automatically propagated across asynchronous boundaries in RxJava unless explicitly wrapped.

[ruslan-kurchenko]

Detailed Analysis and Reproduction Case

I'm experiencing this exact issue in a production Micronaut/Kotlin application using ADK Java 0.3.0 with OpenTelemetry 1.51.0. I've spent considerable time debugging and can provide detailed insights into the root cause and impact.

Environment

• ADK Version: 0.3.0
• OpenTelemetry: 1.51.0 (via Micronaut Tracing)
• Framework: Micronaut 4.8.3 with Kotlin 2.1.0
• RxJava: 3.x (kotlinx-coroutines-rx3 bridge)
• Observability: Jaeger (OTLP exporter)

Reproduction

Scenario: Multi-agent ADK application with FunctionTools that perform external service calls (GraphQL, vector search). Each tool execution should appear as a child span of the ADK invocation span.

What Happens:

1. Main HTTP request creates root span (POST /v1/chat/completions)
2. ChatCompletionsService creates child span (chat.completions.service)
3. ADK execution creates child span (adk.execution)
4. ADK creates invocation span as child of adk.execution ✓
5. When ADK invokes a tool, it creates tool_call [toolName] span as NEW ROOT ✗
6. Our service-level spans (knowledge.search, graphql.findSlots) are also separate root spans ✗

Result: 3+ disconnected traces instead of 1 unified trace.

Root Cause Analysis

After extensive investigation, I've identified the exact issue in ADK source code:

File: core/src/main/java/com/google/adk/flows/llmflows/Functions.java:497

// CURRENT CODE (BROKEN)
Span span = Telemetry.getTracer()
    .spanBuilder("tool_call [" + tool.name() + "]")
    .startSpan();  // ← Creates ROOT span, ignores Context.current()

try (Scope scope = span.makeCurrent()) {
    // Tool execution
}

The Bug:

• startSpan() without .setParent() creates a root span by default
• Even though Context.current() is available (we verified this), ADK doesn't use it
• Same issue in:
  • Functions.java:518 (tool_response span)
  • BaseLlmFlow.java:217 (call_llm span)
  • Potentially other span creation sites

What SHOULD Happen:

// FIXED CODE
Span span = Telemetry.getTracer()
    .spanBuilder("tool_call [" + tool.name() + "]")
    .setParent(Context.current())  // ← Use current context as parent!
    .startSpan();

Impact

Without Fix:

• Cannot measure end-to-end request latency accurately
• Tool execution timing appears in separate traces
• Debugging performance issues requires manual trace correlation
• Distributed tracing across services breaks at ADK boundaries

With Fix:

• Complete hierarchical trace showing:
  • HTTP request → ADK execution → Tool calls → External services
  • Accurate timing for each operation
  • Proper parent-child span relationships
  • Full distributed trace across all hops

Attempted Workarounds

I've tried several approaches, all unsuccessful due to ADK's span creation not respecting context:

1. OpenTelemetry Kotlin Extension (asContextElement())

   • Added to all coroutine dispatcher switches
   • Works for our service-level spans
   • Doesn't affect ADK's internal span creation

2. RxJava Context Propagation

   • Captured context before runAsync()
   • Used Context.makeCurrent() in doOnSubscribe
   • Maintained scope through doFinally
   • Context IS available, but ADK doesn't use it

3. GlobalOpenTelemetry Configuration

   • Enabled global autoconfigure
   • Set OTEL_SERVICE_NAME to unify service naming
   • All spans now have same service name
   • But still separate traces (different trace IDs)

4. Tool Wrapper Context Injection

   • Wrapped all FunctionTool calls with context propagation
   • Ensured Context.current() is set before tool execution
   • ADK still creates root spans

Evidence from Telemetry.java

The Telemetry.java class HAS methods that check Span.current():

public static void traceToolCall(Map<String, Object> args) {
    Span span = Span.current();  // ← This WORKS when context is set
    if (span == null || !span.getSpanContext().isValid()) {
        log.trace("traceToolCall: No valid span in current context.");
        return;
    }
    // ... adds attributes to current span
}

This proves Context.current() IS available during tool execution. The problem is that span CREATION doesn't use it.

Proposed Fix

Minimal Change Required:

In all span creation sites where child relationships are expected, change:

Span span = tracer.spanBuilder(spanName).startSpan();

To:

Span span = tracer.spanBuilder(spanName)
    .setParent(Context.current())
    .startSpan();

Files to Modify:

1. core/src/main/java/com/google/adk/flows/llmflows/Functions.java

   • Line 497: tool_call span creation
   • Line 518: tool_response span creation
   • Line 201: Merged tool_response span

2. core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java

   • Line 217: call_llm span creation
   • Line 474: send_data span creation

3. Potentially core/src/main/java/com/google/adk/agents/BaseAgent.java

   • Line 208, 343: agent_run spans (check if these need parent linking)

Test Case

After fix, running this code should produce a single unified trace:

val span = tracer.spanBuilder("parent").startSpan()
span.makeCurrent().use {
    runner.runAsync(session, userMessage, runConfig)
        .subscribe { event ->
            // ADK creates tool_call spans here
            // Should be children of "parent" span
        }
}

Expected Jaeger Output:

parent [total time]
└── invocation
    ├── tool_call [toolName] [tool time] ← Child, not root!
    │   └── tool_response
    └── call_llm [llm time]

Request

Could the ADK team prioritize this fix? It significantly impacts observability for production deployments. The fix appears straightforward (adding .setParent(Context.current())) and would greatly improve the developer experience with ADK.

Happy to provide additional details, test cases, or collaborate on a PR if helpful!

[Poggecci]

Thank you for investigating this and sending a PR @ruslan-kurchenko ! I've replied there with a small request before I can get this checked in. Please address and then I'll happily merge the changes should everything be passing👍

[ruslan-kurchenko]

Thank you, @Poggecci. I have updated the code as you suggested.

Just curious, how often are patches like these released?

[Poggecci]

We are planning a 0.3.1 release this week, but haven't settled on a timeline for continuous releases.