LLM Observability context lost when using Kitlin coroutines

[grantlerduck]

Tracer Version(s)

1.58.2

Java Version(s)

21.0.7

JVM Vendor

Amazon Corretto

Bug Report

LLM Observability context lost when using Kotlin coroutines

When using LLM Observability with Kotlin coroutines, the LLMObsState context is lost when coroutines suspend and resume on different threads. This breaks parent-child span relationships and can cause spans to have incorrect or missing parent IDs.

Environment

• dd-trace-java version: 1.58.2
• Kotlin with coroutines
• Spring WebFlux / suspend functions

Root Cause

LLMObsState relies on datadog.context.Context which is backed by ThreadLocalContextManager:

// ThreadLocalContextManager.java
final class ThreadLocalContextManager implements ContextManager {
  private static final ThreadLocal<Context[]> CURRENT_HOLDER =
      ThreadLocal.withInitial(() -> new Context[] {EmptyContext.INSTANCE});

  @Override
  public Context current() {
    return CURRENT_HOLDER.get()[0];  // Plain ThreadLocal - not coroutine-safe
  }
}

When DDLLMObsSpan is created, it attaches to the current context and sets the parent:

// DDLLMObsSpan.java constructor
AgentSpanContext parent = LLMObsState.getLLMObsParentContext();  // Returns null on different thread
// ...
this.scope = LLMObsState.attach();
LLMObsState.setLLMObsParentContext(this.span.context());

Impact

This is particularly problematic for:

• Batch LLM operations (parallel embeddings, parallel completions)
• Agent workflows with parallel tool execution
• Any async/concurrent LLM workloads using Kotlin coroutines or reactive frameworks such as Spring WebFlux

Potential Solutions

1. Leverage existing Kotlin coroutine instrumentation - dd-trace-java has coroutine support for APM traces. Could similar context propagation be applied to LLMObsState?

2. Use CoroutineContext elements - Allow users to propagate LLM Obs context explicitly via coroutine context, similar to how MDC context can be propagated.

3. Integrate with datadog.context.Context propagation - If there's existing machinery for propagating Context across async boundaries, ensure LLMObsState benefits from it.

Contribution

I would be happy to contribute a fix for this if the approach is welcomed. Please let me know the preferred direction and any guidance on implementation.

Expected Behavior

Child LLM spans created within coroutines should correctly reference their parent span, maintaining the trace hierarchy.

Actual Behavior

• LLMObsState.getLLMObsParentContext() returns null when called from a different thread
• Child spans get parent_id="undefined" instead of the actual parent span ID
• Trace hierarchy is broken in the Datadog LLM Observability UI

Reproduction Code

Reproduction Scenario

// Parent span created on Thread-1
val parentSpan = LLMObs.startAgentSpan("agent", mlApp, sessionId)

coroutineScope {
    // Child spans created in parallel coroutines
    // These may run on Thread-2, Thread-3, etc.
    listOf(task1, task2, task3).map { task ->
        async([Dispatchers.IO](http://dispatchers.io/)) {
            // LLMObsState.getLLMObsParentContext() returns null here
            // because ThreadLocal from Thread-1 is not available
            val childSpan = LLMObs.startLLMSpan("llm.inference", model, provider, mlApp, sessionId)
            // ... do work
            childSpan.finish()
        }
    }.awaitAll()
}

parentSpan.finish()

[mcculls]

Quick note: context tracking is not backed by ThreadLocalContextManager by default, that's just the fallback if no other manager is installed.

The tracer installs ContinuableScopeManager to track context which supports tracking Kotlin coroutines via the instrumentation in https://github.com/DataDog/dd-trace-java/tree/master/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines - this already uses CoroutineContext

[mcculls]

Wrt. the reproducible scenario - if possible could you provide a zip which can be unpacked, built, and run?

That would help us identify why the context is not being carried over in this specific scenario using LLMObs, compared to our existing coroutine tests.

[grantlerduck]

Thank you for the fast reply.

Okay, then I was confused by the default. My apologies.

Maybe the issues is elsewhere. Let. me provider more context.
I am using Koog which supports parallel tool calling and I am intercepting the llm and tool calls to create the llm spans in the context of the agent run loop which is using asynchronous suspended functions.

The LLMObsState in the same context the APM tracer, right?

[grantlerduck]

I investigated further and observed the following behavior when using LLM Observability with Kotlin coroutines:

  1. With a parent agent span: When I create an agent span that wraps multiple LLM and tool calls, I see "trace ID mismatch"
  errors in the logs. The agent span appears in Datadog with no children, despite LLM and tool spans being created.
  2. First spans work, subsequent don't: The first LLM span and its first tool call are correctly correlated. However,
  subsequent LLM spans and tool calls (especially those running in parallel via async) appear as separate, uncorrelated traces.
  3. Metadata/tags only on first spans: Only the agent span (and sometimes the first LLM span) retain their tags and metadata
  (like tenant_id, prompt_name). Subsequent spans appear with no custom tags.
  4. Parallel execution is key: The issue is most pronounced when tool calls execute in parallel using Kotlin's async/await.
  Sequential execution seems to work better.

  It appears that while the LLMObs parent context is being found correctly (the code attempts to set the parent), the
  underlying APM trace context is not being propagated across coroutine boundaries. This causes new traces to be created
  instead of continuing the existing one, leading to the trace ID mismatch and broken hierarchy.

[ygree]

The problem is that startAgentSpan does not activate the scope; it relies on an existing one. Creating the scope should resolve the issue. The easiest way to do so is to annotate the entire method that calls startAgentSpan with @Trace.