DDLLMObsSpanTest.groovy

package datadog.trace.llmobs.domain

import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace

import datadog.trace.agent.tooling.TracerInstaller
import datadog.trace.api.DDTags
import datadog.trace.api.DDTraceApiInfo
import datadog.trace.api.IdGenerationStrategy
import datadog.trace.api.WellKnownTags
import datadog.trace.api.telemetry.LLMObsMetricCollector
import datadog.trace.api.llmobs.LLMObs
import datadog.trace.api.llmobs.LLMObsSpan
import datadog.trace.api.llmobs.LLMObsTags
import datadog.trace.bootstrap.instrumentation.api.AgentSpan
import datadog.trace.bootstrap.instrumentation.api.AgentTracer
import datadog.trace.bootstrap.instrumentation.api.Tags
import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString
import datadog.trace.core.CoreTracer
import datadog.trace.test.util.DDSpecification
import org.apache.groovy.util.Maps
import spock.lang.Shared

class DDLLMObsSpanTest  extends DDSpecification{
  @SuppressWarnings('PropertyName')
  @Shared
  AgentTracer.TracerAPI TEST_TRACER

  void setupSpec() {
    TEST_TRACER =
      Spy(
      CoreTracer.builder()
      .idGenerationStrategy(IdGenerationStrategy.fromName("SEQUENTIAL"))
      .build())
    TracerInstaller.forceInstallGlobalTracer(TEST_TRACER)

    TEST_TRACER.startSpan(*_) >> {
      def agentSpan = callRealMethod()
      agentSpan
    }
  }

  void cleanupSpec() {
    TEST_TRACER?.close()
  }

  void setup() {
    assert TEST_TRACER.activeSpan() == null: "Span is active before test has started: " + TEST_TRACER.activeSpan()
    TEST_TRACER.flush()
  }

  void cleanup() {
    TEST_TRACER.flush()
  }

  // Prefix for tags
  private static final String LLMOBS_TAG_PREFIX = "_ml_obs_tag."
  // Prefix for metrics
  private static final String LLMOBS_METRIC_PREFIX = "_ml_obs_metric."

  // internal tags to be prefixed
  private static final String INPUT = LLMOBS_TAG_PREFIX + "input"
  private static final String OUTPUT = LLMOBS_TAG_PREFIX + "output"
  private static final String METADATA = LLMOBS_TAG_PREFIX + LLMObsTags.METADATA
  private static final String TOOL_DEFINITIONS = LLMOBS_TAG_PREFIX + LLMObsTags.TOOL_DEFINITIONS


