chore: add the tags returned by the service to the ai_guard span (#6892)

Author: manuel-alvarez-alvarez

packages/dd-trace/src/aiguard/sdk.js

@@ -138,10 +138,11 @@ class AIGuard extends NoopAIGuard {
           span.setTag(AI_GUARD_TOOL_NAME_TAG_KEY, name)
         }
       }
+      const metaStruct = {
+        messages: this.#truncate(messages)
+      }
       span.meta_struct = {
-        [AI_GUARD_META_STRUCT_KEY]: {
-          messages: this.#truncate(messages)
-        }
+        [AI_GUARD_META_STRUCT_KEY]: metaStruct
       }
       let response
       try {
@@ -166,14 +167,15 @@ class AIGuard extends NoopAIGuard {
           `AI Guard service call failed, status ${response.status}`,
           { errors: response.body?.errors })
       }
-      let action, reason, blockingEnabled
+      let action, reason, tags, blockingEnabled
       try {
         const attr = response.body.data.attributes
         if (!attr.action) {
           throw new Error('Action missing from response')
         }
         action = attr.action
         reason = attr.reason
+        tags = attr.tags
         blockingEnabled = attr.is_blocking_enabled ?? false
       } catch (e) {
         appsecMetrics.count(AI_GUARD_TELEMETRY_REQUESTS, { error: true }).inc(1)
@@ -182,7 +184,12 @@ class AIGuard extends NoopAIGuard {
       const shouldBlock = block && blockingEnabled && action !== ALLOW
       appsecMetrics.count(AI_GUARD_TELEMETRY_REQUESTS, { action, error: false, block: shouldBlock }).inc(1)
       span.setTag(AI_GUARD_ACTION_TAG_KEY, action)
-      span.setTag(AI_GUARD_REASON_TAG_KEY, reason)
+      if (reason) {
+        span.setTag(AI_GUARD_REASON_TAG_KEY, reason)
+      }
+      if (tags?.length > 0) {
+        metaStruct.attack_categories = tags
+      }
       if (shouldBlock) {
         span.setTag(AI_GUARD_BLOCKED_TAG_KEY, 'true')
         throw new AIGuardAbortError(reason)

packages/dd-trace/test/aiguard/index.spec.js

@@ -138,9 +138,9 @@ describe('AIGuard SDK', () => {
   }
 
   const testSuite = [
-    { action: 'ALLOW', reason: 'Go ahead' },
-    { action: 'DENY', reason: 'Nope' },
-    { action: 'ABORT', reason: 'Kill it with fire' }
+    { action: 'ALLOW', reason: 'Go ahead', tags: [] },
+    { action: 'DENY', reason: 'Nope', tags: ['deny_everything', 'test_deny'] },
+    { action: 'ABORT', reason: 'Kill it with fire', tags: ['alarm_tag', 'abort_everything'] }
   ].flatMap(r => [
     { ...r, blocking: true },
     { ...r, blocking: false },
@@ -150,9 +150,9 @@ describe('AIGuard SDK', () => {
     { ...r, suite: 'prompt', target: 'prompt', messages: prompt }
   ])
 
-  for (const { action, reason, blocking, suite, target, messages } of testSuite) {
+  for (const { action, reason, tags, blocking, suite, target, messages } of testSuite) {
     it(`test evaluate '${suite}' with ${action} action (blocking: ${blocking})`, async () => {
-      mockFetch({ body: { data: { attributes: { action, reason, is_blocking_enabled: blocking } } } })
+      mockFetch({ body: { data: { attributes: { action, reason, tags, is_blocking_enabled: blocking } } } })
       const shouldBlock = action !== 'ALLOW' && blocking
 
       if (shouldBlock) {
@@ -175,7 +175,10 @@ describe('AIGuard SDK', () => {
         ...(target === 'tool' ? { 'ai_guard.tool_name': 'calc' } : {}),
         ...(shouldBlock ? { 'ai_guard.blocked': 'true', 'error.type': 'AIGuardAbortError' } : {})
       },
-      { messages })
+      {
+        messages,
+        ...(tags.length > 0 ? { attack_categories: tags } : {})
+      })
     })
   }