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 15b4822f26f..1fcde926d9d 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 @@ -270,9 +270,9 @@ public Evaluation evaluate(final List messages, final Options options) WafMetricCollector.get().aiGuardRequest(action, shouldBlock); if (shouldBlock) { span.setTag(BLOCKED_TAG, true); - throw new AIGuardAbortError(action, reason, tags); + throw new AIGuardAbortError(action, reason, tags, sdsFindings); } - return new Evaluation(action, reason, tags); + return new Evaluation(action, reason, tags, sdsFindings); } } catch (AIGuardAbortError e) { span.addThrowable(e); 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 efe62c60b5e..ccacbb2f603 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 @@ -210,11 +210,13 @@ class AIGuardInternalTests extends DDSpecification { error.action == suite.action error.reason == suite.reason error.tags == suite.tags + error.sds == [] } else { error == null eval.action == suite.action eval.reason == suite.reason eval.tags == suite.tags + eval.sds == [] } assertTelemetry('ai_guard.requests', "action:$suite.action", "block:$throwAbortError", 'error:false') @@ -366,7 +368,7 @@ class AIGuardInternalTests extends DDSpecification { Map receivedMeta when: - aiguard.evaluate(PROMPT, AIGuard.Options.DEFAULT) + final result = aiguard.evaluate(PROMPT, AIGuard.Options.DEFAULT) then: 1 * span.setMetaStruct(AIGuardInternal.META_STRUCT_TAG, _) >> { @@ -374,6 +376,7 @@ class AIGuardInternalTests extends DDSpecification { return span } receivedMeta.sds == sdsFindings + result.sds == sdsFindings } void 'test evaluate with empty sds findings'() { @@ -382,7 +385,7 @@ class AIGuardInternalTests extends DDSpecification { Map receivedMeta when: - aiguard.evaluate(PROMPT, AIGuard.Options.DEFAULT) + final result = aiguard.evaluate(PROMPT, AIGuard.Options.DEFAULT) then: 1 * span.setMetaStruct(AIGuardInternal.META_STRUCT_TAG, _) >> { @@ -390,11 +393,33 @@ class AIGuardInternalTests extends DDSpecification { return span } !receivedMeta.containsKey('sds') + result.sds == (sdsFindings ?: []) where: sdsFindings << [null, []] } + void 'test evaluate with sds findings in abort error'() { + given: + final sdsFindings = [ + [ + rule_display_name: 'Credit Card Number', + rule_tag: 'credit_card', + category: 'pii', + matched_text: '4111111111111111', + location: [start_index: 10, end_index_exclusive: 26, path: 'messages[0].content[0].text'] + ] + ] + final aiguard = mockClient(200, [data: [attributes: [action: 'ABORT', reason: 'PII detected', tags: ['pii'], sds_findings: sdsFindings, is_blocking_enabled: true]]]) + + when: + aiguard.evaluate(PROMPT, new AIGuard.Options().block(true)) + + then: + final error = thrown(AIGuard.AIGuardAbortError) + error.sds == sdsFindings + } + void 'test missing tool name'() { given: final aiguard = mockClient(200, [data: [attributes: [action: 'ALLOW', reason: 'Just do it']]]) diff --git a/dd-trace-api/src/main/java/datadog/trace/api/aiguard/AIGuard.java b/dd-trace-api/src/main/java/datadog/trace/api/aiguard/AIGuard.java index 18aa7d4ca9b..6c09ce24b9b 100644 --- a/dd-trace-api/src/main/java/datadog/trace/api/aiguard/AIGuard.java +++ b/dd-trace-api/src/main/java/datadog/trace/api/aiguard/AIGuard.java @@ -69,12 +69,15 @@ public static class AIGuardAbortError extends RuntimeException { private final Action action; private final String reason; private final List tags; + private final List sds; - public AIGuardAbortError(final Action action, final String reason, final List tags) { + public AIGuardAbortError( + final Action action, final String reason, final List tags, final List sds) { super(reason); this.action = action; this.reason = reason; this.tags = tags; + this.sds = sds != null ? sds : Collections.emptyList(); } public Action getAction() { @@ -88,6 +91,10 @@ public String getReason() { public List getTags() { return tags; } + + public List getSds() { + return sds; + } } /** @@ -149,6 +156,7 @@ public static class Evaluation { final Action action; final String reason; final List tags; + final List sds; /** * Creates a new evaluation result. @@ -156,11 +164,14 @@ public static class Evaluation { * @param action the recommended action for the evaluated content * @param reason human-readable explanation for the decision * @param tags list of tags associated with the evaluation (e.g. indirect-prompt-injection) + * @param sds list of Sensitive Data Scanner findings */ - public Evaluation(final Action action, final String reason, final List tags) { + public Evaluation( + final Action action, final String reason, final List tags, final List sds) { this.action = action; this.reason = reason; this.tags = tags; + this.sds = sds != null ? sds : Collections.emptyList(); } /** @@ -189,6 +200,15 @@ public String getReason() { public List getTags() { return tags; } + + /** + * Returns the list of Sensitive Data Scanner findings. + * + * @return list of SDS findings. + */ + public List getSds() { + return sds; + } } /** diff --git a/dd-trace-api/src/main/java/datadog/trace/api/aiguard/noop/NoOpEvaluator.java b/dd-trace-api/src/main/java/datadog/trace/api/aiguard/noop/NoOpEvaluator.java index 44d88c68878..9fa6c5013b3 100644 --- a/dd-trace-api/src/main/java/datadog/trace/api/aiguard/noop/NoOpEvaluator.java +++ b/dd-trace-api/src/main/java/datadog/trace/api/aiguard/noop/NoOpEvaluator.java @@ -13,6 +13,6 @@ public final class NoOpEvaluator implements Evaluator { @Override public Evaluation evaluate(final List messages, final Options options) { - return new Evaluation(ALLOW, "AI Guard is not enabled", emptyList()); + return new Evaluation(ALLOW, "AI Guard is not enabled", emptyList(), emptyList()); } }