feat(llmobs): add WithAnnotatedCostTagKeys (DataDog#4728)

### What does this PR do?

Adds `llmobs.WithAnnotatedCostTagKeys`, allowing users to mark existing LLMObs span tag keys for propagation to LLM cost and token metrics.

Behavior mirrors dd-trace-py and dd-trace-js:
- Cost tag keys are accepted through manual annotation APIs.
- Keys must reference tags already present on the LLMObs span, either from the same annotation or a previous annotation.
- Valid keys are deduplicated across annotations.
- Valid cost tags are serialized on span events at `meta.metadata._dd.cost_tags`.
- Invalid/missing tag keys are skipped with a warning and do not fail annotation.
- Empty cost tag lists do not create `metadata._dd`.

This also adds LLMObs telemetry counters for cost-tag usage:
- `cost_tags.annotated`
- `cost_tags.submitted`

### Motivation

Bring dd-trace-go behavior in line with:
- DataDog/dd-trace-py#17628
- DataDog/dd-trace-js#8175

This lets users break down generated LLM cost and token metrics by custom span dimensions such as team, project, service line, or a run/debug identifier.

### Testing

Added focused coverage for:
- Public `WithAnnotatedCostTagKeys` API behavior.
- Same-call and prior-annotation tag references.
- Deduplication across annotations.
- Missing tag references being skipped.
- Empty cost tag lists not creating `metadata._dd`.
- Preserving existing `metadata._dd` fields.
- Cost tags on non-LLM spans serializing cleanly.
- Cost tags annotated after span finish being dropped.
- Option ordering with `WithAnnotatedSessionID`.


### Reviewer's Checklist
<!--
* Authors can use this list as a reference to ensure that there are no problems
  during the review but the signing off is to be done by the reviewer(s).
-->

- [ ] Changed code has unit tests for its functionality at or near 100% coverage.
- [ ] [System-Tests](https://github.com/DataDog/system-tests/) covering this feature have been added and enabled with the va.b.c-dev version tag.
- [ ] There is a benchmark for any new code, or changes to existing code.
- [ ] If this interacts with the agent in a new way, a system test has been added.
- [ ] New code is free of linting errors. You can check this by running `make lint` locally.
- [ ] New code doesn't break existing tests. You can check this by running `make test` locally.
- [ ] Add an appropriate team label so this PR gets put in the right place for the release notes.
- [ ] All generated files are up to date. You can check this by running `make generate` locally.
- [ ] Non-trivial go.mod changes, e.g. adding new modules, are reviewed by @DataDog/dd-trace-go-guild. Make sure all nested modules are up to date by running `make fix-modules` locally.

Unsure? Have a question? Request a review!


Co-authored-by: xinyuan.guo <xinyuan.guo@datadoghq.com>

Author: XG-xin

internal/llmobs/llmobs.go

@@ -96,6 +96,7 @@ type llmobsContext struct {
 	metadata map[string]any
 	metrics  map[string]float64
 	tags     map[string]string
+	costTags []string
 
 	// agent specific
 	agentManifest string
@@ -491,15 +492,14 @@ func (l *LLMObs) llmobsSpanEvent(span *Span) *transport.LLMObsSpanEvent {
 	}
 
 	metadata := span.llmCtx.metadata
-	if metadata == nil {
+	if len(metadata) > 0 {
+		metadata = maps.Clone(metadata)
+	} else {
 		metadata = make(map[string]any)
 	}
 	if spanKind == SpanKindAgent && span.llmCtx.agentManifest != "" {
 		metadata["agent_manifest"] = span.llmCtx.agentManifest
 	}
-	if len(metadata) > 0 {
-		meta["metadata"] = metadata
-	}
 
 	input := make(map[string]any)
 	output := make(map[string]any)
@@ -618,6 +618,12 @@ func (l *LLMObs) llmobsSpanEvent(span *Span) *transport.LLMObsSpanEvent {
 	}
 
 	maps.Copy(tags, span.llmCtx.tags)
+
+	setMetadataCostTags(metadata, validateCostTags(span, tags))
+	if len(metadata) > 0 {
+		meta["metadata"] = metadata
+	}
+
 	tagsSlice := make([]string, 0, len(tags))
 	for k, v := range tags {
 		tagsSlice = append(tagsSlice, fmt.Sprintf("%s:%s", k, v))
@@ -675,6 +681,53 @@ func (l *LLMObs) llmobsSpanEvent(span *Span) *transport.LLMObsSpanEvent {
 	return ev
 }
 
+// validateCostTags filters the span's annotated cost tags against the final
+// event tag set, drops entries that don't reference an emitted tag key, and
+// emits the cost-tags-submitted telemetry. It must run after the final tags
+// map is fully assembled, so cost tags referencing SDK-injected keys (e.g.
+// session_id from WithSessionID or propagation, integration, ml_app) are
+// accepted.
+func validateCostTags(span *Span, finalTags map[string]string) []string {
+	costTags := span.llmCtx.costTags
+	if len(costTags) == 0 {
+		return nil
+	}
+
+	validated := make([]string, 0, len(costTags))
+	missing := 0
+	for _, costTag := range costTags {
+		if _, ok := finalTags[costTag]; !ok {
+			log.Warn("llmobs: cost_tags entry %q must reference a key present in span tags. Skipping entry.", costTag)
+			missing++
+			continue
+		}
+		validated = append(validated, costTag)
+	}
+
+	if missing > 0 {
+		trackCostTagsSubmitted(span, missing, "annotate", "error", "missing_span_tag")
+	}
+	if len(validated) > 0 {
+		trackCostTagsSubmitted(span, len(validated), "annotate", "success", "none")
+	}
+	return validated
+}
+
+func setMetadataCostTags(metadata map[string]any, costTags []string) {
+	if len(costTags) == 0 {
+		return
+	}
+
+	ddMetadata, ok := metadata["_dd"].(map[string]any)
+	if ok {
+		ddMetadata = maps.Clone(ddMetadata)
+	} else {
+		ddMetadata = make(map[string]any)
+	}
+	ddMetadata["cost_tags"] = append([]string(nil), costTags...)
+	metadata["_dd"] = ddMetadata
+}
+
 func dropSpanEventIO(ev *transport.LLMObsSpanEvent) bool {
 	if ev == nil {
 		return false

internal/llmobs/llmobs_test.go

@@ -866,6 +866,158 @@ func TestSpanAnnotate(t *testing.T) {
 	}
 }
 
+func TestSpanAnnotateCostTags(t *testing.T) {
+	t.Run("sets-cost-tags-on-span-event-metadata", func(t *testing.T) {
+		tt, ll := testTracer(t)
+		span, _ := ll.StartSpan(context.Background(), llmobs.SpanKindLLM, "", llmobs.StartSpanConfig{})
+		span.Annotate(llmobs.SpanAnnotations{
+			Tags:     map[string]string{"team": "ml", "feature": "chatbot", "debug_id": "abc"},
+			CostTags: []string{"team", "feature"},
+		})
+		span.Finish(llmobs.FinishSpanConfig{})
+
+		llmSpans := tt.WaitForLLMObsSpans(t, 1)
+		assert.Equal(t, []any{"team", "feature"}, costTagsFromSpan(t, llmSpans[0]))
+	})
+
+	t.Run("dedupes-cost-tags-across-annotations", func(t *testing.T) {
+		tt, ll := testTracer(t)
+		span, _ := ll.StartSpan(context.Background(), llmobs.SpanKindLLM, "", llmobs.StartSpanConfig{})
+		span.Annotate(llmobs.SpanAnnotations{
+			Tags:     map[string]string{"team": "ml", "feature": "chatbot"},
+			CostTags: []string{"team", "feature", "team"},
+		})
+		span.Annotate(llmobs.SpanAnnotations{
+			Tags:     map[string]string{"project": "alpha"},
+			CostTags: []string{"feature", "project"},
+		})
+		span.Finish(llmobs.FinishSpanConfig{})
+
+		llmSpans := tt.WaitForLLMObsSpans(t, 1)
+		assert.Equal(t, []any{"team", "feature", "project"}, costTagsFromSpan(t, llmSpans[0]))
+	})
+
+	t.Run("skips-cost-tags-that-do-not-reference-span-tags", func(t *testing.T) {
+		tt, ll := testTracer(t)
+		span, _ := ll.StartSpan(context.Background(), llmobs.SpanKindLLM, "", llmobs.StartSpanConfig{})
+		span.Annotate(llmobs.SpanAnnotations{
+			Tags:     map[string]string{"team": "ml"},
+			CostTags: []string{"team", "missing"},
+		})
+		span.Finish(llmobs.FinishSpanConfig{})
+
+		llmSpans := tt.WaitForLLMObsSpans(t, 1)
+		assert.Equal(t, []any{"team"}, costTagsFromSpan(t, llmSpans[0]))
+	})
+
+	t.Run("references-tags-from-a-prior-annotation", func(t *testing.T) {
+		tt, ll := testTracer(t)
+		span, _ := ll.StartSpan(context.Background(), llmobs.SpanKindLLM, "", llmobs.StartSpanConfig{})
+		span.Annotate(llmobs.SpanAnnotations{Tags: map[string]string{"team": "ml"}})
+		span.Annotate(llmobs.SpanAnnotations{CostTags: []string{"team"}})
+		span.Finish(llmobs.FinishSpanConfig{})
+
+		llmSpans := tt.WaitForLLMObsSpans(t, 1)
+		assert.Equal(t, []any{"team"}, costTagsFromSpan(t, llmSpans[0]))
+	})
+
+	t.Run("empty-list-does-not-create-metadata-dd", func(t *testing.T) {
+		tt, ll := testTracer(t)
+		span, _ := ll.StartSpan(context.Background(), llmobs.SpanKindLLM, "", llmobs.StartSpanConfig{})
+		span.Annotate(llmobs.SpanAnnotations{
+			Tags:     map[string]string{"team": "ml"},
+			CostTags: []string{},
+		})
+		span.Finish(llmobs.FinishSpanConfig{})
+
+		llmSpans := tt.WaitForLLMObsSpans(t, 1)
+		assert.NotContains(t, llmSpans[0].Meta, "metadata")
+	})
+
+	t.Run("preserves-existing-metadata-dd-fields", func(t *testing.T) {
+		tt, ll := testTracer(t)
+		span, _ := ll.StartSpan(context.Background(), llmobs.SpanKindLLM, "", llmobs.StartSpanConfig{})
+		span.Annotate(llmobs.SpanAnnotations{
+			Metadata: map[string]any{"_dd": map[string]any{"existing": "value"}},
+			Tags:     map[string]string{"team": "ml"},
+			CostTags: []string{"team"},
+		})
+		span.Finish(llmobs.FinishSpanConfig{})
+
+		llmSpans := tt.WaitForLLMObsSpans(t, 1)
+		metadata := llmSpans[0].Meta["metadata"].(map[string]any)
+		ddMetadata := metadata["_dd"].(map[string]any)
+		assert.Equal(t, "value", ddMetadata["existing"])
+		assert.Equal(t, []any{"team"}, ddMetadata["cost_tags"])
+	})
+
+	t.Run("serializes-cost-tags-on-non-llm-span", func(t *testing.T) {
+		tt, ll := testTracer(t)
+		span, _ := ll.StartSpan(context.Background(), llmobs.SpanKindWorkflow, "", llmobs.StartSpanConfig{})
+		span.Annotate(llmobs.SpanAnnotations{
+			Tags:     map[string]string{"team": "ml"},
+			CostTags: []string{"team"},
+		})
+		span.Finish(llmobs.FinishSpanConfig{})
+
+		llmSpans := tt.WaitForLLMObsSpans(t, 1)
+		assert.Equal(t, "workflow", llmSpans[0].Meta["span.kind"])
+		assert.Equal(t, []any{"team"}, costTagsFromSpan(t, llmSpans[0]))
+	})
+
+	t.Run("references-session-id-from-start-config", func(t *testing.T) {
+		tt, ll := testTracer(t)
+		span, _ := ll.StartSpan(context.Background(), llmobs.SpanKindLLM, "", llmobs.StartSpanConfig{
+			SessionID: "session-123",
+		})
+		span.Annotate(llmobs.SpanAnnotations{
+			CostTags: []string{"session_id"},
+		})
+		span.Finish(llmobs.FinishSpanConfig{})
+
+		llmSpans := tt.WaitForLLMObsSpans(t, 1)
+		assert.Equal(t, "session-123", llmSpans[0].SessionID)
+		assert.Equal(t, []any{"session_id"}, costTagsFromSpan(t, llmSpans[0]))
+	})
+
+	t.Run("references-sdk-injected-tags", func(t *testing.T) {
+		tt, ll := testTracer(t)
+		span, _ := ll.StartSpan(context.Background(), llmobs.SpanKindLLM, "", llmobs.StartSpanConfig{
+			MLApp: "custom-app",
+		})
+		span.Annotate(llmobs.SpanAnnotations{
+			CostTags: []string{"ml_app", "language"},
+		})
+		span.Finish(llmobs.FinishSpanConfig{})
+
+		llmSpans := tt.WaitForLLMObsSpans(t, 1)
+		assert.Equal(t, []any{"ml_app", "language"}, costTagsFromSpan(t, llmSpans[0]))
+	})
+
+	t.Run("drops-cost-tags-annotated-after-finish", func(t *testing.T) {
+		tt, ll := testTracer(t)
+		span, _ := ll.StartSpan(context.Background(), llmobs.SpanKindLLM, "", llmobs.StartSpanConfig{})
+		span.Annotate(llmobs.SpanAnnotations{Tags: map[string]string{"team": "ml"}})
+		span.Finish(llmobs.FinishSpanConfig{})
+		span.Annotate(llmobs.SpanAnnotations{CostTags: []string{"team"}})
+
+		llmSpans := tt.WaitForLLMObsSpans(t, 1)
+		assert.NotContains(t, llmSpans[0].Meta, "metadata")
+	})
+}
+
+func costTagsFromSpan(t *testing.T, span llmobstransport.LLMObsSpanEvent) []any {
+	t.Helper()
+
+	metadata, ok := span.Meta["metadata"].(map[string]any)
+	require.True(t, ok, "metadata should be present")
+	ddMetadata, ok := metadata["_dd"].(map[string]any)
+	require.True(t, ok, "metadata._dd should be present")
+	costTags, ok := ddMetadata["cost_tags"].([]any)
+	require.True(t, ok, "metadata._dd.cost_tags should be present")
+	return costTags
+}
+
 func TestSpanTruncation(t *testing.T) {
 	t.Run("text-input", func(t *testing.T) {
 		tt, ll := testTracer(t)

internal/llmobs/span.go

@@ -8,6 +8,7 @@ package llmobs
 import (
 	"encoding/json"
 	"maps"
+	"slices"
 	"sync"
 	"time"
 
@@ -220,6 +221,9 @@ type SpanAnnotations struct {
 	Metrics map[string]float64
 	// Tags contains string tags key-value pairs.
 	Tags map[string]string
+	// CostTags contains tag keys to propagate to LLMObs cost and token metrics.
+	// Each key must reference a tag already present on the span.
+	CostTags []string
 }
 
 // Span represents an LLMObs span with its associated metadata and context.
@@ -358,6 +362,15 @@ func (s *Span) Annotate(a SpanAnnotations) {
 		}
 	}
 
+	if a.CostTags != nil {
+		trackCostTagsAnnotated(s, "annotate")
+		for _, costTag := range a.CostTags {
+			if !slices.Contains(s.llmCtx.costTags, costTag) {
+				s.llmCtx.costTags = append(s.llmCtx.costTags, costTag)
+			}
+		}
+	}
+
 	if a.Prompt != nil {
 		if s.spanKind != SpanKindLLM {
 			log.Warn("llmobs: input prompt can only be annotated on llm spans, ignoring")

internal/llmobs/telemetry.go

@@ -26,6 +26,8 @@ const (
 	telemetryMetricAnnotations       = "annotations"
 	telemetryMetricEvalsSubmitted    = "evals_submitted"
 	telemetryMetricUserFlushes       = "user_flush"
+	telemetryMetricCostTagsAnnotated = "cost_tags.annotated"
+	telemetryMetricCostTagsSubmitted = "cost_tags.submitted"
 )
 
 var telemetryErrorTypes = map[error]string{
@@ -143,6 +145,26 @@ func trackSpanAnnotations(span *Span, err error) {
 	telemetry.Count(telemetry.NamespaceMLObs, telemetryMetricAnnotations, tags).Submit(1)
 }
 
+func trackCostTagsAnnotated(span *Span, source string) {
+	if telemetry.Disabled() {
+		return
+	}
+
+	telemetry.Count(telemetry.NamespaceMLObs, telemetryMetricCostTagsAnnotated, costTagsTelemetryTags(span, source)).Submit(1)
+}
+
+func trackCostTagsSubmitted(span *Span, count int, source, state, reason string) {
+	if telemetry.Disabled() || count == 0 {
+		return
+	}
+
+	tags := append(costTagsTelemetryTags(span, source),
+		"state:"+state,
+		"reason:"+reason,
+	)
+	telemetry.Count(telemetry.NamespaceMLObs, telemetryMetricCostTagsSubmitted, tags).Submit(float64(count))
+}
+
 func trackSubmitEvaluationMetric(metric *transport.LLMObsMetric, err error) {
 	if telemetry.Disabled() {
 		return
@@ -169,6 +191,29 @@ func trackUserFlush() {
 	userFlushHandle.Submit(1)
 }
 
+func costTagsTelemetryTags(span *Span, source string) []string {
+	spanKind := "N/A"
+	mlApp := "N/A"
+	modelProvider := "N/A"
+	if span != nil {
+		if span.spanKind != "" {
+			spanKind = string(span.spanKind)
+		}
+		if span.mlApp != "" {
+			mlApp = span.mlApp
+		}
+		if span.llmCtx.modelProvider != "" {
+			modelProvider = span.llmCtx.modelProvider
+		}
+	}
+	return []string{
+		"span_kind:" + spanKind,
+		"source:" + source,
+		"ml_app:" + mlApp,
+		"model_provider:" + modelProvider,
+	}
+}
+
 func spanEventTags(event *transport.LLMObsSpanEvent) []string {
 	spanKind := "N/A"
 	if meta, ok := event.Meta["span.kind"]; ok {