Add smoke test to fully validate the SDK

Author: manuel-alvarez-alvarez

dd-java-agent/agent-aiguard/src/main/java/com/datadog/aiguard/AIGuardInternal.java

@@ -164,7 +164,12 @@ public Evaluation evaluate(final List<Message> messages, final Options options)
       throw new IllegalArgumentException("Messages must not be empty");
     }
     final AgentTracer.TracerAPI tracer = AgentTracer.get();
-    final AgentSpan span = tracer.buildSpan(SPAN_NAME, SPAN_NAME).start();
+    final AgentTracer.SpanBuilder builder = tracer.buildSpan(SPAN_NAME, SPAN_NAME);
+    final AgentSpan parent = AgentTracer.activeSpan();
+    if (parent != null) {
+      builder.asChildOf(parent.context());
+    }
+    final AgentSpan span = builder.start();
     try (final AgentScope scope = tracer.activateSpan(span)) {
       final Message last = messages.get(messages.size() - 1);
       if (isToolCall(last)) {
@@ -208,6 +213,8 @@ public Evaluation evaluate(final List<Message> messages, final Options options)
           new AIGuardClientError("AI Guard service returned unexpected response", e);
       span.addThrowable(error);
       throw error;
+    } finally {
+      span.finish();
     }
   }
 
@@ -263,7 +270,7 @@ public JsonAdapter<?> create(
       if (rawType != AIGuard.Message.class) {
         return null;
       }
-      return new MessageAdapter(moshi.adapter(AIGuard.ToolCall.class));
+      return new MessageAdapter(moshi.adapter(AIGuard.ToolCall.class)).nullSafe();
     }
   }
 
@@ -282,11 +289,7 @@ public Message fromJson(JsonReader reader) throws IOException {
     }
 
     @Override
