diff --git a/docs/understand/weblogs/end-to-end_weblog.md b/docs/understand/weblogs/end-to-end_weblog.md index c5fb2a076cc..bbe8b2fc2dc 100644 --- a/docs/understand/weblogs/end-to-end_weblog.md +++ b/docs/understand/weblogs/end-to-end_weblog.md @@ -1197,6 +1197,8 @@ This endpoint triggers AI Guard SDK evaluations using the `evaluate()` method fr - **Content-Type:** `application/json` - **Body:** List of AI Guard `Message` objects to be evaluated - **Header:** `X-AI-Guard-Block` (optional, default: `false`) - Controls whether blocking is enabled +- **Header:** `X-User-Id` (optional) - If present and non-empty, sets `usr.id` on the local root span +- **Header:** `X-Session-Id` (optional) - If present and non-empty, sets `session.id` on the local root span **Request Body Example:** ```json diff --git a/manifests/cpp_httpd.yml b/manifests/cpp_httpd.yml index d49cf2a0941..2b82bb84d63 100644 --- a/manifests/cpp_httpd.yml +++ b/manifests/cpp_httpd.yml @@ -12,6 +12,7 @@ # use `>1.0.0` to indicate that requirement. manifest: tests/ai_guard/test_ai_guard_sdk.py::Test_AIGuardEvent_Tag: missing_feature + tests/ai_guard/test_ai_guard_sdk.py::Test_AnomalyDetectionTags: missing_feature tests/ai_guard/test_ai_guard_sdk.py::Test_ClientIPTagsCollected: missing_feature tests/ai_guard/test_ai_guard_sdk.py::Test_ContentParts: missing_feature tests/ai_guard/test_ai_guard_sdk.py::Test_Evaluation: missing_feature diff --git a/manifests/cpp_nginx.yml b/manifests/cpp_nginx.yml index 08f067b8722..62e9d15d551 100644 --- a/manifests/cpp_nginx.yml +++ b/manifests/cpp_nginx.yml @@ -10,6 +10,7 @@ # use `>1.0.0` to indicate that requirement. manifest: tests/ai_guard/test_ai_guard_sdk.py::Test_AIGuardEvent_Tag: missing_feature + tests/ai_guard/test_ai_guard_sdk.py::Test_AnomalyDetectionTags: missing_feature tests/ai_guard/test_ai_guard_sdk.py::Test_ClientIPTagsCollected: missing_feature tests/ai_guard/test_ai_guard_sdk.py::Test_ContentParts: missing_feature tests/ai_guard/test_ai_guard_sdk.py::Test_Evaluation: missing_feature diff --git a/manifests/dotnet.yml b/manifests/dotnet.yml index 207d6d06f29..dc0ae52b560 100644 --- a/manifests/dotnet.yml +++ b/manifests/dotnet.yml @@ -2,6 +2,7 @@ --- manifest: tests/ai_guard/test_ai_guard_sdk.py::Test_AIGuardEvent_Tag: missing_feature + tests/ai_guard/test_ai_guard_sdk.py::Test_AnomalyDetectionTags: missing_feature tests/ai_guard/test_ai_guard_sdk.py::Test_ClientIPTagsCollected: missing_feature tests/ai_guard/test_ai_guard_sdk.py::Test_ContentParts: missing_feature tests/ai_guard/test_ai_guard_sdk.py::Test_Evaluation: missing_feature diff --git a/manifests/golang.yml b/manifests/golang.yml index 2ca38deab19..19a686fc203 100644 --- a/manifests/golang.yml +++ b/manifests/golang.yml @@ -2,6 +2,7 @@ --- manifest: tests/ai_guard/test_ai_guard_sdk.py::Test_AIGuardEvent_Tag: missing_feature (APPSEC-62216) + tests/ai_guard/test_ai_guard_sdk.py::Test_AnomalyDetectionTags: missing_feature tests/ai_guard/test_ai_guard_sdk.py::Test_ClientIPTagsCollected: missing_feature tests/ai_guard/test_ai_guard_sdk.py::Test_ContentParts: missing_feature tests/ai_guard/test_ai_guard_sdk.py::Test_Evaluation: missing_feature diff --git a/manifests/java.yml b/manifests/java.yml index 97093384bdb..8661a50865d 100644 --- a/manifests/java.yml +++ b/manifests/java.yml @@ -2,6 +2,10 @@ --- manifest: tests/ai_guard/test_ai_guard_sdk.py::Test_AIGuardEvent_Tag: missing_feature (APPSEC-62215) + tests/ai_guard/test_ai_guard_sdk.py::Test_AnomalyDetectionTags: + - weblog_declaration: + "*": irrelevant (just one weblog is enough to test the SDK) + "spring-boot": missing_feature (APPSEC-62449) tests/ai_guard/test_ai_guard_sdk.py::Test_ClientIPTagsCollected: - weblog_declaration: "*": irrelevant (just one weblog is enough to test the SDK) diff --git a/manifests/nodejs.yml b/manifests/nodejs.yml index bdf84f6d7a7..e3bfa200203 100644 --- a/manifests/nodejs.yml +++ b/manifests/nodejs.yml @@ -91,6 +91,10 @@ refs: - &ref_6_0_0 '>=6.0.0-pre' manifest: tests/ai_guard/test_ai_guard_sdk.py::Test_AIGuardEvent_Tag: missing_feature (APPSEC-62217) + tests/ai_guard/test_ai_guard_sdk.py::Test_AnomalyDetectionTags: + - weblog_declaration: + "*": irrelevant (just one weblog is enough to test the SDK) + express4: missing_feature (APPSEC-62452) tests/ai_guard/test_ai_guard_sdk.py::Test_ClientIPTagsCollected: - weblog_declaration: "*": irrelevant (just one weblog is enough to test the SDK) diff --git a/manifests/php.yml b/manifests/php.yml index 408b1cf8855..59d8028707e 100644 --- a/manifests/php.yml +++ b/manifests/php.yml @@ -2,6 +2,7 @@ --- manifest: tests/ai_guard/test_ai_guard_sdk.py::Test_AIGuardEvent_Tag: missing_feature + tests/ai_guard/test_ai_guard_sdk.py::Test_AnomalyDetectionTags: missing_feature tests/ai_guard/test_ai_guard_sdk.py::Test_ClientIPTagsCollected: missing_feature tests/ai_guard/test_ai_guard_sdk.py::Test_ContentParts: missing_feature tests/ai_guard/test_ai_guard_sdk.py::Test_Evaluation: missing_feature diff --git a/manifests/python.yml b/manifests/python.yml index fb15dd4f281..56e2b123964 100644 --- a/manifests/python.yml +++ b/manifests/python.yml @@ -8,6 +8,10 @@ manifest: - weblog_declaration: tornado: v4.4.0-rc2 tests/ai_guard/test_ai_guard_sdk.py::Test_AIGuardEvent_Tag: missing_feature (APPSEC-62216) + tests/ai_guard/test_ai_guard_sdk.py::Test_AnomalyDetectionTags: + - weblog_declaration: + "*": irrelevant (just one weblog is enough to test the SDK) + "flask-poc": missing_feature (APPSEC-62450) tests/ai_guard/test_ai_guard_sdk.py::Test_ClientIPTagsCollected: - weblog_declaration: "*": irrelevant (just one weblog is enough to test the SDK) diff --git a/manifests/ruby.yml b/manifests/ruby.yml index 47506646760..64d483155dd 100644 --- a/manifests/ruby.yml +++ b/manifests/ruby.yml @@ -2,6 +2,16 @@ --- manifest: tests/ai_guard/test_ai_guard_sdk.py::Test_AIGuardEvent_Tag: missing_feature (APPSEC-62218) + tests/ai_guard/test_ai_guard_sdk.py::Test_AnomalyDetectionTags: + - weblog_declaration: + "*": missing_feature (APPSEC-62451) + rails42: irrelevant + rack: irrelevant + sinatra14: irrelevant + sinatra22: irrelevant + sinatra32: irrelevant + sinatra41: irrelevant + uds-sinatra: irrelevant tests/ai_guard/test_ai_guard_sdk.py::Test_ClientIPTagsCollected: - weblog_declaration: "*": missing_feature (APPSEC-62200) diff --git a/tests/ai_guard/test_ai_guard_sdk.py b/tests/ai_guard/test_ai_guard_sdk.py index 7dc74467c15..a3a7e66c8e3 100644 --- a/tests/ai_guard/test_ai_guard_sdk.py +++ b/tests/ai_guard/test_ai_guard_sdk.py @@ -493,6 +493,82 @@ def test_sds_in_response(self): assert _assert_key(location, "path") +@rfc("https://datadoghq.atlassian.net/wiki/x/KIApiQE") +@features.ai_guard +@scenarios.ai_guard +class Test_AnomalyDetectionTags: + """Test that anomaly detection attributes are propagated from the root span into every AI Guard span.""" + + PUBLIC_IP = "5.6.7.9" + USER_ID = "u12345" + SESSION_ID = "s12345" + + def _assert_span(self, root_span: DataDogLibrarySpan): + def validate(span: DataDogLibrarySpan): + if span["resource"] != "ai_guard": + return False + + meta = span["meta"] + + # Tags copied from the root span must be present on every AI Guard span + _assert_key(meta, "ai_guard.http.client_ip") + _assert_key(meta, "ai_guard.network.client.ip") + _assert_key(meta, "ai_guard.http.useragent") + _assert_key(meta, "ai_guard.usr.id", self.USER_ID) + _assert_key(meta, "ai_guard.session.id", self.SESSION_ID) + + # Values must match what is on the root span + root_meta = root_span["meta"] + assert meta["ai_guard.http.client_ip"] == root_meta.get("http.client_ip"), ( + f"ai_guard.http.client_ip mismatch: {meta['ai_guard.http.client_ip']} != {root_meta.get('http.client_ip')}" + ) + assert meta["ai_guard.network.client.ip"] == root_meta.get("network.client.ip"), ( + f"ai_guard.network.client.ip mismatch: {meta['ai_guard.network.client.ip']} != {root_meta.get('network.client.ip')}" + ) + assert meta["ai_guard.http.useragent"] == root_meta.get("http.useragent"), ( + f"ai_guard.http.useragent mismatch: {meta['ai_guard.http.useragent']} != {root_meta.get('http.useragent')}" + ) + assert meta["ai_guard.usr.id"] == root_meta.get("usr.id"), ( + f"ai_guard.usr.id mismatch: {meta['ai_guard.usr.id']} != {root_meta.get('usr.id')}" + ) + assert meta["ai_guard.session.id"] == root_meta.get("session.id"), ( + f"ai_guard.session.id mismatch: {meta['ai_guard.session.id']} != {root_meta.get('session.id')}" + ) + + return True + + return validate + + def setup_anomaly_detection_tags(self): + self.r = weblog.post( + "/ai_guard/evaluate", + headers={ + "X-Forwarded-For": self.PUBLIC_IP, + "X-User-Id": self.USER_ID, + "X-Session-Id": self.SESSION_ID, + }, + json=MESSAGES["ALLOW"], + ) + + def test_anomaly_detection_tags(self): + """Test that AI Guard spans carry anomaly detection attributes copied from the root span. + + Verifies that http.client_ip, network.client.ip, http.useragent, usr.id and session.id + are all present on the AI Guard span with the ai_guard. prefix, and that their values + match the corresponding tags on the local root span. + """ + assert self.r.status_code == 200 + + root_span = interfaces.library.get_root_span(self.r) + assert root_span, "No root span found" + + interfaces.library.validate_one_span( + self.r, + validator=self._assert_span(root_span=root_span), + full_trace=True, + ) + + @features.ai_guard @scenarios.ai_guard class Test_AIGuardEvent_Tag: diff --git a/utils/build/docker/java/spring-boot/src/main/java/com/datadoghq/system_tests/springboot/ai_guard/AIGuardController.java b/utils/build/docker/java/spring-boot/src/main/java/com/datadoghq/system_tests/springboot/ai_guard/AIGuardController.java index 0a5ac98ae35..6f137a92471 100644 --- a/utils/build/docker/java/spring-boot/src/main/java/com/datadoghq/system_tests/springboot/ai_guard/AIGuardController.java +++ b/utils/build/docker/java/spring-boot/src/main/java/com/datadoghq/system_tests/springboot/ai_guard/AIGuardController.java @@ -5,6 +5,9 @@ import com.fasterxml.jackson.databind.JsonNode; import datadog.trace.api.aiguard.AIGuard; import datadog.trace.api.aiguard.AIGuard.Evaluation; +import datadog.trace.api.interceptor.MutableSpan; +import io.opentracing.Span; +import io.opentracing.util.GlobalTracer; import org.springframework.boot.autoconfigure.jackson.Jackson2ObjectMapperBuilderCustomizer; import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Configuration; @@ -39,7 +42,19 @@ public Jackson2ObjectMapperBuilderCustomizer mixInCustomizer() { @PostMapping("/ai_guard/evaluate") public ResponseEntity evaluate( @RequestHeader(name = "X-AI-Guard-Block", defaultValue = "false") final boolean block, + @RequestHeader(name = "X-User-Id", required = false) final String userId, + @RequestHeader(name = "X-Session-Id", required = false) final String sessionId, @RequestBody final List data) { + final Span activeSpan = GlobalTracer.get().activeSpan(); + if (activeSpan instanceof MutableSpan) { + final MutableSpan rootSpan = ((MutableSpan) activeSpan).getLocalRootSpan(); + if (userId != null && !userId.isEmpty()) { + rootSpan.setTag("usr.id", userId); + } + if (sessionId != null && !sessionId.isEmpty()) { + rootSpan.setTag("session.id", sessionId); + } + } try { final List messages = data.stream().map(Message::toAIGuard).collect(Collectors.toList()); final Evaluation result = AIGuard.evaluate(messages, new AIGuard.Options().block(block)); diff --git a/utils/build/docker/nodejs/express/app.js b/utils/build/docker/nodejs/express/app.js index 17163269c8a..d3ee6a78d58 100644 --- a/utils/build/docker/nodejs/express/app.js +++ b/utils/build/docker/nodejs/express/app.js @@ -759,6 +759,17 @@ app.post('/ai_guard/evaluate', async (req, res) => { const renameAttrs = ({ tagProbabilities: tag_probs, ...rest }) => ({ ...rest, tag_probs }) const block = req.headers['x-ai-guard-block'] === 'true' const messages = req.body + const userId = req.headers['x-user-id'] + const sessionId = req.headers['x-session-id'] + const rootSpan = tracer.scope().active()?.context()._trace.started[0] + if (rootSpan) { + if (userId) { + rootSpan.setTag('usr.id', userId) + } + if (sessionId) { + rootSpan.setTag('session.id', sessionId) + } + } try { const evaluation = await tracer.aiguard.evaluate(messages, { block }) res.status(200).json(renameAttrs(evaluation)) diff --git a/utils/build/docker/python/flask/app.py b/utils/build/docker/python/flask/app.py index 7744830b9d0..bfcb3ac138f 100644 --- a/utils/build/docker/python/flask/app.py +++ b/utils/build/docker/python/flask/app.py @@ -2187,6 +2187,15 @@ def ai_guard_evaluate(): should_block = flask_request.headers.get("X-AI-Guard-Block", "false").lower() == "true" messages = flask_request.get_json() + user_id = flask_request.headers.get("X-User-Id") + session_id = flask_request.headers.get("X-Session-Id") + root_span = tracer.current_root_span() + if root_span: + if user_id: + root_span.set_tag("usr.id", user_id) + if session_id: + root_span.set_tag("session.id", session_id) + client = new_ai_guard_client(endpoint=os.environ.get("DD_AI_GUARD_ENDPOINT")) evaluation = client.evaluate(messages, Options(block=should_block)) return jsonify(evaluation), 200 diff --git a/utils/build/docker/ruby/rails52/app/controllers/ai_guard_controller.rb b/utils/build/docker/ruby/rails52/app/controllers/ai_guard_controller.rb index 042a529d4a3..08e7641f7eb 100644 --- a/utils/build/docker/ruby/rails52/app/controllers/ai_guard_controller.rb +++ b/utils/build/docker/ruby/rails52/app/controllers/ai_guard_controller.rb @@ -19,9 +19,9 @@ def evaluate Datadog::AIGuard.message(role: message_data[:role]) do |m| message_data[:content].each do |part| case part[:type] - when "text" + when 'text' m.text(part[:text]) - when "image_url" + when 'image_url' m.image_url(part.dig(:image_url, :url)) end end @@ -31,7 +31,15 @@ def evaluate end end - allow_raise = request.headers['X-AI-Guard-Block']&.downcase == "true" + trace = Datadog::Tracing.active_trace + if trace + user_id = request.headers['X-User-Id'] + session_id = request.headers['X-Session-Id'] + trace.set_tag('usr.id', user_id) if user_id.present? + trace.set_tag('session.id', session_id) if session_id.present? + end + + allow_raise = request.headers['X-AI-Guard-Block']&.downcase == 'true' result = Datadog::AIGuard.evaluate(*messages, allow_raise: allow_raise) response_data = { @@ -48,7 +56,7 @@ def evaluate error_data[:tag_probabilities] = e.tag_probabilities if e.respond_to?(:tag_probabilities) error_data[:sds_findings] = e.sds_findings if e.respond_to?(:sds_findings) render json: error_data, status: 403 - rescue => e - render json: {error: e.to_s, type: e.class.name}, status: 500 + rescue StandardError => e + render json: { error: e.to_s, type: e.class.name }, status: 500 end end diff --git a/utils/build/docker/ruby/rails61/app/controllers/ai_guard_controller.rb b/utils/build/docker/ruby/rails61/app/controllers/ai_guard_controller.rb index 042a529d4a3..08e7641f7eb 100644 --- a/utils/build/docker/ruby/rails61/app/controllers/ai_guard_controller.rb +++ b/utils/build/docker/ruby/rails61/app/controllers/ai_guard_controller.rb @@ -19,9 +19,9 @@ def evaluate Datadog::AIGuard.message(role: message_data[:role]) do |m| message_data[:content].each do |part| case part[:type] - when "text" + when 'text' m.text(part[:text]) - when "image_url" + when 'image_url' m.image_url(part.dig(:image_url, :url)) end end @@ -31,7 +31,15 @@ def evaluate end end - allow_raise = request.headers['X-AI-Guard-Block']&.downcase == "true" + trace = Datadog::Tracing.active_trace + if trace + user_id = request.headers['X-User-Id'] + session_id = request.headers['X-Session-Id'] + trace.set_tag('usr.id', user_id) if user_id.present? + trace.set_tag('session.id', session_id) if session_id.present? + end + + allow_raise = request.headers['X-AI-Guard-Block']&.downcase == 'true' result = Datadog::AIGuard.evaluate(*messages, allow_raise: allow_raise) response_data = { @@ -48,7 +56,7 @@ def evaluate error_data[:tag_probabilities] = e.tag_probabilities if e.respond_to?(:tag_probabilities) error_data[:sds_findings] = e.sds_findings if e.respond_to?(:sds_findings) render json: error_data, status: 403 - rescue => e - render json: {error: e.to_s, type: e.class.name}, status: 500 + rescue StandardError => e + render json: { error: e.to_s, type: e.class.name }, status: 500 end end diff --git a/utils/build/docker/ruby/rails72/app/controllers/ai_guard_controller.rb b/utils/build/docker/ruby/rails72/app/controllers/ai_guard_controller.rb index 042a529d4a3..08e7641f7eb 100644 --- a/utils/build/docker/ruby/rails72/app/controllers/ai_guard_controller.rb +++ b/utils/build/docker/ruby/rails72/app/controllers/ai_guard_controller.rb @@ -19,9 +19,9 @@ def evaluate Datadog::AIGuard.message(role: message_data[:role]) do |m| message_data[:content].each do |part| case part[:type] - when "text" + when 'text' m.text(part[:text]) - when "image_url" + when 'image_url' m.image_url(part.dig(:image_url, :url)) end end @@ -31,7 +31,15 @@ def evaluate end end - allow_raise = request.headers['X-AI-Guard-Block']&.downcase == "true" + trace = Datadog::Tracing.active_trace + if trace + user_id = request.headers['X-User-Id'] + session_id = request.headers['X-Session-Id'] + trace.set_tag('usr.id', user_id) if user_id.present? + trace.set_tag('session.id', session_id) if session_id.present? + end + + allow_raise = request.headers['X-AI-Guard-Block']&.downcase == 'true' result = Datadog::AIGuard.evaluate(*messages, allow_raise: allow_raise) response_data = { @@ -48,7 +56,7 @@ def evaluate error_data[:tag_probabilities] = e.tag_probabilities if e.respond_to?(:tag_probabilities) error_data[:sds_findings] = e.sds_findings if e.respond_to?(:sds_findings) render json: error_data, status: 403 - rescue => e - render json: {error: e.to_s, type: e.class.name}, status: 500 + rescue StandardError => e + render json: { error: e.to_s, type: e.class.name }, status: 500 end end diff --git a/utils/build/docker/ruby/rails80/app/controllers/ai_guard_controller.rb b/utils/build/docker/ruby/rails80/app/controllers/ai_guard_controller.rb index c56b35ea3fb..08e7641f7eb 100644 --- a/utils/build/docker/ruby/rails80/app/controllers/ai_guard_controller.rb +++ b/utils/build/docker/ruby/rails80/app/controllers/ai_guard_controller.rb @@ -1,4 +1,3 @@ - # frozen_string_literal: true class AiGuardController < ApplicationController @@ -20,9 +19,9 @@ def evaluate Datadog::AIGuard.message(role: message_data[:role]) do |m| message_data[:content].each do |part| case part[:type] - when "text" + when 'text' m.text(part[:text]) - when "image_url" + when 'image_url' m.image_url(part.dig(:image_url, :url)) end end @@ -32,7 +31,15 @@ def evaluate end end - allow_raise = request.headers['X-AI-Guard-Block']&.downcase == "true" + trace = Datadog::Tracing.active_trace + if trace + user_id = request.headers['X-User-Id'] + session_id = request.headers['X-Session-Id'] + trace.set_tag('usr.id', user_id) if user_id.present? + trace.set_tag('session.id', session_id) if session_id.present? + end + + allow_raise = request.headers['X-AI-Guard-Block']&.downcase == 'true' result = Datadog::AIGuard.evaluate(*messages, allow_raise: allow_raise) response_data = { @@ -49,7 +56,7 @@ def evaluate error_data[:tag_probabilities] = e.tag_probabilities if e.respond_to?(:tag_probabilities) error_data[:sds_findings] = e.sds_findings if e.respond_to?(:sds_findings) render json: error_data, status: 403 - rescue => e - render json: {error: e.to_s, type: e.class.name}, status: 500 + rescue StandardError => e + render json: { error: e.to_s, type: e.class.name }, status: 500 end end