diff --git a/dd-java-agent/agent-aiguard/src/main/java/com/datadog/aiguard/AIGuardInternal.java b/dd-java-agent/agent-aiguard/src/main/java/com/datadog/aiguard/AIGuardInternal.java index a4e98014ce6..a471f5444d4 100644 --- a/dd-java-agent/agent-aiguard/src/main/java/com/datadog/aiguard/AIGuardInternal.java +++ b/dd-java-agent/agent-aiguard/src/main/java/com/datadog/aiguard/AIGuardInternal.java @@ -12,6 +12,7 @@ import com.squareup.moshi.Types; import datadog.communication.http.OkHttpUtils; import datadog.trace.api.Config; +import datadog.trace.api.DDTags; import datadog.trace.api.aiguard.AIGuard; import datadog.trace.api.aiguard.AIGuard.AIGuardAbortError; import datadog.trace.api.aiguard.AIGuard.AIGuardClientError; @@ -218,6 +219,10 @@ public Evaluation evaluate(final List messages, final Options options) builder.asChildOf(parent.context()); } final AgentSpan span = builder.start(); + final AgentSpan localRootSpan = span.getLocalRootSpan(); + if (localRootSpan != null) { + localRootSpan.setTag(DDTags.MANUAL_KEEP, true); + } try (final AgentScope scope = tracer.activateSpan(span)) { final Message last = messages.get(messages.size() - 1); if (isToolCall(last)) { diff --git a/dd-java-agent/agent-aiguard/src/test/groovy/com/datadog/aiguard/AIGuardInternalTests.groovy b/dd-java-agent/agent-aiguard/src/test/groovy/com/datadog/aiguard/AIGuardInternalTests.groovy index e3a5ffd6671..b5a2e7e9da2 100644 --- a/dd-java-agent/agent-aiguard/src/test/groovy/com/datadog/aiguard/AIGuardInternalTests.groovy +++ b/dd-java-agent/agent-aiguard/src/test/groovy/com/datadog/aiguard/AIGuardInternalTests.groovy @@ -6,6 +6,7 @@ import com.fasterxml.jackson.databind.PropertyNamingStrategies import com.squareup.moshi.Moshi import datadog.common.version.VersionInfo import datadog.trace.api.Config +import datadog.trace.api.DDTags import datadog.trace.api.aiguard.AIGuard import datadog.trace.api.telemetry.WafMetricCollector import datadog.trace.bootstrap.instrumentation.api.AgentSpan @@ -72,12 +73,15 @@ class AIGuardInternalTests extends DDSpecification { protected static final PROMPT = TOOL_OUTPUT + [AIGuard.Message.message('assistant', '2 + 2 is 5'), AIGuard.Message.message('user', '')] protected AgentSpan span + protected AgentSpan localRootSpan void setup() { injectEnvConfig('SERVICE', 'ai_guard_test') injectEnvConfig('ENV', 'test') span = Mock(AgentSpan) + localRootSpan = Mock(AgentSpan) + span.getLocalRootSpan() >> localRootSpan final builder = Mock(AgentTracer.SpanBuilder) { start() >> span } @@ -185,6 +189,7 @@ class AIGuardInternalTests extends DDSpecification { then: 1 * span.setTag(AIGuardInternal.TARGET_TAG, suite.target) + 1 * localRootSpan.setTag(DDTags.MANUAL_KEEP, true) if (suite.target == 'tool') { 1 * span.setTag(AIGuardInternal.TOOL_TAG, 'calc') } diff --git a/dd-smoke-tests/appsec/springboot/src/test/groovy/datadog/smoketest/appsec/AIGuardSmokeTest.groovy b/dd-smoke-tests/appsec/springboot/src/test/groovy/datadog/smoketest/appsec/AIGuardSmokeTest.groovy index 973c96324ec..6cf582a8dd1 100644 --- a/dd-smoke-tests/appsec/springboot/src/test/groovy/datadog/smoketest/appsec/AIGuardSmokeTest.groovy +++ b/dd-smoke-tests/appsec/springboot/src/test/groovy/datadog/smoketest/appsec/AIGuardSmokeTest.groovy @@ -69,9 +69,15 @@ class AIGuardSmokeTest extends AbstractAppSecServerSmokeTest { ?.find { it.resource == 'ai_guard' } as DecodedSpan + final rootSpan = traces*.spans + ?.flatten() + ?.find { + it.traceId == span.traceId && it.parentId == 0 + } as DecodedSpan assert span.meta.get('ai_guard.action') == action assert span.meta.get('ai_guard.reason') == reason assert span.meta.get('ai_guard.target') == 'prompt' + assert rootSpan.metrics.get('_sampling_priority_v1') == 2 final messages = span.metaStruct.get('ai_guard').messages as List> assert messages.size() == 2 with(messages[0]) {