  def "test span simple"() {
    setup:
    def test = llmObsSpan(Tags.LLMOBS_WORKFLOW_SPAN_KIND, "test-span")

    when:
    def input = "test input"
    def output = "test output"
    // initial set
    test.annotateIO(input, output)
    test.setMetadata(Maps.of("sport", "baseball", "price_data", Maps.of("gpt4", 100)))
    test.setMetrics(Maps.of("rank", 1))
    test.setMetric("likelihood", 0.1)
    test.setTag("DOMAIN", "north-america")
    test.setTags(Maps.of("bulk1", 1, "bulk2", "2"))
    def errMsg = "mr brady"
    test.setErrorMessage(errMsg)

    then:
    def innerSpan = (AgentSpan)test.span
    Tags.LLMOBS_WORKFLOW_SPAN_KIND.equals(innerSpan.getTag(LLMOBS_TAG_PREFIX + "span.kind"))

    null == innerSpan.getTag("input")
    input.equals(innerSpan.getTag(INPUT))
    null == innerSpan.getTag("output")
    output.equals(innerSpan.getTag(OUTPUT))

    null == innerSpan.getTag("metadata")
    def expectedMetadata = Maps.of("sport", "baseball", "price_data", Maps.of("gpt4", 100))
    expectedMetadata.equals(innerSpan.getTag(METADATA))

    null == innerSpan.getTag("rank")
    def rankMetric = innerSpan.getTag(LLMOBS_METRIC_PREFIX + "rank")
    rankMetric instanceof Number && 1 == (int)rankMetric

    null == innerSpan.getTag("likelihood")
    def likelihoodMetric = innerSpan.getTag(LLMOBS_METRIC_PREFIX + "likelihood")
    likelihoodMetric instanceof Number
    0.1 == (double)likelihoodMetric

    null == innerSpan.getTag("DOMAIN")
    def domain = innerSpan.getTag(LLMOBS_TAG_PREFIX + "DOMAIN")
    domain instanceof String
    "north-america".equals((String)domain)

    null == innerSpan.getTag("bulk1")
    def tagBulk1 = innerSpan.getTag(LLMOBS_TAG_PREFIX + "bulk1")
    tagBulk1 instanceof Number
    1 == ((int)tagBulk1)

    null == innerSpan.getTag("bulk2")
    def tagBulk2 = innerSpan.getTag(LLMOBS_TAG_PREFIX + "bulk2")
    tagBulk2 instanceof String
    "2".equals((String)tagBulk2)

    innerSpan.isError()
    innerSpan.getTag(DDTags.ERROR_MSG) instanceof String
    errMsg.equals(innerSpan.getTag(DDTags.ERROR_MSG))

    null == innerSpan.getTag("env")
    def tagEnv = innerSpan.getTag(LLMOBS_TAG_PREFIX + "env")
    tagEnv instanceof UTF8BytesString
    "test-env" == tagEnv.toString()

    null == innerSpan.getTag("service")
    def tagSvc = innerSpan.getTag(LLMOBS_TAG_PREFIX + "service")
    tagSvc instanceof UTF8BytesString
    "test-svc" == tagSvc.toString()

    null == innerSpan.getTag("version")
    def tagVersion = innerSpan.getTag(LLMOBS_TAG_PREFIX + "version")
    tagVersion instanceof UTF8BytesString
    "v1" == tagVersion.toString()

    DDTraceApiInfo.VERSION == innerSpan.getTag(LLMOBS_TAG_PREFIX + "ddtrace.version")
  }

  def "test span with overwrites"() {
    setup:
    def test = llmObsSpan(Tags.LLMOBS_AGENT_SPAN_KIND, "test-span")

    when:
    def input = "test input"
    // initial set
    test.annotateIO(input, "test output")
    // this should be a no-op
    test.annotateIO("", "")
    // this should replace the initial output
    def expectedOutput = Arrays.asList(Maps.of("role", "user", "content", "how much is gas"))
    test.annotateIO(null, expectedOutput)

    // initial set
    test.setMetadata(Maps.of("sport", "baseball", "price_data", Maps.of("gpt4", 100)))
    // this should replace baseball with hockey
    test.setMetadata(Maps.of("sport", "hockey"))
    // this should add a new key
    test.setMetadata(Maps.of("temperature", 30))

    // initial set
    test.setMetrics(Maps.of("rank", 1))
    // this should replace the metric
    test.setMetric("rank", 10)

    // initial set
    test.setTag("DOMAIN", "north-america")
    // add and replace
    test.setTags(Maps.of("bulk1", 1, "DOMAIN", "europe"))

    def throwableMsg = "false positive"
    test.addThrowable(new Throwable(throwableMsg))
    test.setError(false)

    then:
    def innerSpan = (AgentSpan)test.span
    Tags.LLMOBS_AGENT_SPAN_KIND.equals(innerSpan.getTag(LLMOBS_TAG_PREFIX + "span.kind"))

    null == innerSpan.getTag("input")
    input.equals(innerSpan.getTag(INPUT))
    null == innerSpan.getTag("output")
    expectedOutput.equals(innerSpan.getTag(OUTPUT))

    null == innerSpan.getTag("metadata")
    def expectedMetadata = Maps.of("sport", "hockey", "price_data", Maps.of("gpt4", 100), "temperature", 30)
    expectedMetadata.equals(innerSpan.getTag(METADATA))

    null == innerSpan.getTag("rank")
    def rankMetric = innerSpan.getTag(LLMOBS_METRIC_PREFIX + "rank")
    rankMetric instanceof Number && 10 == (int)rankMetric

    null == innerSpan.getTag("DOMAIN")
    def domain = innerSpan.getTag(LLMOBS_TAG_PREFIX + "DOMAIN")
    domain instanceof String
    "europe".equals((String)domain)

    null == innerSpan.getTag("bulk1")
    def tagBulk1 = innerSpan.getTag(LLMOBS_TAG_PREFIX + "bulk1")
    tagBulk1 instanceof Number
    1 == ((int)tagBulk1)

    !innerSpan.isError()
    innerSpan.getTag(DDTags.ERROR_MSG) instanceof String
    throwableMsg.equals(innerSpan.getTag(DDTags.ERROR_MSG))
    innerSpan.getTag(DDTags.ERROR_STACK) instanceof String
    ((String)innerSpan.getTag(DDTags.ERROR_STACK)).contains(throwableMsg)


    null == innerSpan.getTag("env")
    def tagEnv = innerSpan.getTag(LLMOBS_TAG_PREFIX + "env")
    tagEnv instanceof UTF8BytesString
    "test-env" == tagEnv.toString()

    null == innerSpan.getTag("service")
    def tagSvc = innerSpan.getTag(LLMOBS_TAG_PREFIX + "service")
    tagSvc instanceof UTF8BytesString
    "test-svc" == tagSvc.toString()

    null == innerSpan.getTag("version")
    def tagVersion = innerSpan.getTag(LLMOBS_TAG_PREFIX + "version")
    tagVersion instanceof UTF8BytesString
    "v1" == tagVersion.toString()

    DDTraceApiInfo.VERSION == innerSpan.getTag(LLMOBS_TAG_PREFIX + "ddtrace.version")
  }

  def "test llm span string input formatted to messages"() {
    setup:
    def test = llmObsSpan(Tags.LLMOBS_LLM_SPAN_KIND, "test-span")

    when:
    def input = "test input"
    def output = "test output"
    // initial set
    test.annotateIO(input, output)

    then:
    def innerSpan = (AgentSpan)test.span
    Tags.LLMOBS_LLM_SPAN_KIND.equals(innerSpan.getTag(LLMOBS_TAG_PREFIX + "span.kind"))

    null == innerSpan.getTag("input")
    def spanInput = innerSpan.getTag(INPUT)
    spanInput instanceof List
    ((List)spanInput).size() == 1
    spanInput.get(0) instanceof LLMObs.LLMMessage
    def expectedInputMsg = LLMObs.LLMMessage.from("unknown", input)
    expectedInputMsg.getContent().equals(input)
    expectedInputMsg.getRole().equals("unknown")
    expectedInputMsg.getToolCalls().equals(null)

    null == innerSpan.getTag("output")
    def spanOutput = innerSpan.getTag(OUTPUT)
    spanOutput instanceof List
    ((List)spanOutput).size() == 1
    spanOutput.get(0) instanceof LLMObs.LLMMessage
    def expectedOutputMsg = LLMObs.LLMMessage.from("unknown", output)
    expectedOutputMsg.getContent().equals(output)
    expectedOutputMsg.getRole().equals("unknown")
    expectedOutputMsg.getToolCalls().equals(null)


    null == innerSpan.getTag("env")
    def tagEnv = innerSpan.getTag(LLMOBS_TAG_PREFIX + "env")
    tagEnv instanceof UTF8BytesString
    "test-env" == tagEnv.toString()

    null == innerSpan.getTag("service")
    def tagSvc = innerSpan.getTag(LLMOBS_TAG_PREFIX + "service")
    tagSvc instanceof UTF8BytesString
    "test-svc" == tagSvc.toString()

    null == innerSpan.getTag("version")
    def tagVersion = innerSpan.getTag(LLMOBS_TAG_PREFIX + "version")
    tagVersion instanceof UTF8BytesString
    "v1" == tagVersion.toString()

    DDTraceApiInfo.VERSION == innerSpan.getTag(LLMOBS_TAG_PREFIX + "ddtrace.version")
  }

