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 96051ea0ff8..15b4822f26f 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 @@ -72,6 +72,7 @@ public BadConfigurationException(final String message) { static final String META_STRUCT_TAG = "ai_guard"; static final String META_STRUCT_MESSAGES = "messages"; static final String META_STRUCT_CATEGORIES = "attack_categories"; + static final String META_STRUCT_SDS = "sds"; public static void install() { final Config config = Config.get(); @@ -252,6 +253,8 @@ public Evaluation evaluate(final List messages, final Options options) final String reason = (String) result.get("reason"); @SuppressWarnings("unchecked") final List tags = (List) result.get("tags"); + @SuppressWarnings("unchecked") + final List sdsFindings = (List) result.get("sds_findings"); span.setTag(ACTION_TAG, action); if (reason != null) { span.setTag(REASON_TAG, reason); @@ -259,6 +262,9 @@ public Evaluation evaluate(final List messages, final Options options) if (tags != null && !tags.isEmpty()) { metaStruct.put(META_STRUCT_CATEGORIES, tags); } + if (sdsFindings != null && !sdsFindings.isEmpty()) { + metaStruct.put(META_STRUCT_SDS, sdsFindings); + } final boolean shouldBlock = isBlockingEnabled(options, result.get("is_blocking_enabled")) && action != Action.ALLOW; WafMetricCollector.get().aiGuardRequest(action, shouldBlock); 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 1a9a2b1e9f1..efe62c60b5e 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 @@ -344,6 +344,57 @@ class AIGuardInternalTests extends DDSpecification { messages << [[], null] } + void 'test evaluate with sds findings'() { + 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'] + ], + [ + rule_display_name: 'Social Security Number', + rule_tag: 'ssn', + category: 'pii', + matched_text: '123-45-6789', + location: [start_index: 30, end_index_exclusive: 41, path: 'messages[1].tool_calls[0].function.arguments'] + ] + ] + final aiguard = mockClient(200, [data: [attributes: [action: 'ALLOW', reason: 'It is fine', sds_findings: sdsFindings]]]) + Map receivedMeta + + when: + aiguard.evaluate(PROMPT, AIGuard.Options.DEFAULT) + + then: + 1 * span.setMetaStruct(AIGuardInternal.META_STRUCT_TAG, _) >> { + receivedMeta = it[1] as Map + return span + } + receivedMeta.sds == sdsFindings + } + + void 'test evaluate with empty sds findings'() { + given: + final aiguard = mockClient(200, [data: [attributes: [action: 'ALLOW', reason: 'It is fine', sds_findings: sdsFindings]]]) + Map receivedMeta + + when: + aiguard.evaluate(PROMPT, AIGuard.Options.DEFAULT) + + then: + 1 * span.setMetaStruct(AIGuardInternal.META_STRUCT_TAG, _) >> { + receivedMeta = it[1] as Map + return span + } + !receivedMeta.containsKey('sds') + + where: + sdsFindings << [null, []] + } + void 'test missing tool name'() { given: final aiguard = mockClient(200, [data: [attributes: [action: 'ALLOW', reason: 'Just do it']]])