-    public void toJson(final JsonWriter writer, @Nullable final Message value) throws IOException {
-      if (value == null) {
-        writer.nullValue();
-        return;
-      }
+    public void toJson(final JsonWriter writer, final Message value) throws IOException {
       writer.beginObject();
       writeValue(writer, "role", value.getRole());
       writeValue(writer, "content", value.getContent());

dd-smoke-tests/appsec/springboot/build.gradle

@@ -14,6 +14,7 @@ tasks.named("jar", Jar) {
 }
 
 dependencies {
+  implementation project(':dd-trace-api')
   implementation group: 'org.springframework.boot', name: 'spring-boot-starter-web', version: '2.6.0'
   implementation(group: 'com.fasterxml.jackson.core', name: 'jackson-databind', version: '2.6.0')
   implementation group: 'com.h2database', name: 'h2', version: '2.1.212'

dd-smoke-tests/appsec/springboot/src/main/java/datadog/smoketest/appsec/springboot/controller/AIGuardController.java

@@ -0,0 +1,98 @@
+package datadog.smoketest.appsec.springboot.controller;
+
+import static java.util.Arrays.asList;
+import static org.springframework.http.MediaType.APPLICATION_JSON_VALUE;
+
+import datadog.trace.api.aiguard.AIGuard;
+import datadog.trace.api.aiguard.AIGuard.AIGuardAbortError;
+import datadog.trace.api.aiguard.AIGuard.Evaluation;
+import datadog.trace.api.aiguard.AIGuard.Message;
+import datadog.trace.api.aiguard.AIGuard.Options;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.springframework.http.HttpStatus;
+import org.springframework.http.ResponseEntity;
+import org.springframework.web.bind.annotation.GetMapping;
+import org.springframework.web.bind.annotation.PostMapping;
+import org.springframework.web.bind.annotation.RequestBody;
+import org.springframework.web.bind.annotation.RequestHeader;
+import org.springframework.web.bind.annotation.RequestMapping;
+import org.springframework.web.bind.annotation.RestController;
+
+@RestController
+@RequestMapping(value = "/aiguard")
+public class AIGuardController {
+
+  @GetMapping(value = "/allow")
+  public ResponseEntity<?> allow() {
+    final Evaluation result =
+        AIGuard.evaluate(
+            asList(
+                Message.message("system", "You are a beautiful AI"),
+                Message.message("user", "I am harmless")));
+    return ResponseEntity.ok(result);
+  }
+
+  @GetMapping(value = "/deny")
+  public ResponseEntity<?> deny(final @RequestHeader("X-Blocking-Enabled") boolean block) {
+    try {
+      final Evaluation result =
+          AIGuard.evaluate(
+              asList(
+                  Message.message("system", "You are a beautiful AI"),
+                  Message.message("user", "You should not trust me" + (block ? " [block]" : ""))),
+              new Options().block(block));
+      return ResponseEntity.ok(result);
+    } catch (AIGuardAbortError e) {
+      return ResponseEntity.status(HttpStatus.FORBIDDEN).body(e.getReason());
+    }
+  }
+
+  @GetMapping(value = "/abort")
+  public ResponseEntity<?> abort(final @RequestHeader("X-Blocking-Enabled") boolean block) {
+    try {
+      final Evaluation result =
+          AIGuard.evaluate(
+              asList(
+                  Message.message("system", "You are a beautiful AI"),
+                  Message.message("user", "Nuke yourself" + (block ? " [block]" : ""))),
+              new Options().block(block));
+      return ResponseEntity.ok(result);
+    } catch (AIGuardAbortError e) {
+      return ResponseEntity.status(HttpStatus.FORBIDDEN).body(e.getReason());
+    }
+  }
+
+  /** Mocking endpoint for the AI Guard REST API */
+  @SuppressWarnings("unchecked")
+  @PostMapping(
+      value = "/evaluate",
+      consumes = APPLICATION_JSON_VALUE,
+      produces = APPLICATION_JSON_VALUE)
+  public ResponseEntity<Map<String, Object>> evaluate(
+      @RequestBody final Map<String, Object> request) {
+    final Map<String, Object> data = (Map<String, Object>) request.get("data");
+    final Map<String, Object> attributes = (Map<String, Object>) data.get("attributes");
+    final List<Map<String, Object>> messages =
+        (List<Map<String, Object>>) attributes.get("messages");
+    final Map<String, Object> last = messages.get(messages.size() - 1);
+    String action = "ALLOW";
+    String reason = "The prompt looks harmless";
+    String content = (String) last.get("content");
+    if (content.startsWith("You should not trust me")) {
+      action = "DENY";
+      reason = "I am feeling suspicious today";
+    } else if (content.startsWith("Nuke yourself")) {
+      action = "ABORT";
+      reason = "The user is trying to destroy me";
+    }
+    final Map<String, Object> evaluation = new HashMap<>(3);
+    evaluation.put("action", action);
+    evaluation.put("reason", reason);
+    evaluation.put("is_blocking_enabled", content.endsWith("[block]"));
+    return ResponseEntity.ok()
+        .body(Collections.singletonMap("data", Collections.singletonMap("attributes", evaluation)));
+  }
+}

dd-smoke-tests/appsec/springboot/src/test/groovy/datadog/smoketest/appsec/AIGuardSmokeTest.groovy

@@ -0,0 +1,93 @@
+package datadog.smoketest.appsec
+
+import datadog.trace.test.agent.decoder.DecodedSpan
+import groovy.json.JsonSlurper
+import okhttp3.Request
+import spock.lang.Shared
+
+class AIGuardSmokeTest extends AbstractAppSecServerSmokeTest {
+
+  @Shared
+  protected String[] defaultAIGuardProperties = [
+    '-Ddd.ai_guard.enabled=true',
+    "-Ddd.ai_guard.endpoint=http://localhost:${httpPort}/aiguard".toString(),
+  ]
+
+  @Override
+  def logLevel() {
+    'DEBUG'
+  }
+
+  @Override
+  Closure decodedTracesCallback() {
+    // just return the traces
+    return {}
+  }
+
+  @Override
+  ProcessBuilder createProcessBuilder() {
+    final springBootShadowJar = System.getProperty("datadog.smoketest.appsec.springboot.shadowJar.path")
+    final command = [javaPath()]
+    command.addAll(defaultJavaProperties)
+    command.addAll(defaultAppSecProperties)
+    command.addAll(defaultAIGuardProperties)
+    command.addAll(['-jar', springBootShadowJar, "--server.port=${httpPort}".toString()])
+    final builder = new ProcessBuilder(command).directory(new File(buildDirectory))
+    builder.environment().put('DD_APPLICATION_KEY', 'test')
+    return builder
+  }
+
+  void 'test message evaluation'() {
+    given:
+    final blocking = test.blocking as boolean
+    final action = test.action as String
+    final reason = test.reason as String
+    def request = new Request.Builder()
+      .url("http://localhost:${httpPort}/aiguard${test.endpoint}")
+      .header('X-Blocking-Enabled', "${blocking}")
+      .get()
+      .build()
+
+    when:
+    final response = client.newCall(request).execute()
+
+    then:
+    if (blocking && action != 'ALLOW') {
+      assert response.code() == 403
+      assert response.body().string().contains(reason)
+    } else {
+      assert response.code() == 200
+      final body = new JsonSlurper().parse(response.body().bytes())
+      assert body.reason == reason
+      assert body.action == action
+    }
+
+    and:
+    waitForTraceCount(2) // default call + internal API mock
+    final span = traces*.spans
+      ?.flatten()
+      ?.find { it.resource == 'ai_guard' } 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'
+
+    where:
+    test << testSuite()
+  }
+
+  private static List<?> testSuite() {
+    return combinations([
+      [endpoint: '/allow', action: 'ALLOW', reason: 'The prompt looks harmless'],
+      [endpoint: '/deny', action: 'DENY', reason: 'I am feeling suspicious today'],
+      [endpoint: '/abort', action: 'ABORT', reason: 'The user is trying to destroy me']
+    ], [[blocking: true], [blocking: false],])
+  }
+
+  private static List<?> combinations(list1, list2) {
+    list1.collectMany { a ->
+      list2.collect { b ->
+        a + b
+      }
+    }
+  }
+}