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..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 { @@ -770,6 +782,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 +848,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/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 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..5f1fc363202 --- /dev/null +++ b/internal/llmobs/start_span_agent_test.go @@ -0,0 +1,157 @@ +// 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" +) + +// 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 +} + +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 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() + + 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) + } + }) +} diff --git a/llmobs/llmobs.go b/llmobs/llmobs.go index bdf5eff0a92..932120f653c 100644 --- a/llmobs/llmobs.go +++ b/llmobs/llmobs.go @@ -451,3 +451,4 @@ func parseAnnotateOptions(opts ...AnnotateOption) illmobs.SpanAnnotations { func isSpanKind(s Span, target illmobs.SpanKind) bool { return illmobs.SpanKind(s.Kind()) == target } +