  def "test llm span with messages"() {
    setup:
    def test = llmObsSpan(Tags.LLMOBS_LLM_SPAN_KIND, "test-span")

    when:
    def inputMsg =  LLMObs.LLMMessage.from("user", "input")
    def outputMsg = LLMObs.LLMMessage.from("assistant", "output", Arrays.asList(LLMObs.ToolCall.from("weather-tool", "function", "6176241000", Maps.of("location", "paris"))))
    // initial set
    test.annotateIO(Arrays.asList(inputMsg), Arrays.asList(outputMsg))

    then:
    def innerSpan = (AgentSpan)test.span
    Tags.LLMOBS_LLM_SPAN_KIND.equals(innerSpan.getTag(LLMOBS_TAG_PREFIX + "span.kind"))

    null == innerSpan.getTag("input")
    def spanInput = innerSpan.getTag(INPUT)
    spanInput instanceof List
    ((List)spanInput).size() == 1
    def spanInputMsg = spanInput.get(0)
    spanInputMsg instanceof LLMObs.LLMMessage
    spanInputMsg.getContent().equals(inputMsg.getContent())
    spanInputMsg.getRole().equals("user")
    spanInputMsg.getToolCalls().equals(null)

    null == innerSpan.getTag("output")
    def spanOutput = innerSpan.getTag(OUTPUT)
    spanOutput instanceof List
    ((List)spanOutput).size() == 1
    def spanOutputMsg = spanOutput.get(0)
    spanOutputMsg instanceof LLMObs.LLMMessage
    spanOutputMsg.getContent().equals(outputMsg.getContent())
    spanOutputMsg.getRole().equals("assistant")
    spanOutputMsg.getToolCalls().size() == 1
    def toolCall = spanOutputMsg.getToolCalls().get(0)
    toolCall.getName().equals("weather-tool")
    toolCall.getType().equals("function")
    toolCall.getToolId().equals("6176241000")
    def expectedToolArgs = Maps.of("location", "paris")
    toolCall.getArguments().equals(expectedToolArgs)

    null == innerSpan.getTag("env")
    def tagEnv = innerSpan.getTag(LLMOBS_TAG_PREFIX + "env")
    tagEnv instanceof UTF8BytesString
    "test-env" == tagEnv.toString()

    null == innerSpan.getTag("service")
    def tagSvc = innerSpan.getTag(LLMOBS_TAG_PREFIX + "service")
    tagSvc instanceof UTF8BytesString
    "test-svc" == tagSvc.toString()

    null == innerSpan.getTag("version")
    def tagVersion = innerSpan.getTag(LLMOBS_TAG_PREFIX + "version")
    tagVersion instanceof UTF8BytesString
    "v1" == tagVersion.toString()

    DDTraceApiInfo.VERSION == innerSpan.getTag(LLMOBS_TAG_PREFIX + "ddtrace.version")
  }

  def "test llm span with tool definitions"() {
    setup:
    def test = llmObsSpan(Tags.LLMOBS_LLM_SPAN_KIND, "test-span")
    def schema = Maps.of(
      "type", "object",
      "properties", Maps.of("location", Maps.of("type", "string")))
    def toolDefinition =
      LLMObs.ToolDefinition.from("get_weather", "Get the weather by location", schema, "1.2.3")

    when:
    test.setToolDefinitions(Arrays.asList(
      toolDefinition,
      LLMObs.ToolDefinition.from("get_time"),
      LLMObs.ToolDefinition.from(""),
      LLMObs.ToolDefinition.from(null),
      null))

    then:
    def innerSpan = (AgentSpan)test.span
    def toolDefinitions = innerSpan.getTag(TOOL_DEFINITIONS)
    toolDefinitions instanceof List
    toolDefinitions.size() == 2
    toolDefinitions.get(0).is(toolDefinition)
    toolDefinitions.get(1).name == "get_time"
    toolDefinitions.get(1).description == null
    toolDefinitions.get(1).schema == null
    toolDefinitions.get(1).version == null
  }

  def "tool definitions overwrite prior annotations and ignore invalid definitions"() {
    setup:
    def test = llmObsSpan(Tags.LLMOBS_LLM_SPAN_KIND, "test-span")
    test.setToolDefinitions(Arrays.asList(LLMObs.ToolDefinition.from("first")))

    when:
    test.setToolDefinitions(Arrays.asList(
      LLMObs.ToolDefinition.from("second")))

    then:
    def innerSpan = (AgentSpan)test.span
    innerSpan.getTag(TOOL_DEFINITIONS)*.name == ["second"]

    when:
    test.setToolDefinitions(Arrays.asList(
      LLMObs.ToolDefinition.from(""),
      LLMObs.ToolDefinition.from(null),
      null))

    then:
    innerSpan.getTag(TOOL_DEFINITIONS)*.name == ["second"]
  }

