Skip to content

feat(llmobs): add agent attribution to Java SDK - #12238

Open
yahya-mouman wants to merge 27 commits into
masterfrom
yahya/llmobs-agent-attribution
Open

feat(llmobs): add agent attribution to Java SDK#12238
yahya-mouman wants to merge 27 commits into
masterfrom
yahya/llmobs-agent-attribution

Conversation

@yahya-mouman

@yahya-mouman yahya-mouman commented Aug 19, 2026

Copy link
Copy Markdown
Contributor

Summary

Agent attribution tracks the nearest `agent`-kind ancestor span for every LLMObs span and emits it as `meta.agent_attribution = {"pagent_name": str|null, "pagent_span_id": str}`. This feature is already merged in Python (#18788), JavaScript (#9175), and Go. This PR brings the Java SDK to parity.

Wire contract

  • Serialized field: `meta.agent_attribution = {"pagent_name": str|null, "pagent_span_id": str}`

Two-case in-process resolution (O(1) at span start)

  1. This span is `kind=agent` → self-attribution (own span ID + name)
  2. In-process LLMObs parent → inherit from `LLMObsContext` thread-local context

A trace-ID consistency gate prevents stale `LLMObsContext` values leaking across async boundaries: attribution is only inherited when the context's trace ID matches the new span's trace ID.

Changes

File Change
`internal-api/.../LLMObsContext.java` Two new `ContextKey` fields + 4-arg `attach()` overload for pagent propagation
`dd-trace-api/.../LLMObsTags.java` Two new constants: `PAGENT_SPAN_ID` and `PAGENT_NAME`
`dd-trace-core/.../LLMObsSpanMapper.java` Serializes two internal pagent tags as a single `agent_attribution` map in meta
`agent-llmobs/.../DDLLMObsSpan.java` Resolution logic in constructor; wire-safe agent name validation (char-scan, no allocation)
`openai-java/.../CommonTags.java` Two new tag constants for pagent propagation on auto-instrumented OpenAI spans
`openai-java/.../OpenAiDecorator.java` Inherit pagent attribution from active `LLMObsContext` during `doAfterStart`
`agent-llmobs/.../DDLLMObsSpanAgentAttributionTest.java` New JUnit 5 unit tests: self-attribution, unsafe name, tilde rejection, in-process inheritance, transitive, no-ancestor, stale-context gate
`agent-llmobs/.../AgentAttributionIntegrationTest.java` New JUnit 5 integration tests through the public `LLMObs.start*Span()` API; includes OpenAI decorator context simulation

Test plan

  • `./gradlew :dd-java-agent:agent-llmobs:test` — `DDLLMObsSpanAgentAttributionTest` and `AgentAttributionIntegrationTest` cover all resolution cases
  • `./gradlew :dd-trace-core:test` — mapper tests verify `agent_attribution` wire format including explicit-null name path
  • `./gradlew spotlessCheck` — format check
  • Manual: create an agent → llm → tool chain and verify `meta.agent_attribution` appears on tool and llm spans pointing to the agent

🤖 Generated with Claude Code

@yahya-mouman yahya-mouman added type: feature Enhancements and improvements comp: core Tracer core labels Aug 19, 2026
@dd-octo-sts

dd-octo-sts Bot commented Aug 19, 2026

Copy link
Copy Markdown
Contributor

🟢 Java Benchmark SLOs — All performance SLOs passed

Suite Status
Startup 🟢 pass

SLO thresholds are defined here based on automatically generated metrics. A warning is raised when results are within 5% of the threshold.

PR vs. master results
Scenario Candidate master Δ (95% CI of mean)
startup:insecure-bank:iast:Agent 14.80 s 14.71 s [-0.4%; +1.6%] (no difference)
startup:insecure-bank:tracing:Agent 13.66 s 13.59 s [-0.4%; +1.4%] (no difference)
startup:petclinic:appsec:Agent 17.00 s 17.00 s [-1.0%; +1.0%] (no difference)
startup:petclinic:iast:Agent 17.04 s 17.10 s [-1.1%; +0.4%] (no difference)
startup:petclinic:profiling:Agent 16.79 s 17.10 s [-2.9%; -0.9%] (maybe better)
startup:petclinic:sca:Agent 16.91 s 16.60 s [+1.0%; +2.9%] (maybe worse)
startup:petclinic:tracing:Agent 16.20 s 16.13 s [-0.6%; +1.4%] (no difference)

Commit: df8a0678 · CI Pipeline · Benchmarking Platform UI


Load and DaCapo benchmarks can be triggered manually in the GitLab pipeline. Results will appear in the Benchmarking Platform UI after completion.

@datadog-prod-us1-5

This comment has been minimized.

@yahya-mouman

Copy link
Copy Markdown
Contributor Author

@codex

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 65fb79c534

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

Comment thread dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/PTagsCodec.java Outdated
@yahya-mouman

yahya-mouman commented Aug 25, 2026

Copy link
Copy Markdown
Contributor Author

Deferred follow-up items (from Codex review)

Three issues flagged as P1 in the Codex review are intentionally deferred to separate PRs to keep this one reviewable:

1. W3C incoming decodeW3CPTagsCodec.fromHeaderValue does not extract llmobs_pagent_span_id / llmobs_pagent_name into the named fields on PTags. The Datadog-format codec handles this; W3C decode is a separate PR.

2. Auto-instrumented LLM spansOpenAiDecorator and similar auto-instrumentation decorators don't set pagent tags. Attribution only works for manually-created DDLLMObsSpans in this PR. Follow-up will plumb pagent through the decorators.

3. Standalone agent without APM parent — When a kind=agent DDLLMObsSpan is the first span in a process (no APM root span), span.getLocalRootSpan() is the span itself and rootCtx is its own context, so stamping PTags works but the span cannot receive distributed pagent from upstream. This edge case is low-priority and deferred.

4. #12347

@yahya-mouman
yahya-mouman force-pushed the yahya/llmobs-agent-attribution branch from 2a9e133 to 4fbafd9 Compare August 26, 2026 13:41
@yahya-mouman

Copy link
Copy Markdown
Contributor Author

@codex

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 4fbafd9c86

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

Comment thread internal-api/src/main/java/datadog/trace/api/llmobs/LLMObsContext.java Outdated

@sabrenner sabrenner left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

lgtm overall!

yahya-mouman and others added 17 commits August 31, 2026 16:47
Extends agent attribution (pagent_name/pagent_span_id) to dd-trace-java,
matching the existing implementation in dd-trace-py, dd-trace-js, and dd-trace-go.

Every LLMObs span now carries meta.agent_attribution = {pagent_name, pagent_span_id}
identifying its nearest agent-kind ancestor, resolved O(1) at span start.

Changes:
- LLMObsPropagationAccess: new bridge interface in internal-api allowing agent-llmobs
  to read/write _dd.p.llmobs_pagent_* propagation tags on the APM span context
  without a direct dd-trace-core dependency
- LLMObsContext: adds PAGENT_SPAN_ID_KEY/PAGENT_NAME_KEY context keys and extended
  attach() overload so agent attribution propagates in-process to descendants
- PropagationTags + PTags: adds getParentAgentSpanId/Name and updateParentAgentSpanId/Name
  with volatile TagValue fields and header cache invalidation
- PTagsCodec: defines PARENT_AGENT_SPAN_ID_TAG/PARENT_AGENT_NAME_TAG constants,
  emits both in headerValue() and fillTagMap()
- DatadogPTagsCodec: extracts _dd.p.llmobs_pagent_* from incoming x-datadog-tags header
- DDSpanContext: implements LLMObsPropagationAccess by delegating to getPropagationTags()
- DDLLMObsSpan: resolves attribution at span start (agent spans write themselves;
  non-agent spans inherit from context; distributed case reads from root span PTags);
  wire-safe validation for agent names (printable ASCII, no commas/semicolons, ≤256 bytes)
- LLMObsSpanMapper: serializes agent_attribution as a structured sub-map in meta,
  emitting pagent_name as explicit null when name was dropped

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Tests cover the three resolution cases in DDLLMObsSpan (self-as-agent,
in-process context inheritance, distributed propagation) and the
serializer's agent_attribution block emission in LLMObsSpanMapper,
including the explicit-null name path when only pagent_span_id is set.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…guard

Fix metaSize to subtract 1 whenever pagent_name is in tagsToRemapToMeta
(not only when both pagent fields are present) to prevent an off-by-one
if name is set without span_id, making the formula symmetric with the
actual skip logic in the serializer loop.

Replace getBytes(UTF_8) in agentNameWireSafe with a char-by-char scan,
eliminating the per-agent-span byte array allocation. Because the loop
rejects c > 0x7E, every passing char is single-byte UTF-8 so length()
is an exact byte-count proxy for the 256-byte limit.

Drop the redundant resolvedPagentSpanId != null guard on the agent-kind
propagation block - always non-null at that point.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…pagent inheritance

The stale-context trace-ID gate was correct for parent_id/session_id (APM-trace
concepts) but wrong for pagent: LLMObsContext scopes are explicitly closed in
finish(), so cross-trace leakage is impossible. Gating on trace IDs broke
in-process inheritance because each DDLLMObsSpan creates its own APM trace when
no APM scope is active.

Also hardens LLMObsSpanMapper pagent_span_id validation (non-empty String check),
fixes metaSize formula for pagent_name, removes byte-array alloc in agentNameWireSafe,
and clears W3C cache in PTagsFactory on pagent updates.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
The stale-context gate (same applied to parent_id and session_id) must also
cover pagent: without it a stale LLMObsContext from a different trace leaked
across an async boundary would attribute spans to an unrelated agent.

Tests now establish a root APM scope so all LLMObs spans share one APM trace,
matching production behavior where the DD agent always activates a root scope.
This makes the trace-ID consistency check reliable in tests.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…o auto-instrumented spans

W3CPTagsCodec.fromHeaderValue now extracts llmobs_pagent_span_id and
llmobs_pagent_name into the named PTags fields, matching what DatadogPTagsCodec
already does. Without this, tracecontext-only hops lost attribution.

OpenAiDecorator.doAfterStart now inherits agent attribution from LLMObsContext,
following the same pattern already used for session_id. Auto-instrumented LLM
spans inside a manual agent span now get the pagent_* tags set correctly.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Two cases the attribution test suite was missing:
- inner agent finish() restores outer agent's PTags (nested agent support)
- stale LLMObsContext from a different APM trace does not leak pagent attribution

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Three correctness fixes from review:

- Reject tilde (0x7E) in agentNameWireSafe: W3C tracestate encoding
  rewrites ~ to _, which would silently corrupt the agent name after a
  tracecontext hop. Now treated as unsafe, matching the stricter bound.

- Clear PAGENT_NAME_KEY in LLMObsContext.attach when name is null:
  previously an unsafe-named inner agent left the outer agent's name in
  the context, so tool spans under the inner agent inherited a mismatched
  ID/name pair. Context.with(key, null) removes the key.

- Fix metaSize off-by-one in LLMObsSpanMapper: when pagent_span_id is
  present but invalid (non-String or empty), hasAgentAttribution=false
  skips the entry during serialization but the old formula still counted
  it in tagsToRemapToMeta.size(), producing a malformed msgpack map.

Tests added for tilde rejection and the unsafe-name-clear case.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
When a manual agent span is started without an ambient APM root (e.g. a
script or CLI with no DD agent HTTP server instrumentation), its local
root is itself. Outgoing HTTP calls instrumented after it create a fresh
APM root with empty PTags, so the pagent values stamped on the agent's
root context are never injected into outgoing headers.

Fix: detect the standalone case (kind=agent AND span.getLocalRootSpan()==span)
and activate the underlying APM span as an AgentScope. Auto-instrumented
outgoing spans then become children of the agent span, share its APM trace,
and pick up the pagent PTags on injection. The scope is closed in finish().

In the production case (DD agent's HTTP server APM root already active),
getLocalRootSpan() != span, so the activation is skipped and existing
behavior is preserved.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…PTags

Removes the LLMObsPropagationAccess bridge and all associated propagation
tag plumbing. The shared-root PTags approach races for parallel agent subtrees,
and the Java LLMObs SDK currently has no distributed tracing support — so
cross-service propagation is deferred to a follow-up PR.

In-process attribution via LLMObsContext is unaffected; meta.agent_attribution
serialization is preserved.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Two tests created AGENT-kind DDLLMObsSpan instances without calling finish().
DDLLMObsSpan activates a standaloneApmScope when an agent span is its own APM
local root (no ambient trace). Without finish(), that scope persisted across
tests, causing all subsequent tests to fail on the activeSpan() == null guard
in setup().

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Verifies agent attribution end-to-end through the public LLMObs.start*Span()
API (the same path a user's code takes) rather than DDLLMObsSpan directly.
Installs a RealSpanFactory backed by DDLLMObsSpan — mirroring what
LLMObsSystem does when the agent boots with DD_LLMOBS_ENABLED=true.

Scenarios covered:
- agent self-attributes
- LLM under agent inherits attribution
- tool transitively inherits (agent → llm → tool)
- nested agents: inner overrides outer for descendants; outer restores after inner finishes
- no agent ancestor → no attribution tags
- unsafe name (comma): ID set, name null; propagated to children
- tilde in name: rejected by wire-safe guard
- realistic router → executor multi-agent workflow

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
- Inline result variable in DatadogPTagsCodec.createValid call (left
  as a two-step assign-then-return by an earlier reverted commit)
- Inline result variable in W3CPTagsCodec new W3CPTags call (same)
- Rename hasInvalidPagentSpanId → hasInvalidParentAgentSpanId for
  readability per reviewer nit

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Rename all code variables that abbreviated the concept as `pagent`
to the fully-spelled-out `parentAgent` prefix for readability, per
sabrenner's review nit. Tag-name strings and constants (pagent_span_id,
pagent_name in the wire format) are unchanged.

Affected identifiers:
  resolvedPagentSpanId/Name  → resolvedParentAgentSpanId/Name (DDLLMObsSpan)
  pagentSpanId/Name params   → parentAgentSpanId/Name (LLMObsContext.attach)
  pagentSpanIdVal            → parentAgentSpanIdVal (LLMObsSpanMapper)
  pagentSpanId/Name locals   → parentAgentSpanId/Name (OpenAiDecorator)
  pagentSpanId/Name helpers  → parentAgentSpanId/Name (integration tests)
  expectedPagent*/outerPagent* → expectedParentAgent*/outerParentAgent* (unit tests)

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Removes the standalone APM scope activation from DDLLMObsSpan and its
two tests. This feature (ensuring outgoing instrumented calls nest under
a standalone agent span's APM trace) is correct but out of scope for
this PR, which is focused on in-process pagent attribution. It will be
reintroduced in a follow-up with proper distributed tracing tests.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
yahya-mouman and others added 2 commits August 31, 2026 16:47
… in OpenAiDecorator

The rename refactor (spell out pagent as parentAgent) updated the local
variable declarations but missed the usages two lines below, causing
undefined symbol compile errors.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Remove trailing blank lines flagged by google-java-format:
- DDLLMObsSpan.java: blank line inside constructor body before closing brace
- DDLLMObsSpanAgentAttributionTest.java: blank line before class closing brace

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
@yahya-mouman
yahya-mouman force-pushed the yahya/llmobs-agent-attribution branch from 4554204 to 3b6a952 Compare August 31, 2026 14:56
@yahya-mouman
yahya-mouman marked this pull request as ready for review August 31, 2026 15:04
@yahya-mouman
yahya-mouman requested review from a team as code owners August 31, 2026 15:04
@yahya-mouman
yahya-mouman requested review from ValentinZakharov and removed request for a team August 31, 2026 15:04
@dd-octo-sts dd-octo-sts Bot added the tag: ai generated Largely based on code generated by an AI or LLM label Aug 31, 2026

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 3b6a9529a0

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

Without an active APM scope, each DDLLMObsSpan starts its own APM trace.
The trace-ID consistency gate in DDLLMObsSpan then blocks in-process agent
attribution inheritance because parent.getTraceId() != span.getTraceId().

In production the Datadog agent always activates an APM scope before user
code runs, so all LLMObs spans within a request share one trace. The test
must mirror this setup via @BeforeEach/@AfterEach APM scope lifecycle.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

@datadog-prod-us1-5 datadog-prod-us1-5 Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Datadog Autotest: FAIL

Standalone public SDK nesting starts each manual span on a new trace, so child spans lose agent attribution. Stale LLMObs contexts can also copy an agent ID across traces into manual or OpenAI spans.

Open Bits AI session

🤖 Datadog Autotest · Commit 3b6a952 · What is Autotest? · @DataDog review to ask questions · Any feedback? Reach out in #autotest

// leaked across an async boundary would otherwise attribute a span to an agent from a
// different trace. In production the DD agent always establishes a root APM scope, so
// all LLMObs spans within a request share one trace and this check passes.
if (null != parent && parent.getTraceId() == span.getTraceId()) {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

P1 Keep standalone child spans on the agent trace

Standalone public SDK users lose agent attribution from all manual child spans.

Assertion details
  • Input: Start a manual agent span and then start a manual LLM or tool span without an active APM scope.
  • Expected: A child LLMObs span must use the open manual agent span as its parent and inherit its agent ID and name.
  • Actual: The child span starts a new trace. The trace check then rejects the agent context and omits agent attribution.

Was this helpful? React 👍 or 👎
🤖 Datadog Autotest · What is Autotest? · @DataDog review to ask questions · Any feedback? Reach out in #autotest

Comment thread internal-api/src/main/java/datadog/trace/api/llmobs/LLMObsContext.java Outdated
yahya-mouman and others added 7 commits August 31, 2026 17:48
…ntract, add OpenAI test

- Add PAGENT_SPAN_ID/PAGENT_NAME to LLMObsTags; derive all pagent tag
  string literals from it (DDLLMObsSpan, CommonTags, LLMObsSpanMapper)
- Fix agentNameWireSafe Javadoc: correct range to 0x20-0x7D (exclusive
  of tilde) and rejection condition to c >= 0x7E
- Document Context.with(key, null) null-removes-key contract in
  LLMObsContext.attach() comment with reference to the Context API
- Add autoInstrumentedSpanInAgentScopeReadsAttributionFromContext test
  simulating what OpenAiDecorator.doAfterStart() reads from LLMObsContext

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…ionTest

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
LLMObsContext: always write both pagent keys in 4-arg attach(), so
null clears stale values from an outer scope. Previously a non-agent
span that had its attribution blocked by the trace-ID gate would still
propagate the outer context's pagent keys to its same-trace children.

OpenAiDecorator: gate pagent inheritance on trace-ID consistency,
mirroring the check in DDLLMObsSpan. A stale LLMObsContext from a
different async trace must not stamp its agent ID onto a new OpenAI span.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…striction

Agent names are now accepted as-is — the character restrictions (tilde,
non-ASCII, comma, semicolon) existed solely for x-datadog-tags header
propagation, which was removed from this PR. The msgpack intake mapper
accepts any string.

For agent spans, annotateAgentManifest() now syncs the pagent name to
the manifest name (manifest > span name fallback). Both the internal
tag on the agent span and the LLMObsContext scope are updated so
descendants started after the call inherit the manifest name.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…ifest

Closing the original scope while a newer scope is on top corrupts the
context stack: when the new scope is eventually closed it restores the
old (already-closed) scope's context rather than the outer empty one,
leaking pagent tags into subsequent tests.

Only update PAGENT_NAME_TAG_INTERNAL (for the serializer's wire output);
context propagation to children keeps the span name set at construction.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

comp: core Tracer core tag: ai generated Largely based on code generated by an AI or LLM type: feature Enhancements and improvements

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants