From d3d5bb94a0452059f8712185265c3f714b62c791 Mon Sep 17 00:00:00 2001 From: yahya-mouman Date: Thu, 16 Jul 2026 12:58:52 +0200 Subject: [PATCH 1/6] feat(llmobs): resolve nearest agent ancestor at span start Adds parentAgentName/parentAgentSpanID fields to Span, ParentAgentName/ ParentAgentSpanID fields to PropagatedLLMSpan, and wires resolveParentAgent into StartSpan so every span records its nearest agent ancestor in O(1). Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: yahya-mouman --- internal/llmobs/context.go | 6 ++ internal/llmobs/llmobs.go | 25 ++++++ internal/llmobs/resolve_parent_agent_test.go | 73 +++++++++++++++ internal/llmobs/span.go | 6 ++ internal/llmobs/start_span_agent_test.go | 94 ++++++++++++++++++++ 5 files changed, 204 insertions(+) create mode 100644 internal/llmobs/resolve_parent_agent_test.go create mode 100644 internal/llmobs/start_span_agent_test.go diff --git a/internal/llmobs/context.go b/internal/llmobs/context.go index 79ae2dd9194..8ea4e8d3bb9 100644 --- a/internal/llmobs/context.go +++ b/internal/llmobs/context.go @@ -20,6 +20,12 @@ type PropagatedLLMSpan struct { TraceID string // SpanID is the span ID. SpanID string + // ParentAgentName is the name of the nearest agent ancestor, propagated across + // process boundaries. Empty when the upstream hop sent an id-only attribution. + ParentAgentName string + // ParentAgentSpanID is the span ID of the nearest agent ancestor, propagated + // across process boundaries. Empty when there is no agent ancestor. + ParentAgentSpanID string } // PropagatedLLMSpanFromContext retrieves a PropagatedLLMSpan from the context. diff --git a/internal/llmobs/llmobs.go b/internal/llmobs/llmobs.go index bb333040c1b..43dcf5b94eb 100644 --- a/internal/llmobs/llmobs.go +++ b/internal/llmobs/llmobs.go @@ -770,6 +770,29 @@ func dropSpanEventIO(ev *transport.LLMObsSpanEvent) bool { return droppedIO } +// resolveParentAgent returns the nearest agent ancestor's (name, spanID) for a +// span about to start, given its resolved parent and/or propagated parent. +// +// Resolution is O(1): the parent already resolved its own attribution when it +// started, so a non-agent parent simply hands down what it inherited. +// +// parent is an agent span -> (parent.name, parent.SpanID()) +// parent is any other kind -> (parent.parentAgentName, parent.parentAgentSpanID) +// no local parent, propagated parent -> (propagated.ParentAgentName, propagated.ParentAgentSpanID) +// neither -> ("", "") +func resolveParentAgent(parent *Span, propagated *PropagatedLLMSpan) (name string, spanID string) { + if parent != nil { + if parent.spanKind == SpanKindAgent { + return parent.name, parent.SpanID() + } + return parent.parentAgentName, parent.parentAgentSpanID + } + if propagated != nil { + return propagated.ParentAgentName, propagated.ParentAgentSpanID + } + return "", "" +} + // StartSpan starts a new LLMObs span with the given kind, name, and configuration. // Returns the created span and a context containing the span. func (l *LLMObs) StartSpan(ctx context.Context, kind SpanKind, name string, cfg StartSpanConfig) (*Span, context.Context) { @@ -813,6 +836,8 @@ func (l *LLMObs) StartSpan(ctx context.Context, kind SpanKind, name string, cfg span.llmTraceID = newLLMObsTraceID() } + span.parentAgentName, span.parentAgentSpanID = resolveParentAgent(span.parent, span.propagated) + span.mlApp = cfg.MLApp span.spanKind = kind span.sessionID = cfg.SessionID diff --git a/internal/llmobs/resolve_parent_agent_test.go b/internal/llmobs/resolve_parent_agent_test.go new file mode 100644 index 00000000000..3f6f11ebc13 --- /dev/null +++ b/internal/llmobs/resolve_parent_agent_test.go @@ -0,0 +1,73 @@ +// Unless explicitly stated otherwise all files in this repository are licensed +// under the Apache License Version 2.0. +// This product includes software developed at Datadog (https://www.datadoghq.com/). +// Copyright 2025 Datadog, Inc. + +package llmobs + +import "testing" + +func TestResolveParentAgent(t *testing.T) { + t.Run("no-parent", func(t *testing.T) { + name, id := resolveParentAgent(nil, nil) + if name != "" || id != "" { + t.Fatalf("expected empty, got name=%q id=%q", name, id) + } + }) + + t.Run("parent-is-agent", func(t *testing.T) { + parent := &Span{name: "my_agent", spanKind: SpanKindAgent} + parent.apm = testAPMSpan{spanID: "111"} + name, id := resolveParentAgent(parent, nil) + if name != "my_agent" || id != "111" { + t.Fatalf("expected (my_agent,111), got (%q,%q)", name, id) + } + }) + + t.Run("parent-other-kind-inherits", func(t *testing.T) { + parent := &Span{ + name: "tool_x", + spanKind: SpanKindTool, + parentAgentName: "top_agent", + parentAgentSpanID: "222", + } + parent.apm = testAPMSpan{spanID: "333"} + name, id := resolveParentAgent(parent, nil) + if name != "top_agent" || id != "222" { + t.Fatalf("expected inherited (top_agent,222), got (%q,%q)", name, id) + } + }) + + t.Run("propagated-parent", func(t *testing.T) { + prop := &PropagatedLLMSpan{ParentAgentName: "remote_agent", ParentAgentSpanID: "444"} + name, id := resolveParentAgent(nil, prop) + if name != "remote_agent" || id != "444" { + t.Fatalf("expected (remote_agent,444), got (%q,%q)", name, id) + } + }) + + t.Run("parent-takes-precedence-over-propagated", func(t *testing.T) { + parent := &Span{name: "local_agent", spanKind: SpanKindAgent} + parent.apm = testAPMSpan{spanID: "555"} + prop := &PropagatedLLMSpan{ParentAgentName: "remote_agent", ParentAgentSpanID: "444"} + name, id := resolveParentAgent(parent, prop) + if name != "local_agent" || id != "555" { + t.Fatalf("expected local (local_agent,555), got (%q,%q)", name, id) + } + }) +} + +// testAPMSpan is a minimal APMSpan implementation returning fixed values. +// All methods are implemented to avoid nil-pointer panics when embedded-interface +// patterns would otherwise dispatch to a nil value. +type testAPMSpan struct { + spanID string + traceID string +} + +func (s testAPMSpan) Finish(_ FinishAPMSpanConfig) {} +func (s testAPMSpan) AddLink(_ SpanLink) {} +func (s testAPMSpan) SpanID() string { return s.spanID } +func (s testAPMSpan) TraceID() string { return s.traceID } +func (s testAPMSpan) SetBaggageItem(_ string, _ string) {} +func (s testAPMSpan) BaggageItem(_ string) string { return "" } diff --git a/internal/llmobs/span.go b/internal/llmobs/span.go index 20af6e683bc..7c6bea56018 100644 --- a/internal/llmobs/span.go +++ b/internal/llmobs/span.go @@ -253,6 +253,12 @@ type Span struct { finishTime time.Time spanLinks []SpanLink + + // parentAgentName and parentAgentSpanID identify the nearest agent ancestor. + // Both are set exactly once in StartSpan and never mutated, so concurrent + // reads (e.g. from Annotate) are safe without holding the mutex. + parentAgentName string + parentAgentSpanID string } func (s *Span) Name() string { diff --git a/internal/llmobs/start_span_agent_test.go b/internal/llmobs/start_span_agent_test.go new file mode 100644 index 00000000000..2c95db700f5 --- /dev/null +++ b/internal/llmobs/start_span_agent_test.go @@ -0,0 +1,94 @@ +// Unless explicitly stated otherwise all files in this repository are licensed +// under the Apache License Version 2.0. +// This product includes software developed at Datadog (https://www.datadoghq.com/). +// Copyright 2025 Datadog, Inc. + +package llmobs + +import ( + "context" + "fmt" + "testing" + + "github.com/DataDog/dd-trace-go/v2/internal/llmobs/config" +) + +// resolveTestTracer is a minimal Tracer that returns testAPMSpans with unique IDs. +type resolveTestTracer struct { + next int +} + +func (tr *resolveTestTracer) StartSpan(ctx context.Context, _ string, _ StartAPMSpanConfig) (APMSpan, context.Context) { + tr.next++ + return testAPMSpan{spanID: fmt.Sprintf("id-%d", tr.next)}, ctx +} + +func newTestLLMObsForResolve(t *testing.T) *LLMObs { + t.Helper() + return &LLMObs{ + Config: &config.Config{Enabled: true, MLApp: "gotest"}, + Tracer: &resolveTestTracer{}, + } +} + +func TestStartSpanResolvesParentAgent(t *testing.T) { + l := newTestLLMObsForResolve(t) + ctx := context.Background() + + t.Run("tool-under-agent", func(t *testing.T) { + agent, agentCtx := l.StartSpan(ctx, SpanKindAgent, "my_agent", StartSpanConfig{}) + tool, _ := l.StartSpan(agentCtx, SpanKindTool, "my_tool", StartSpanConfig{}) + if tool.parentAgentName != "my_agent" { + t.Fatalf("tool.parentAgentName = %q, want my_agent", tool.parentAgentName) + } + if tool.parentAgentSpanID != agent.SpanID() { + t.Fatalf("tool.parentAgentSpanID = %q, want %q", tool.parentAgentSpanID, agent.SpanID()) + } + }) + + t.Run("agent-workflow-tool-indirect-nesting", func(t *testing.T) { + agent, agentCtx := l.StartSpan(ctx, SpanKindAgent, "my_agent", StartSpanConfig{}) + wf, wfCtx := l.StartSpan(agentCtx, SpanKindWorkflow, "wf", StartSpanConfig{}) + tool, _ := l.StartSpan(wfCtx, SpanKindTool, "tool", StartSpanConfig{}) + if wf.parentAgentSpanID != agent.SpanID() || wf.parentAgentName != "my_agent" { + t.Fatalf("workflow should attribute to top agent, got (%q,%q)", wf.parentAgentName, wf.parentAgentSpanID) + } + if tool.parentAgentSpanID != agent.SpanID() || tool.parentAgentName != "my_agent" { + t.Fatalf("tool should attribute to top agent, got (%q,%q)", tool.parentAgentName, tool.parentAgentSpanID) + } + }) + + t.Run("sub-agent-under-agent", func(t *testing.T) { + outer, outerCtx := l.StartSpan(ctx, SpanKindAgent, "outer_agent", StartSpanConfig{}) + inner, innerCtx := l.StartSpan(outerCtx, SpanKindAgent, "inner_agent", StartSpanConfig{}) + innerTool, _ := l.StartSpan(innerCtx, SpanKindTool, "inner_tool", StartSpanConfig{}) + if inner.parentAgentSpanID != outer.SpanID() || inner.parentAgentName != "outer_agent" { + t.Fatalf("inner agent should attribute to outer, got (%q,%q)", inner.parentAgentName, inner.parentAgentSpanID) + } + if innerTool.parentAgentSpanID != inner.SpanID() || innerTool.parentAgentName != "inner_agent" { + t.Fatalf("inner tool should attribute to inner agent, got (%q,%q)", innerTool.parentAgentName, innerTool.parentAgentSpanID) + } + }) + + t.Run("top-level-agent-has-no-attribution", func(t *testing.T) { + agent, _ := l.StartSpan(ctx, SpanKindAgent, "top_agent", StartSpanConfig{}) + if agent.parentAgentName != "" || agent.parentAgentSpanID != "" { + t.Fatalf("top-level agent must have empty attribution, got (%q,%q)", agent.parentAgentName, agent.parentAgentSpanID) + } + }) + + t.Run("top-level-llm-has-no-attribution", func(t *testing.T) { + llm, _ := l.StartSpan(ctx, SpanKindLLM, "top_llm", StartSpanConfig{}) + if llm.parentAgentName != "" || llm.parentAgentSpanID != "" { + t.Fatalf("top-level llm must have empty attribution, got (%q,%q)", llm.parentAgentName, llm.parentAgentSpanID) + } + }) + + t.Run("workflow-then-tool-no-agent-anywhere", func(t *testing.T) { + _, wfCtx := l.StartSpan(ctx, SpanKindWorkflow, "wf", StartSpanConfig{}) + tool, _ := l.StartSpan(wfCtx, SpanKindTool, "tool", StartSpanConfig{}) + if tool.parentAgentName != "" || tool.parentAgentSpanID != "" { + t.Fatalf("tool with no agent ancestor must have empty attribution, got (%q,%q)", tool.parentAgentName, tool.parentAgentSpanID) + } + }) +} From a2e5d9e00c5d794103d25466956489dcfa870c18 Mon Sep 17 00:00:00 2001 From: yahya-mouman Date: Thu, 16 Jul 2026 13:05:17 +0200 Subject: [PATCH 2/6] feat(llmobs): serialize meta.agent_attribution block Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: yahya-mouman --- internal/llmobs/llmobs.go | 12 ++++++ internal/llmobs/llmobs_test.go | 70 ++++++++++++++++++++++++++++++++++ 2 files changed, 82 insertions(+) diff --git a/internal/llmobs/llmobs.go b/internal/llmobs/llmobs.go index 43dcf5b94eb..d09a4dc885f 100644 --- a/internal/llmobs/llmobs.go +++ b/internal/llmobs/llmobs.go @@ -596,6 +596,18 @@ func (l *LLMObs) llmobsSpanEvent(span *Span) *transport.LLMObsSpanEvent { meta["output"] = output } + if span.parentAgentSpanID != "" { + var agentName any = span.parentAgentName + if span.parentAgentName == "" { + // id-only: emit explicit JSON null, matching the Python/Node wire shape. + agentName = nil + } + meta["agent_attribution"] = map[string]any{ + "pagent_name": agentName, + "pagent_span_id": span.parentAgentSpanID, + } + } + spanID := span.apm.SpanID() parentID := defaultParentID if span.parent != nil { diff --git a/internal/llmobs/llmobs_test.go b/internal/llmobs/llmobs_test.go index f403cc7a459..6fdac2154ed 100644 --- a/internal/llmobs/llmobs_test.go +++ b/internal/llmobs/llmobs_test.go @@ -2278,6 +2278,76 @@ func traceHandler(h http.Handler) http.Handler { }) } +func TestAgentAttributionSerialization(t *testing.T) { + t.Run("tool-under-agent-has-attribution", func(t *testing.T) { + _, coll, ll := testTracer(t) + ctx := context.Background() + + agent, ctx := ll.StartSpan(ctx, llmobs.SpanKindAgent, "my_agent", llmobs.StartSpanConfig{}) + tool, _ := ll.StartSpan(ctx, llmobs.SpanKindTool, "my_tool", llmobs.StartSpanConfig{}) + tool.Finish(llmobs.FinishSpanConfig{}) + agent.Finish(llmobs.FinishSpanConfig{}) + tracer.Flush() + + toolSpan := coll.RequireSpan(t, "my_tool") + attr, ok := toolSpan.Meta["agent_attribution"].(map[string]any) + require.True(t, ok, "agent_attribution must be present on tool span") + assert.Equal(t, "my_agent", attr["pagent_name"]) + assert.Equal(t, agent.SpanID(), attr["pagent_span_id"]) + }) + + t.Run("top-level-agent-omits-attribution", func(t *testing.T) { + _, coll, ll := testTracer(t) + ctx := context.Background() + + agent, _ := ll.StartSpan(ctx, llmobs.SpanKindAgent, "top_agent", llmobs.StartSpanConfig{}) + agent.Finish(llmobs.FinishSpanConfig{}) + tracer.Flush() + + agentSpan := coll.RequireSpan(t, "top_agent") + _, ok := agentSpan.Meta["agent_attribution"] + assert.False(t, ok, "top-level agent must omit agent_attribution") + }) + + t.Run("no-agent-anywhere-omits-attribution", func(t *testing.T) { + _, coll, ll := testTracer(t) + ctx := context.Background() + + wf, ctx := ll.StartSpan(ctx, llmobs.SpanKindWorkflow, "wf", llmobs.StartSpanConfig{}) + tool, _ := ll.StartSpan(ctx, llmobs.SpanKindTool, "tool", llmobs.StartSpanConfig{}) + tool.Finish(llmobs.FinishSpanConfig{}) + wf.Finish(llmobs.FinishSpanConfig{}) + tracer.Flush() + + _, wfOK := coll.RequireSpan(t, "wf").Meta["agent_attribution"] + _, toolOK := coll.RequireSpan(t, "tool").Meta["agent_attribution"] + assert.False(t, wfOK, "workflow with no agent ancestor must omit agent_attribution") + assert.False(t, toolOK, "tool with no agent ancestor must omit agent_attribution") + }) + + t.Run("id-only-emits-null-name", func(t *testing.T) { + // Simulate an id-only propagated parent: name empty, id set. + _, coll, ll := testTracer(t) + prop := &llmobs.PropagatedLLMSpan{ + MLApp: mlApp, + ParentAgentSpanID: "9999", + ParentAgentName: "", // id-only + } + ctx := llmobs.ContextWithPropagatedLLMSpan(context.Background(), prop) + + child, _ := ll.StartSpan(ctx, llmobs.SpanKindTool, "child_tool", llmobs.StartSpanConfig{}) + child.Finish(llmobs.FinishSpanConfig{}) + tracer.Flush() + + attr, ok := coll.RequireSpan(t, "child_tool").Meta["agent_attribution"].(map[string]any) + require.True(t, ok, "agent_attribution must be present when id-only") + assert.Equal(t, "9999", attr["pagent_span_id"]) + v, present := attr["pagent_name"] + require.True(t, present, "pagent_name key must be present (explicit null), not absent") + assert.Nil(t, v, "pagent_name must be JSON null when name is empty") + }) +} + func traceClient(c *http.Client) *http.Client { c.Transport = &tracedRT{base: c.Transport} return c From a48aa5b78a3ac7927405c6982b7843870c69353c Mon Sep 17 00:00:00 2001 From: yahya-mouman Date: Thu, 16 Jul 2026 13:10:11 +0200 Subject: [PATCH 3/6] feat(llmobs): add agentNameWireSafe wire-safety check Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: yahya-mouman --- internal/llmobs/util.go | 32 +++++++++++++++++++++++++++++ internal/llmobs/util_test.go | 39 ++++++++++++++++++++++++++++++++++++ 2 files changed, 71 insertions(+) create mode 100644 internal/llmobs/util.go create mode 100644 internal/llmobs/util_test.go diff --git a/internal/llmobs/util.go b/internal/llmobs/util.go new file mode 100644 index 00000000000..b0891d66b87 --- /dev/null +++ b/internal/llmobs/util.go @@ -0,0 +1,32 @@ +// Unless explicitly stated otherwise all files in this repository are licensed +// under the Apache License Version 2.0. +// This product includes software developed at Datadog (https://www.datadoghq.com/). +// Copyright 2025 Datadog, Inc. + +package llmobs + +// agentNameWireSafe reports whether name can be safely written as an +// x-datadog-tags tag value. The rules match the shared cross-language contract: +// +// - reject if the name exceeds 256 bytes +// - reject any byte outside the printable ASCII range [0x20, 0x7E] +// - reject a comma (the tagset delimiter) +// +// An equals sign is legal in tagset values (only illegal in keys) and is NOT +// rejected. This gate applies only to the wire tag; meta.agent_attribution +// always uses the real name. +func agentNameWireSafe(name string) bool { + if len(name) > 256 { + return false + } + for i := 0; i < len(name); i++ { + b := name[i] + if b < 0x20 || b > 0x7E { + return false + } + if b == ',' { + return false + } + } + return true +} diff --git a/internal/llmobs/util_test.go b/internal/llmobs/util_test.go new file mode 100644 index 00000000000..c92a9763043 --- /dev/null +++ b/internal/llmobs/util_test.go @@ -0,0 +1,39 @@ +// Unless explicitly stated otherwise all files in this repository are licensed +// under the Apache License Version 2.0. +// This product includes software developed at Datadog (https://www.datadoghq.com/). +// Copyright 2025 Datadog, Inc. + +package llmobs + +import ( + "strings" + "testing" +) + +func TestAgentNameWireSafe(t *testing.T) { + cases := []struct { + name string + input string + want bool + }{ + {"plain-safe", "my_agent", true}, + {"empty", "", true}, + {"contains-comma", "my,agent", false}, + {"contains-equals-is-safe", "key=val", true}, + {"space-is-safe", "my agent", true}, + {"tilde-boundary-safe", "~", true}, + {"space-boundary-safe", " ", true}, + {"tab-non-printable", "my\tagent", false}, + {"newline-non-printable", "my\nagent", false}, + {"multibyte-utf8-unsafe", "agént", false}, + {"len-256-safe", strings.Repeat("a", 256), true}, + {"len-257-unsafe", strings.Repeat("a", 257), false}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := agentNameWireSafe(tc.input); got != tc.want { + t.Fatalf("agentNameWireSafe(%q) = %v, want %v", tc.input, got, tc.want) + } + }) + } +} From 3b96907e45474adfed8b204e6ea655087f095f06 Mon Sep 17 00:00:00 2001 From: yahya-mouman Date: Thu, 16 Jul 2026 13:15:03 +0200 Subject: [PATCH 4/6] feat(llmobs): add InjectContext/ExtractContext for agent attribution Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: yahya-mouman --- internal/llmobs/propagation.go | 106 +++++++++++++++ internal/llmobs/propagation_test.go | 192 ++++++++++++++++++++++++++++ llmobs/llmobs.go | 19 +++ 3 files changed, 317 insertions(+) create mode 100644 internal/llmobs/propagation.go create mode 100644 internal/llmobs/propagation_test.go diff --git a/internal/llmobs/propagation.go b/internal/llmobs/propagation.go new file mode 100644 index 00000000000..ecf1e9cbbb4 --- /dev/null +++ b/internal/llmobs/propagation.go @@ -0,0 +1,106 @@ +// Unless explicitly stated otherwise all files in this repository are licensed +// under the Apache License Version 2.0. +// This product includes software developed at Datadog (https://www.datadoghq.com/). +// Copyright 2025 Datadog, Inc. + +package llmobs + +import ( + "context" + "strings" +) + +const ( + headerKeyDatadogTags = "x-datadog-tags" + tagKeyLLMObsPAgentSpanID = "_dd.p.llmobs_pagent_span_id" + tagKeyLLMObsPAgentName = "_dd.p.llmobs_pagent_name" + + // maxDatadogTagsLen bounds the total x-datadog-tags value length. The name + // tag is only appended when the result stays within this budget. + maxDatadogTagsLen = 512 +) + +// InjectContext writes the active span's outbound agent attribution into the +// carrier's x-datadog-tags value. Exported so the public llmobs package (which +// cannot read Span's unexported fields) can delegate here. +func InjectContext(ctx context.Context, carrier map[string]string) { + span, ok := ActiveLLMSpanFromContext(ctx) + if !ok || span == nil { + return + } + injectAgentAttribution(span, carrier) +} + +// ExtractContext reads agent attribution from the carrier and attaches it to a +// PropagatedLLMSpan on the returned context, preserving any fields already set. +func ExtractContext(ctx context.Context, carrier map[string]string) context.Context { + name, spanID, present := extractAgentAttribution(carrier) + if !present { + return ctx + } + + // Update, do not overwrite: preserve any existing MLApp/TraceID/SpanID. + var prop PropagatedLLMSpan + if existing, ok := PropagatedLLMSpanFromContext(ctx); ok && existing != nil { + prop = *existing + } + prop.ParentAgentName = name + prop.ParentAgentSpanID = spanID + return ContextWithPropagatedLLMSpan(ctx, &prop) +} + +// injectAgentAttribution resolves the span's outbound agent attribution and +// merges the resulting tags into the carrier's x-datadog-tags value. +// +// active span is an agent -> (self.name, self.SpanID()) +// otherwise -> (span.parentAgentName, span.parentAgentSpanID) +func injectAgentAttribution(span *Span, carrier map[string]string) { + var name, spanID string + if span.spanKind == SpanKindAgent { + name, spanID = span.name, span.SpanID() + } else { + name, spanID = span.parentAgentName, span.parentAgentSpanID + } + if spanID == "" { + return // no agent ancestor: write nothing + } + + existing := carrier[headerKeyDatadogTags] + parts := make([]string, 0, 2) + if existing != "" { + parts = append(parts, existing) + } + parts = append(parts, tagKeyLLMObsPAgentSpanID+"="+spanID) + + // Append the name only if wire-safe and within the total length budget. + if name != "" && agentNameWireSafe(name) { + candidate := append(parts, tagKeyLLMObsPAgentName+"="+name) + if joined := strings.Join(candidate, ","); len(joined) <= maxDatadogTagsLen { + parts = candidate + } + } + carrier[headerKeyDatadogTags] = strings.Join(parts, ",") +} + +// extractAgentAttribution parses the carrier's x-datadog-tags value for the +// llmobs agent attribution tags. present is true if the span id tag was found. +func extractAgentAttribution(carrier map[string]string) (name string, spanID string, present bool) { + header, ok := carrier[headerKeyDatadogTags] + if !ok || header == "" { + return "", "", false + } + for _, pair := range strings.Split(header, ",") { + k, v, found := strings.Cut(pair, "=") + if !found { + continue + } + switch k { + case tagKeyLLMObsPAgentSpanID: + spanID = v + present = true + case tagKeyLLMObsPAgentName: + name = v + } + } + return name, spanID, present +} diff --git a/internal/llmobs/propagation_test.go b/internal/llmobs/propagation_test.go new file mode 100644 index 00000000000..0d33946278b --- /dev/null +++ b/internal/llmobs/propagation_test.go @@ -0,0 +1,192 @@ +// Unless explicitly stated otherwise all files in this repository are licensed +// under the Apache License Version 2.0. +// This product includes software developed at Datadog (https://www.datadoghq.com/). +// Copyright 2025 Datadog, Inc. + +package llmobs + +import ( + "context" + "strings" + "testing" +) + +// tagValue returns the value of key within a comma-delimited x-datadog-tags string. +func tagValue(header, key string) (string, bool) { + for _, pair := range strings.Split(header, ",") { + k, v, found := strings.Cut(pair, "=") + if found && k == key { + return v, true + } + } + return "", false +} + +func TestInjectAgentAttribution(t *testing.T) { + t.Run("inject-from-agent", func(t *testing.T) { + agent := &Span{name: "my_agent", spanKind: SpanKindAgent} + agent.apm = testAPMSpan{spanID: "1000"} + carrier := map[string]string{} + injectAgentAttribution(agent, carrier) + + tags := carrier["x-datadog-tags"] + id, ok := tagValue(tags, "_dd.p.llmobs_pagent_span_id") + if !ok || id != "1000" { + t.Fatalf("expected span id 1000, got %q (ok=%v) in %q", id, ok, tags) + } + name, ok := tagValue(tags, "_dd.p.llmobs_pagent_name") + if !ok || name != "my_agent" { + t.Fatalf("expected name my_agent, got %q (ok=%v)", name, ok) + } + }) + + t.Run("inject-from-tool-under-agent-inherits", func(t *testing.T) { + tool := &Span{name: "tool", spanKind: SpanKindTool, parentAgentName: "my_agent", parentAgentSpanID: "1000"} + tool.apm = testAPMSpan{spanID: "2000"} + carrier := map[string]string{} + injectAgentAttribution(tool, carrier) + + tags := carrier["x-datadog-tags"] + id, _ := tagValue(tags, "_dd.p.llmobs_pagent_span_id") + name, _ := tagValue(tags, "_dd.p.llmobs_pagent_name") + if id != "1000" || name != "my_agent" { + t.Fatalf("tool should inject inherited agent (my_agent,1000), got (%q,%q)", name, id) + } + }) + + t.Run("no-agent-in-chain-writes-nothing", func(t *testing.T) { + tool := &Span{name: "tool", spanKind: SpanKindTool} + tool.apm = testAPMSpan{spanID: "2000"} + carrier := map[string]string{} + injectAgentAttribution(tool, carrier) + + if _, ok := carrier["x-datadog-tags"]; ok { + if _, hasID := tagValue(carrier["x-datadog-tags"], "_dd.p.llmobs_pagent_span_id"); hasID { + t.Fatalf("no agent ancestor: span_id tag must not be written, got %q", carrier["x-datadog-tags"]) + } + } + }) + + t.Run("unsafe-name-comma-writes-id-only", func(t *testing.T) { + agent := &Span{name: "bad,name", spanKind: SpanKindAgent} + agent.apm = testAPMSpan{spanID: "1000"} + carrier := map[string]string{} + injectAgentAttribution(agent, carrier) + + tags := carrier["x-datadog-tags"] + if _, ok := tagValue(tags, "_dd.p.llmobs_pagent_span_id"); !ok { + t.Fatalf("expected span id written, got %q", tags) + } + if _, ok := tagValue(tags, "_dd.p.llmobs_pagent_name"); ok { + t.Fatalf("unsafe name must be omitted, got %q", tags) + } + }) + + t.Run("name-with-equals-propagates", func(t *testing.T) { + agent := &Span{name: "k=v", spanKind: SpanKindAgent} + agent.apm = testAPMSpan{spanID: "1000"} + carrier := map[string]string{} + injectAgentAttribution(agent, carrier) + + name, ok := tagValue(carrier["x-datadog-tags"], "_dd.p.llmobs_pagent_name") + if !ok || name != "k=v" { + t.Fatalf("name with '=' must propagate, got %q (ok=%v)", name, ok) + } + }) + + t.Run("oversized-name-writes-id-only", func(t *testing.T) { + agent := &Span{name: strings.Repeat("a", 257), spanKind: SpanKindAgent} + agent.apm = testAPMSpan{spanID: "1000"} + carrier := map[string]string{} + injectAgentAttribution(agent, carrier) + + if _, ok := tagValue(carrier["x-datadog-tags"], "_dd.p.llmobs_pagent_span_id"); !ok { + t.Fatalf("expected span id written for oversized name") + } + if _, ok := tagValue(carrier["x-datadog-tags"], "_dd.p.llmobs_pagent_name"); ok { + t.Fatalf("oversized name must be omitted") + } + }) + + t.Run("merges-with-existing-tags", func(t *testing.T) { + agent := &Span{name: "my_agent", spanKind: SpanKindAgent} + agent.apm = testAPMSpan{spanID: "1000"} + carrier := map[string]string{"x-datadog-tags": "_dd.p.dm=-0"} + injectAgentAttribution(agent, carrier) + + tags := carrier["x-datadog-tags"] + if _, ok := tagValue(tags, "_dd.p.dm"); !ok { + t.Fatalf("existing tag _dd.p.dm must be preserved, got %q", tags) + } + if _, ok := tagValue(tags, "_dd.p.llmobs_pagent_span_id"); !ok { + t.Fatalf("new tag must be appended, got %q", tags) + } + }) +} + +func TestExtractAgentAttribution(t *testing.T) { + t.Run("extract-both", func(t *testing.T) { + carrier := map[string]string{ + "x-datadog-tags": "_dd.p.llmobs_pagent_span_id=1000,_dd.p.llmobs_pagent_name=my_agent", + } + name, id, present := extractAgentAttribution(carrier) + if !present || id != "1000" || name != "my_agent" { + t.Fatalf("expected (my_agent,1000,true), got (%q,%q,%v)", name, id, present) + } + }) + + t.Run("extract-id-only", func(t *testing.T) { + carrier := map[string]string{ + "x-datadog-tags": "_dd.p.llmobs_pagent_span_id=1000", + } + name, id, present := extractAgentAttribution(carrier) + if !present || id != "1000" || name != "" { + t.Fatalf("expected id-only (\"\",1000,true), got (%q,%q,%v)", name, id, present) + } + }) + + t.Run("extract-absent", func(t *testing.T) { + carrier := map[string]string{"x-datadog-tags": "_dd.p.dm=-0"} + _, _, present := extractAgentAttribution(carrier) + if present { + t.Fatalf("expected present=false when no llmobs tags") + } + }) +} + +func TestExtractContextUpdatesExistingPropagated(t *testing.T) { + // A prior ExtractContext (or manual set) put ML app + trace id on the context. + base := ContextWithPropagatedLLMSpan(context.Background(), &PropagatedLLMSpan{ + MLApp: "existing_app", + TraceID: "trace123", + SpanID: "span456", + }) + carrier := map[string]string{ + "x-datadog-tags": "_dd.p.llmobs_pagent_span_id=1000,_dd.p.llmobs_pagent_name=my_agent", + } + out := ExtractContext(base, carrier) + prop, ok := PropagatedLLMSpanFromContext(out) + if !ok { + t.Fatal("expected a propagated span on the returned context") + } + if prop.MLApp != "existing_app" || prop.TraceID != "trace123" || prop.SpanID != "span456" { + t.Fatalf("existing propagated fields must be preserved, got %+v", prop) + } + if prop.ParentAgentName != "my_agent" || prop.ParentAgentSpanID != "1000" { + t.Fatalf("agent fields must be set, got name=%q id=%q", prop.ParentAgentName, prop.ParentAgentSpanID) + } +} + +func TestInjectContextReadsActiveSpan(t *testing.T) { + agent := &Span{name: "my_agent", spanKind: SpanKindAgent} + agent.apm = testAPMSpan{spanID: "1000"} + ctx := contextWithActiveLLMSpan(context.Background(), agent) + carrier := map[string]string{} + InjectContext(ctx, carrier) + + name, _ := tagValue(carrier["x-datadog-tags"], "_dd.p.llmobs_pagent_name") + id, _ := tagValue(carrier["x-datadog-tags"], "_dd.p.llmobs_pagent_span_id") + if name != "my_agent" || id != "1000" { + t.Fatalf("InjectContext should inject active agent (my_agent,1000), got (%q,%q)", name, id) + } +} diff --git a/llmobs/llmobs.go b/llmobs/llmobs.go index bdf5eff0a92..5d984aa6156 100644 --- a/llmobs/llmobs.go +++ b/llmobs/llmobs.go @@ -451,3 +451,22 @@ func parseAnnotateOptions(opts ...AnnotateOption) illmobs.SpanAnnotations { func isSpanKind(s Span, target illmobs.SpanKind) bool { return illmobs.SpanKind(s.Kind()) == target } + +// InjectContext writes the active LLMObs span's agent attribution into the +// carrier so a downstream service can resolve agent_attribution for its spans. +// Call this alongside the APM tracer's Inject when making a distributed request. +// +// The carrier is a map of header name to value (e.g. an HTTP header map). Agent +// attribution is written into the x-datadog-tags value, merging with any value +// already present. +func InjectContext(ctx context.Context, carrier map[string]string) { + illmobs.InjectContext(ctx, carrier) +} + +// ExtractContext reads agent attribution written by an upstream InjectContext +// from the carrier and returns a context carrying it, so spans started from the +// returned context inherit the upstream agent attribution. Call this alongside +// the APM tracer's Extract when handling a distributed request. +func ExtractContext(ctx context.Context, carrier map[string]string) context.Context { + return illmobs.ExtractContext(ctx, carrier) +} From 56f640c165a953c227f322d71d25ec563c41ac90 Mon Sep 17 00:00:00 2001 From: yahya-mouman Date: Thu, 16 Jul 2026 13:18:16 +0200 Subject: [PATCH 5/6] test(llmobs): add end-to-end agent attribution round-trip Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: yahya-mouman --- internal/llmobs/propagation_test.go | 41 +++++++++++++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/internal/llmobs/propagation_test.go b/internal/llmobs/propagation_test.go index 0d33946278b..f70c7794ff5 100644 --- a/internal/llmobs/propagation_test.go +++ b/internal/llmobs/propagation_test.go @@ -190,3 +190,44 @@ func TestInjectContextReadsActiveSpan(t *testing.T) { t.Fatalf("InjectContext should inject active agent (my_agent,1000), got (%q,%q)", name, id) } } + +func TestAgentAttributionRoundTrip(t *testing.T) { + l := newTestLLMObsForResolve(t) + + t.Run("full-round-trip-with-name", func(t *testing.T) { + // Upstream: start an agent, inject from its context. + agent, upCtx := l.StartSpan(context.Background(), SpanKindAgent, "my_agent", StartSpanConfig{}) + carrier := map[string]string{} + InjectContext(upCtx, carrier) + + // Downstream: extract into a fresh context, start a child span. + downCtx := ExtractContext(context.Background(), carrier) + child, _ := l.StartSpan(downCtx, SpanKindTool, "downstream_tool", StartSpanConfig{}) + + if child.parentAgentName != "my_agent" { + t.Fatalf("child.parentAgentName = %q, want my_agent", child.parentAgentName) + } + if child.parentAgentSpanID != agent.SpanID() { + t.Fatalf("child.parentAgentSpanID = %q, want %q", child.parentAgentSpanID, agent.SpanID()) + } + }) + + t.Run("full-round-trip-id-only-yields-null-name", func(t *testing.T) { + // Upstream agent has a wire-unsafe name -> id-only propagates. + agent, upCtx := l.StartSpan(context.Background(), SpanKindAgent, "bad,name", StartSpanConfig{}) + carrier := map[string]string{} + InjectContext(upCtx, carrier) + + downCtx := ExtractContext(context.Background(), carrier) + child, _ := l.StartSpan(downCtx, SpanKindTool, "downstream_tool", StartSpanConfig{}) + + if child.parentAgentSpanID != agent.SpanID() { + t.Fatalf("child.parentAgentSpanID = %q, want %q", child.parentAgentSpanID, agent.SpanID()) + } + if child.parentAgentName != "" { + t.Fatalf("child.parentAgentName must be empty (id-only), got %q", child.parentAgentName) + } + // And the serialized child span emits pagent_name: null (verified in Task 2's + // id-only serialization test; here we confirm the in-process field state). + }) +} From 6ba23ee26d31441b8998851a647feffc332bec79 Mon Sep 17 00:00:00 2001 From: yahya-mouman Date: Mon, 20 Jul 2026 16:31:23 +0200 Subject: [PATCH 6/6] refactor(llmobs): drop InjectContext/ExtractContext, consolidate tests Defer distributed agent attribution propagation to a follow-up PR. The in-process attribution (resolveParentAgent + meta.agent_attribution serialization) is self-contained and ships without the distributed API. - Delete propagation.go and its tests (InjectContext, ExtractContext, injectAgentAttribution, extractAgentAttribution, agentNameWireSafe) - Remove InjectContext/ExtractContext from the public llmobs package - Merge resolve_parent_agent_test.go into start_span_agent_test.go Co-Authored-By: Claude Sonnet 4.6 --- internal/llmobs/propagation.go | 106 --------- internal/llmobs/propagation_test.go | 233 ------------------- internal/llmobs/resolve_parent_agent_test.go | 73 ------ internal/llmobs/start_span_agent_test.go | 63 +++++ internal/llmobs/util.go | 32 --- internal/llmobs/util_test.go | 39 ---- llmobs/llmobs.go | 18 -- 7 files changed, 63 insertions(+), 501 deletions(-) delete mode 100644 internal/llmobs/propagation.go delete mode 100644 internal/llmobs/propagation_test.go delete mode 100644 internal/llmobs/resolve_parent_agent_test.go delete mode 100644 internal/llmobs/util.go delete mode 100644 internal/llmobs/util_test.go diff --git a/internal/llmobs/propagation.go b/internal/llmobs/propagation.go deleted file mode 100644 index ecf1e9cbbb4..00000000000 --- a/internal/llmobs/propagation.go +++ /dev/null @@ -1,106 +0,0 @@ -// Unless explicitly stated otherwise all files in this repository are licensed -// under the Apache License Version 2.0. -// This product includes software developed at Datadog (https://www.datadoghq.com/). -// Copyright 2025 Datadog, Inc. - -package llmobs - -import ( - "context" - "strings" -) - -const ( - headerKeyDatadogTags = "x-datadog-tags" - tagKeyLLMObsPAgentSpanID = "_dd.p.llmobs_pagent_span_id" - tagKeyLLMObsPAgentName = "_dd.p.llmobs_pagent_name" - - // maxDatadogTagsLen bounds the total x-datadog-tags value length. The name - // tag is only appended when the result stays within this budget. - maxDatadogTagsLen = 512 -) - -// InjectContext writes the active span's outbound agent attribution into the -// carrier's x-datadog-tags value. Exported so the public llmobs package (which -// cannot read Span's unexported fields) can delegate here. -func InjectContext(ctx context.Context, carrier map[string]string) { - span, ok := ActiveLLMSpanFromContext(ctx) - if !ok || span == nil { - return - } - injectAgentAttribution(span, carrier) -} - -// ExtractContext reads agent attribution from the carrier and attaches it to a -// PropagatedLLMSpan on the returned context, preserving any fields already set. -func ExtractContext(ctx context.Context, carrier map[string]string) context.Context { - name, spanID, present := extractAgentAttribution(carrier) - if !present { - return ctx - } - - // Update, do not overwrite: preserve any existing MLApp/TraceID/SpanID. - var prop PropagatedLLMSpan - if existing, ok := PropagatedLLMSpanFromContext(ctx); ok && existing != nil { - prop = *existing - } - prop.ParentAgentName = name - prop.ParentAgentSpanID = spanID - return ContextWithPropagatedLLMSpan(ctx, &prop) -} - -// injectAgentAttribution resolves the span's outbound agent attribution and -// merges the resulting tags into the carrier's x-datadog-tags value. -// -// active span is an agent -> (self.name, self.SpanID()) -// otherwise -> (span.parentAgentName, span.parentAgentSpanID) -func injectAgentAttribution(span *Span, carrier map[string]string) { - var name, spanID string - if span.spanKind == SpanKindAgent { - name, spanID = span.name, span.SpanID() - } else { - name, spanID = span.parentAgentName, span.parentAgentSpanID - } - if spanID == "" { - return // no agent ancestor: write nothing - } - - existing := carrier[headerKeyDatadogTags] - parts := make([]string, 0, 2) - if existing != "" { - parts = append(parts, existing) - } - parts = append(parts, tagKeyLLMObsPAgentSpanID+"="+spanID) - - // Append the name only if wire-safe and within the total length budget. - if name != "" && agentNameWireSafe(name) { - candidate := append(parts, tagKeyLLMObsPAgentName+"="+name) - if joined := strings.Join(candidate, ","); len(joined) <= maxDatadogTagsLen { - parts = candidate - } - } - carrier[headerKeyDatadogTags] = strings.Join(parts, ",") -} - -// extractAgentAttribution parses the carrier's x-datadog-tags value for the -// llmobs agent attribution tags. present is true if the span id tag was found. -func extractAgentAttribution(carrier map[string]string) (name string, spanID string, present bool) { - header, ok := carrier[headerKeyDatadogTags] - if !ok || header == "" { - return "", "", false - } - for _, pair := range strings.Split(header, ",") { - k, v, found := strings.Cut(pair, "=") - if !found { - continue - } - switch k { - case tagKeyLLMObsPAgentSpanID: - spanID = v - present = true - case tagKeyLLMObsPAgentName: - name = v - } - } - return name, spanID, present -} diff --git a/internal/llmobs/propagation_test.go b/internal/llmobs/propagation_test.go deleted file mode 100644 index f70c7794ff5..00000000000 --- a/internal/llmobs/propagation_test.go +++ /dev/null @@ -1,233 +0,0 @@ -// Unless explicitly stated otherwise all files in this repository are licensed -// under the Apache License Version 2.0. -// This product includes software developed at Datadog (https://www.datadoghq.com/). -// Copyright 2025 Datadog, Inc. - -package llmobs - -import ( - "context" - "strings" - "testing" -) - -// tagValue returns the value of key within a comma-delimited x-datadog-tags string. -func tagValue(header, key string) (string, bool) { - for _, pair := range strings.Split(header, ",") { - k, v, found := strings.Cut(pair, "=") - if found && k == key { - return v, true - } - } - return "", false -} - -func TestInjectAgentAttribution(t *testing.T) { - t.Run("inject-from-agent", func(t *testing.T) { - agent := &Span{name: "my_agent", spanKind: SpanKindAgent} - agent.apm = testAPMSpan{spanID: "1000"} - carrier := map[string]string{} - injectAgentAttribution(agent, carrier) - - tags := carrier["x-datadog-tags"] - id, ok := tagValue(tags, "_dd.p.llmobs_pagent_span_id") - if !ok || id != "1000" { - t.Fatalf("expected span id 1000, got %q (ok=%v) in %q", id, ok, tags) - } - name, ok := tagValue(tags, "_dd.p.llmobs_pagent_name") - if !ok || name != "my_agent" { - t.Fatalf("expected name my_agent, got %q (ok=%v)", name, ok) - } - }) - - t.Run("inject-from-tool-under-agent-inherits", func(t *testing.T) { - tool := &Span{name: "tool", spanKind: SpanKindTool, parentAgentName: "my_agent", parentAgentSpanID: "1000"} - tool.apm = testAPMSpan{spanID: "2000"} - carrier := map[string]string{} - injectAgentAttribution(tool, carrier) - - tags := carrier["x-datadog-tags"] - id, _ := tagValue(tags, "_dd.p.llmobs_pagent_span_id") - name, _ := tagValue(tags, "_dd.p.llmobs_pagent_name") - if id != "1000" || name != "my_agent" { - t.Fatalf("tool should inject inherited agent (my_agent,1000), got (%q,%q)", name, id) - } - }) - - t.Run("no-agent-in-chain-writes-nothing", func(t *testing.T) { - tool := &Span{name: "tool", spanKind: SpanKindTool} - tool.apm = testAPMSpan{spanID: "2000"} - carrier := map[string]string{} - injectAgentAttribution(tool, carrier) - - if _, ok := carrier["x-datadog-tags"]; ok { - if _, hasID := tagValue(carrier["x-datadog-tags"], "_dd.p.llmobs_pagent_span_id"); hasID { - t.Fatalf("no agent ancestor: span_id tag must not be written, got %q", carrier["x-datadog-tags"]) - } - } - }) - - t.Run("unsafe-name-comma-writes-id-only", func(t *testing.T) { - agent := &Span{name: "bad,name", spanKind: SpanKindAgent} - agent.apm = testAPMSpan{spanID: "1000"} - carrier := map[string]string{} - injectAgentAttribution(agent, carrier) - - tags := carrier["x-datadog-tags"] - if _, ok := tagValue(tags, "_dd.p.llmobs_pagent_span_id"); !ok { - t.Fatalf("expected span id written, got %q", tags) - } - if _, ok := tagValue(tags, "_dd.p.llmobs_pagent_name"); ok { - t.Fatalf("unsafe name must be omitted, got %q", tags) - } - }) - - t.Run("name-with-equals-propagates", func(t *testing.T) { - agent := &Span{name: "k=v", spanKind: SpanKindAgent} - agent.apm = testAPMSpan{spanID: "1000"} - carrier := map[string]string{} - injectAgentAttribution(agent, carrier) - - name, ok := tagValue(carrier["x-datadog-tags"], "_dd.p.llmobs_pagent_name") - if !ok || name != "k=v" { - t.Fatalf("name with '=' must propagate, got %q (ok=%v)", name, ok) - } - }) - - t.Run("oversized-name-writes-id-only", func(t *testing.T) { - agent := &Span{name: strings.Repeat("a", 257), spanKind: SpanKindAgent} - agent.apm = testAPMSpan{spanID: "1000"} - carrier := map[string]string{} - injectAgentAttribution(agent, carrier) - - if _, ok := tagValue(carrier["x-datadog-tags"], "_dd.p.llmobs_pagent_span_id"); !ok { - t.Fatalf("expected span id written for oversized name") - } - if _, ok := tagValue(carrier["x-datadog-tags"], "_dd.p.llmobs_pagent_name"); ok { - t.Fatalf("oversized name must be omitted") - } - }) - - t.Run("merges-with-existing-tags", func(t *testing.T) { - agent := &Span{name: "my_agent", spanKind: SpanKindAgent} - agent.apm = testAPMSpan{spanID: "1000"} - carrier := map[string]string{"x-datadog-tags": "_dd.p.dm=-0"} - injectAgentAttribution(agent, carrier) - - tags := carrier["x-datadog-tags"] - if _, ok := tagValue(tags, "_dd.p.dm"); !ok { - t.Fatalf("existing tag _dd.p.dm must be preserved, got %q", tags) - } - if _, ok := tagValue(tags, "_dd.p.llmobs_pagent_span_id"); !ok { - t.Fatalf("new tag must be appended, got %q", tags) - } - }) -} - -func TestExtractAgentAttribution(t *testing.T) { - t.Run("extract-both", func(t *testing.T) { - carrier := map[string]string{ - "x-datadog-tags": "_dd.p.llmobs_pagent_span_id=1000,_dd.p.llmobs_pagent_name=my_agent", - } - name, id, present := extractAgentAttribution(carrier) - if !present || id != "1000" || name != "my_agent" { - t.Fatalf("expected (my_agent,1000,true), got (%q,%q,%v)", name, id, present) - } - }) - - t.Run("extract-id-only", func(t *testing.T) { - carrier := map[string]string{ - "x-datadog-tags": "_dd.p.llmobs_pagent_span_id=1000", - } - name, id, present := extractAgentAttribution(carrier) - if !present || id != "1000" || name != "" { - t.Fatalf("expected id-only (\"\",1000,true), got (%q,%q,%v)", name, id, present) - } - }) - - t.Run("extract-absent", func(t *testing.T) { - carrier := map[string]string{"x-datadog-tags": "_dd.p.dm=-0"} - _, _, present := extractAgentAttribution(carrier) - if present { - t.Fatalf("expected present=false when no llmobs tags") - } - }) -} - -func TestExtractContextUpdatesExistingPropagated(t *testing.T) { - // A prior ExtractContext (or manual set) put ML app + trace id on the context. - base := ContextWithPropagatedLLMSpan(context.Background(), &PropagatedLLMSpan{ - MLApp: "existing_app", - TraceID: "trace123", - SpanID: "span456", - }) - carrier := map[string]string{ - "x-datadog-tags": "_dd.p.llmobs_pagent_span_id=1000,_dd.p.llmobs_pagent_name=my_agent", - } - out := ExtractContext(base, carrier) - prop, ok := PropagatedLLMSpanFromContext(out) - if !ok { - t.Fatal("expected a propagated span on the returned context") - } - if prop.MLApp != "existing_app" || prop.TraceID != "trace123" || prop.SpanID != "span456" { - t.Fatalf("existing propagated fields must be preserved, got %+v", prop) - } - if prop.ParentAgentName != "my_agent" || prop.ParentAgentSpanID != "1000" { - t.Fatalf("agent fields must be set, got name=%q id=%q", prop.ParentAgentName, prop.ParentAgentSpanID) - } -} - -func TestInjectContextReadsActiveSpan(t *testing.T) { - agent := &Span{name: "my_agent", spanKind: SpanKindAgent} - agent.apm = testAPMSpan{spanID: "1000"} - ctx := contextWithActiveLLMSpan(context.Background(), agent) - carrier := map[string]string{} - InjectContext(ctx, carrier) - - name, _ := tagValue(carrier["x-datadog-tags"], "_dd.p.llmobs_pagent_name") - id, _ := tagValue(carrier["x-datadog-tags"], "_dd.p.llmobs_pagent_span_id") - if name != "my_agent" || id != "1000" { - t.Fatalf("InjectContext should inject active agent (my_agent,1000), got (%q,%q)", name, id) - } -} - -func TestAgentAttributionRoundTrip(t *testing.T) { - l := newTestLLMObsForResolve(t) - - t.Run("full-round-trip-with-name", func(t *testing.T) { - // Upstream: start an agent, inject from its context. - agent, upCtx := l.StartSpan(context.Background(), SpanKindAgent, "my_agent", StartSpanConfig{}) - carrier := map[string]string{} - InjectContext(upCtx, carrier) - - // Downstream: extract into a fresh context, start a child span. - downCtx := ExtractContext(context.Background(), carrier) - child, _ := l.StartSpan(downCtx, SpanKindTool, "downstream_tool", StartSpanConfig{}) - - if child.parentAgentName != "my_agent" { - t.Fatalf("child.parentAgentName = %q, want my_agent", child.parentAgentName) - } - if child.parentAgentSpanID != agent.SpanID() { - t.Fatalf("child.parentAgentSpanID = %q, want %q", child.parentAgentSpanID, agent.SpanID()) - } - }) - - t.Run("full-round-trip-id-only-yields-null-name", func(t *testing.T) { - // Upstream agent has a wire-unsafe name -> id-only propagates. - agent, upCtx := l.StartSpan(context.Background(), SpanKindAgent, "bad,name", StartSpanConfig{}) - carrier := map[string]string{} - InjectContext(upCtx, carrier) - - downCtx := ExtractContext(context.Background(), carrier) - child, _ := l.StartSpan(downCtx, SpanKindTool, "downstream_tool", StartSpanConfig{}) - - if child.parentAgentSpanID != agent.SpanID() { - t.Fatalf("child.parentAgentSpanID = %q, want %q", child.parentAgentSpanID, agent.SpanID()) - } - if child.parentAgentName != "" { - t.Fatalf("child.parentAgentName must be empty (id-only), got %q", child.parentAgentName) - } - // And the serialized child span emits pagent_name: null (verified in Task 2's - // id-only serialization test; here we confirm the in-process field state). - }) -} diff --git a/internal/llmobs/resolve_parent_agent_test.go b/internal/llmobs/resolve_parent_agent_test.go deleted file mode 100644 index 3f6f11ebc13..00000000000 --- a/internal/llmobs/resolve_parent_agent_test.go +++ /dev/null @@ -1,73 +0,0 @@ -// Unless explicitly stated otherwise all files in this repository are licensed -// under the Apache License Version 2.0. -// This product includes software developed at Datadog (https://www.datadoghq.com/). -// Copyright 2025 Datadog, Inc. - -package llmobs - -import "testing" - -func TestResolveParentAgent(t *testing.T) { - t.Run("no-parent", func(t *testing.T) { - name, id := resolveParentAgent(nil, nil) - if name != "" || id != "" { - t.Fatalf("expected empty, got name=%q id=%q", name, id) - } - }) - - t.Run("parent-is-agent", func(t *testing.T) { - parent := &Span{name: "my_agent", spanKind: SpanKindAgent} - parent.apm = testAPMSpan{spanID: "111"} - name, id := resolveParentAgent(parent, nil) - if name != "my_agent" || id != "111" { - t.Fatalf("expected (my_agent,111), got (%q,%q)", name, id) - } - }) - - t.Run("parent-other-kind-inherits", func(t *testing.T) { - parent := &Span{ - name: "tool_x", - spanKind: SpanKindTool, - parentAgentName: "top_agent", - parentAgentSpanID: "222", - } - parent.apm = testAPMSpan{spanID: "333"} - name, id := resolveParentAgent(parent, nil) - if name != "top_agent" || id != "222" { - t.Fatalf("expected inherited (top_agent,222), got (%q,%q)", name, id) - } - }) - - t.Run("propagated-parent", func(t *testing.T) { - prop := &PropagatedLLMSpan{ParentAgentName: "remote_agent", ParentAgentSpanID: "444"} - name, id := resolveParentAgent(nil, prop) - if name != "remote_agent" || id != "444" { - t.Fatalf("expected (remote_agent,444), got (%q,%q)", name, id) - } - }) - - t.Run("parent-takes-precedence-over-propagated", func(t *testing.T) { - parent := &Span{name: "local_agent", spanKind: SpanKindAgent} - parent.apm = testAPMSpan{spanID: "555"} - prop := &PropagatedLLMSpan{ParentAgentName: "remote_agent", ParentAgentSpanID: "444"} - name, id := resolveParentAgent(parent, prop) - if name != "local_agent" || id != "555" { - t.Fatalf("expected local (local_agent,555), got (%q,%q)", name, id) - } - }) -} - -// testAPMSpan is a minimal APMSpan implementation returning fixed values. -// All methods are implemented to avoid nil-pointer panics when embedded-interface -// patterns would otherwise dispatch to a nil value. -type testAPMSpan struct { - spanID string - traceID string -} - -func (s testAPMSpan) Finish(_ FinishAPMSpanConfig) {} -func (s testAPMSpan) AddLink(_ SpanLink) {} -func (s testAPMSpan) SpanID() string { return s.spanID } -func (s testAPMSpan) TraceID() string { return s.traceID } -func (s testAPMSpan) SetBaggageItem(_ string, _ string) {} -func (s testAPMSpan) BaggageItem(_ string) string { return "" } diff --git a/internal/llmobs/start_span_agent_test.go b/internal/llmobs/start_span_agent_test.go index 2c95db700f5..5f1fc363202 100644 --- a/internal/llmobs/start_span_agent_test.go +++ b/internal/llmobs/start_span_agent_test.go @@ -13,6 +13,19 @@ import ( "github.com/DataDog/dd-trace-go/v2/internal/llmobs/config" ) +// testAPMSpan is a minimal APMSpan implementation returning fixed values. +type testAPMSpan struct { + spanID string + traceID string +} + +func (s testAPMSpan) Finish(_ FinishAPMSpanConfig) {} +func (s testAPMSpan) AddLink(_ SpanLink) {} +func (s testAPMSpan) SpanID() string { return s.spanID } +func (s testAPMSpan) TraceID() string { return s.traceID } +func (s testAPMSpan) SetBaggageItem(_ string, _ string) {} +func (s testAPMSpan) BaggageItem(_ string) string { return "" } + // resolveTestTracer is a minimal Tracer that returns testAPMSpans with unique IDs. type resolveTestTracer struct { next int @@ -31,6 +44,56 @@ func newTestLLMObsForResolve(t *testing.T) *LLMObs { } } +func TestResolveParentAgent(t *testing.T) { + t.Run("no-parent", func(t *testing.T) { + name, id := resolveParentAgent(nil, nil) + if name != "" || id != "" { + t.Fatalf("expected empty, got name=%q id=%q", name, id) + } + }) + + t.Run("parent-is-agent", func(t *testing.T) { + parent := &Span{name: "my_agent", spanKind: SpanKindAgent} + parent.apm = testAPMSpan{spanID: "111"} + name, id := resolveParentAgent(parent, nil) + if name != "my_agent" || id != "111" { + t.Fatalf("expected (my_agent,111), got (%q,%q)", name, id) + } + }) + + t.Run("parent-other-kind-inherits", func(t *testing.T) { + parent := &Span{ + name: "tool_x", + spanKind: SpanKindTool, + parentAgentName: "top_agent", + parentAgentSpanID: "222", + } + parent.apm = testAPMSpan{spanID: "333"} + name, id := resolveParentAgent(parent, nil) + if name != "top_agent" || id != "222" { + t.Fatalf("expected inherited (top_agent,222), got (%q,%q)", name, id) + } + }) + + t.Run("propagated-parent", func(t *testing.T) { + prop := &PropagatedLLMSpan{ParentAgentName: "remote_agent", ParentAgentSpanID: "444"} + name, id := resolveParentAgent(nil, prop) + if name != "remote_agent" || id != "444" { + t.Fatalf("expected (remote_agent,444), got (%q,%q)", name, id) + } + }) + + t.Run("parent-takes-precedence-over-propagated", func(t *testing.T) { + parent := &Span{name: "local_agent", spanKind: SpanKindAgent} + parent.apm = testAPMSpan{spanID: "555"} + prop := &PropagatedLLMSpan{ParentAgentName: "remote_agent", ParentAgentSpanID: "444"} + name, id := resolveParentAgent(parent, prop) + if name != "local_agent" || id != "555" { + t.Fatalf("expected local (local_agent,555), got (%q,%q)", name, id) + } + }) +} + func TestStartSpanResolvesParentAgent(t *testing.T) { l := newTestLLMObsForResolve(t) ctx := context.Background() diff --git a/internal/llmobs/util.go b/internal/llmobs/util.go deleted file mode 100644 index b0891d66b87..00000000000 --- a/internal/llmobs/util.go +++ /dev/null @@ -1,32 +0,0 @@ -// Unless explicitly stated otherwise all files in this repository are licensed -// under the Apache License Version 2.0. -// This product includes software developed at Datadog (https://www.datadoghq.com/). -// Copyright 2025 Datadog, Inc. - -package llmobs - -// agentNameWireSafe reports whether name can be safely written as an -// x-datadog-tags tag value. The rules match the shared cross-language contract: -// -// - reject if the name exceeds 256 bytes -// - reject any byte outside the printable ASCII range [0x20, 0x7E] -// - reject a comma (the tagset delimiter) -// -// An equals sign is legal in tagset values (only illegal in keys) and is NOT -// rejected. This gate applies only to the wire tag; meta.agent_attribution -// always uses the real name. -func agentNameWireSafe(name string) bool { - if len(name) > 256 { - return false - } - for i := 0; i < len(name); i++ { - b := name[i] - if b < 0x20 || b > 0x7E { - return false - } - if b == ',' { - return false - } - } - return true -} diff --git a/internal/llmobs/util_test.go b/internal/llmobs/util_test.go deleted file mode 100644 index c92a9763043..00000000000 --- a/internal/llmobs/util_test.go +++ /dev/null @@ -1,39 +0,0 @@ -// Unless explicitly stated otherwise all files in this repository are licensed -// under the Apache License Version 2.0. -// This product includes software developed at Datadog (https://www.datadoghq.com/). -// Copyright 2025 Datadog, Inc. - -package llmobs - -import ( - "strings" - "testing" -) - -func TestAgentNameWireSafe(t *testing.T) { - cases := []struct { - name string - input string - want bool - }{ - {"plain-safe", "my_agent", true}, - {"empty", "", true}, - {"contains-comma", "my,agent", false}, - {"contains-equals-is-safe", "key=val", true}, - {"space-is-safe", "my agent", true}, - {"tilde-boundary-safe", "~", true}, - {"space-boundary-safe", " ", true}, - {"tab-non-printable", "my\tagent", false}, - {"newline-non-printable", "my\nagent", false}, - {"multibyte-utf8-unsafe", "agént", false}, - {"len-256-safe", strings.Repeat("a", 256), true}, - {"len-257-unsafe", strings.Repeat("a", 257), false}, - } - for _, tc := range cases { - t.Run(tc.name, func(t *testing.T) { - if got := agentNameWireSafe(tc.input); got != tc.want { - t.Fatalf("agentNameWireSafe(%q) = %v, want %v", tc.input, got, tc.want) - } - }) - } -} diff --git a/llmobs/llmobs.go b/llmobs/llmobs.go index 5d984aa6156..932120f653c 100644 --- a/llmobs/llmobs.go +++ b/llmobs/llmobs.go @@ -452,21 +452,3 @@ func isSpanKind(s Span, target illmobs.SpanKind) bool { return illmobs.SpanKind(s.Kind()) == target } -// InjectContext writes the active LLMObs span's agent attribution into the -// carrier so a downstream service can resolve agent_attribution for its spans. -// Call this alongside the APM tracer's Inject when making a distributed request. -// -// The carrier is a map of header name to value (e.g. an HTTP header map). Agent -// attribution is written into the x-datadog-tags value, merging with any value -// already present. -func InjectContext(ctx context.Context, carrier map[string]string) { - illmobs.InjectContext(ctx, carrier) -} - -// ExtractContext reads agent attribution written by an upstream InjectContext -// from the carrier and returns a context carrying it, so spans started from the -// returned context inherit the upstream agent attribution. Call this alongside -// the APM tracer's Extract when handling a distributed request. -func ExtractContext(ctx context.Context, carrier map[string]string) context.Context { - return illmobs.ExtractContext(ctx, carrier) -}