  def "null and empty tool definition lists do not overwrite prior annotations"() {
    setup:
    def test = llmObsSpan(Tags.LLMOBS_LLM_SPAN_KIND, "test-span")
    test.setToolDefinitions(Arrays.asList(LLMObs.ToolDefinition.from("first")))

    when:
    test.setToolDefinitions(toolDefinitions)

    then:
    def innerSpan = (AgentSpan)test.span
    innerSpan.getTag(TOOL_DEFINITIONS)*.name == ["first"]

    where:
    toolDefinitions << [null, Collections.emptyList()]
  }

  def "tool definitions cannot be annotated after the span finishes"() {
    setup:
    def test = llmObsSpan(Tags.LLMOBS_LLM_SPAN_KIND, "test-span")
    test.setToolDefinitions(Arrays.asList(LLMObs.ToolDefinition.from("first")))
    test.finish()

    when:
    test.setToolDefinitions(Arrays.asList(LLMObs.ToolDefinition.from("second")))

    then:
    def innerSpan = (AgentSpan)test.span
    innerSpan.getTag(TOOL_DEFINITIONS)*.name == ["first"]
  }

  def "finish records span.finished telemetry when LLMObs enabled"() {
    setup:
    LLMObsMetricCollector collector = LLMObsMetricCollector.get()
    collector.drain()

    when:
    llmObsSpan(Tags.LLMOBS_WORKFLOW_SPAN_KIND, "workflow-span").finish()

    then:
    def metrics = collector.drain()
    metrics.size() == 1

    and:
    def m = metrics[0]
    m.namespace == 'mlobs'
    m.metricName == 'span.finished'
    m.type == 'count'
    m.value == 1
    m.tags.contains('integration:llmobs')
    m.tags.contains('span_kind:workflow')
    m.tags.contains('autoinstrumented:0')
    m.tags.contains('is_root_span:1')
  }

  def "finish records span.finished telemetry for non-root span when LLMObs enabled"() {
    setup:
    LLMObsMetricCollector collector = LLMObsMetricCollector.get()
    collector.drain()

    when:
    runUnderTrace("parent") {
      llmObsSpan(Tags.LLMOBS_LLM_SPAN_KIND, "child-llm").finish()
    }

    then:
    def metrics = collector.drain()
    metrics.size() == 1

    and:
    def m = metrics[0]
    m.namespace == 'mlobs'
    m.metricName == 'span.finished'
    m.type == 'count'
    m.value == 1
    m.tags.contains('integration:llmobs')
    m.tags.contains('span_kind:llm')
    m.tags.contains('autoinstrumented:0')
    m.tags.contains('is_root_span:0')
  }

  def "span has expected session tag and telemetry has #expectedHasSessionIdTag"() {
    setup:
    LLMObsMetricCollector collector = LLMObsMetricCollector.get()
    collector.drain()

    when:
    llmObsSpan(Tags.LLMOBS_WORKFLOW_SPAN_KIND, "workflow-span", sessionId).finish()

    then:
    def metrics = collector.drain()
    metrics.size() == 1

    and:
    def m = metrics[0]
    m.namespace == 'mlobs'
    m.metricName == 'span.finished'
    m.tags.contains(expectedHasSessionIdTag)

    where:
    sessionId     | expectedHasSessionIdTag
    "session-123" | "has_session_id:1"
    null          | "has_session_id:0"
  }

  def "child LLMObs span inherits session_id from parent context when none is passed"() {
    setup:
    def expectedSessionId = "session-abc-123"
    def parent = llmObsSpan(Tags.LLMOBS_WORKFLOW_SPAN_KIND, "parent-workflow", expectedSessionId)
    // Activate the parent's AgentScope so the child span is created in the same trace.
    // Without this, the child gets a fresh trace_id and the trace-consistency gate in
    // DDLLMObsSpan would (correctly) skip session_id inheritance.
    def parentScope = AgentTracer.activateSpan((AgentSpan) parent.span)

    when:
    // Child created with null sessionId — should inherit from the parent's LLMObsContext.
    def child = llmObsSpan(Tags.LLMOBS_LLM_SPAN_KIND, "child-llm", null)

    then:
    def innerChild = (AgentSpan) child.span
    expectedSessionId == innerChild.getTag(LLMOBS_TAG_PREFIX + LLMObsTags.SESSION_ID)

    cleanup:
    child.finish()
    parentScope.close()
    parent.finish()
  }

  def "child LLMObs span has no session_id when neither parent nor child passes one"() {
    setup:
    def parent = llmObsSpan(Tags.LLMOBS_WORKFLOW_SPAN_KIND, "parent-workflow", null)

    when:
    def child = llmObsSpan(Tags.LLMOBS_LLM_SPAN_KIND, "child-llm", null)

    then:
    def innerChild = (AgentSpan) child.span
    null == innerChild.getTag(LLMOBS_TAG_PREFIX + LLMObsTags.SESSION_ID)

    cleanup:
    child.finish()
    parent.finish()
  }

  def "grandchild LLMObs span transitively inherits session_id through intermediate span"() {
    setup:
    def expectedSessionId = "session-grandparent-xyz"
    def grandparent = llmObsSpan(Tags.LLMOBS_WORKFLOW_SPAN_KIND, "grandparent-workflow", expectedSessionId)
    // Activate each ancestor's AgentScope so descendants stay in the same trace —
    // session_id inheritance is gated on trace-id consistency in DDLLMObsSpan.
    def grandparentScope = AgentTracer.activateSpan((AgentSpan) grandparent.span)
    def parent = llmObsSpan(Tags.LLMOBS_WORKFLOW_SPAN_KIND, "parent-workflow", null)
    def parentScope = AgentTracer.activateSpan((AgentSpan) parent.span)

    when:
    // Grandchild created with null sessionId — should inherit transitively
    // through parent's re-attached LLMObsContext (which itself inherited from grandparent).
    def grandchild = llmObsSpan(Tags.LLMOBS_LLM_SPAN_KIND, "grandchild-llm", null)

    then:
    def innerGrandchild = (AgentSpan) grandchild.span
    expectedSessionId == innerGrandchild.getTag(LLMOBS_TAG_PREFIX + LLMObsTags.SESSION_ID)

    cleanup:
    grandchild.finish()
    parentScope.close()
    parent.finish()
    grandparentScope.close()
    grandparent.finish()
  }

  def "child does NOT inherit session_id when stale LLMObsContext is from a different trace (e.g. async boundary leak)"() {
    setup:
    // Simulates a stale LLMObsContext (e.g. leaked across an async boundary). The parent's
    // LLMObsContext is attached, but its AgentScope is deliberately NOT activated — so the
    // next span we create starts a fresh trace and the trace-consistency gate must skip
    // session_id inheritance.
    def parent = llmObsSpan(Tags.LLMOBS_WORKFLOW_SPAN_KIND, "stale-workflow", "stale-session-id")

    when:
    def child = llmObsSpan(Tags.LLMOBS_LLM_SPAN_KIND, "child-llm", null)

    then:
    def innerParent = (AgentSpan) parent.span
    def innerChild = (AgentSpan) child.span
    // Sanity: traces differ — confirms the scenario is set up correctly.
    innerParent.getTraceId() != innerChild.getTraceId()
    // Session_id from the stale-trace context must NOT leak into the new span.
    null == innerChild.getTag(LLMOBS_TAG_PREFIX + LLMObsTags.SESSION_ID)

    cleanup:
    child.finish()
    parent.finish()
  }

  def "global dd_tags are included in LLMObs span tags"() {
    setup:
    injectSysConfig("trace.global.tags", "team:backend,owner:ml-platform")
    def test = llmObsSpan(Tags.LLMOBS_WORKFLOW_SPAN_KIND, "test-span")

    expect:
    def innerSpan = (AgentSpan) test.span
    innerSpan.getTag(LLMOBS_TAG_PREFIX + "team") == "backend"
    innerSpan.getTag(LLMOBS_TAG_PREFIX + "owner") == "ml-platform"
  }

  private LLMObsSpan llmObsSpan(String kind, name) {
    llmObsSpan(kind, name, null)
  }

  private LLMObsSpan llmObsSpan(String kind, name, String sessionId) {
    def tags = new WellKnownTags("test-runtime-1", "host-1", "test-env", "test-svc", "v1", "java")
    new DDLLMObsSpan(kind, name, "test-ml-app", sessionId, "test-svc", tags)
  }
}