feat(ai): Add support for AI token span data (#3250)

Upstream in SDKs (getsentry/sentry-python#2791)
we added strings like "ai.prompt_tokens.used"

This adds them as known data fields, otherwise they are removed as PII
since most configurations consider any superstring of "token" to be a
password.

CHANGELOG.md

@@ -7,6 +7,7 @@
 - Add support for Reporting API for CSP reports ([#3277](https://github.com/getsentry/relay/pull/3277))
 - Extract op and description while converting opentelemetry spans to sentry spans. ([#3287](https://github.com/getsentry/relay/pull/3287))
 - Drop `event_id` and `remote_addr` from all outcomes. ([#3319](https://github.com/getsentry/relay/pull/3319))
+- Support for AI token metrics ([#3250](https://github.com/getsentry/relay/pull/3250))
 
 **Internal**:
 

relay-event-schema/src/protocol/span.rs

@@ -270,6 +270,26 @@ pub struct SpanData {
     #[metastructure(field = "http.response.status_code", legacy_alias = "status_code")]
     pub http_response_status_code: Annotated<Value>,
 
+    /// The input messages to an AI model call
+    #[metastructure(field = "ai.input_messages")]
+    pub ai_input_messages: Annotated<Value>,
+
+    /// The number of tokens used to generate the response to an AI call
+    #[metastructure(field = "ai.completion_tokens.used", pii = "false")]
+    pub ai_completion_tokens_used: Annotated<Value>,
+
+    /// The number of tokens used to process a request for an AI call
+    #[metastructure(field = "ai.prompt_tokens.used", pii = "false")]
+    pub ai_prompt_tokens_used: Annotated<Value>,
+
+    /// The total number of tokens used to for an AI call
+    #[metastructure(field = "ai.total_tokens.used", pii = "false")]
+    pub ai_total_tokens_used: Annotated<Value>,
+
+    /// The responses to an AI model call
+    #[metastructure(field = "ai.responses")]
+    pub ai_responses: Annotated<Value>,
+
     /// Label identifying a thread from where the span originated.
     #[metastructure(field = "thread.name")]
     pub thread_name: Annotated<Value>,
@@ -606,6 +626,11 @@ mod tests {
             resource_render_blocking_status: ~,
             server_address: ~,
             http_response_status_code: ~,
+            ai_input_messages: ~,
+            ai_completion_tokens_used: ~,
+            ai_prompt_tokens_used: ~,
+            ai_total_tokens_used: ~,
+            ai_responses: ~,
             thread_name: ~,
             transaction: ~,
             ui_component_name: ~,

relay-pii/src/processor.rs

@@ -472,7 +472,7 @@ mod tests {
         Addr, Breadcrumb, DebugImage, DebugMeta, ExtraValue, Headers, LogEntry, Message,
         NativeDebugImage, Request, Span, TagEntry, Tags, TraceContext,
     };
-    use relay_protocol::{assert_annotated_snapshot, get_value, FromValue, Object};
+    use relay_protocol::{assert_annotated_snapshot, get_value, FromValue, Object, Val};
     use serde_json::json;
 
     use super::*;
@@ -1337,6 +1337,43 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_ai_token_values() {
+        let mut span = Span::from_value(
+            json!({
+                "data": {
+                    "ai.total_tokens.used": 30,
+                    "ai.prompt_tokens.used": 20,
+                    "ai.completion_tokens.used": 10,
+                }
+            })
+            .into(),
+        );
+
+        let pii_config = serde_json::from_value::<PiiConfig>(json!({
+            "applications": {
+                "$object": ["@password"],
+            }
+        }))
+        .expect("invalid json config");
+
+        let mut pii_processor = PiiProcessor::new(pii_config.compiled());
+        process_value(&mut span, &mut pii_processor, ProcessingState::root()).unwrap();
+
+        assert_eq!(
+            Val::from(get_value!(span.data.ai_total_tokens_used!)).as_u64(),
+            Some(30),
+        );
+        assert_eq!(
+            Val::from(get_value!(span.data.ai_prompt_tokens_used!)).as_u64(),
+            Some(20),
+        );
+        assert_eq!(
+            Val::from(get_value!(span.data.ai_completion_tokens_used!)).as_u64(),
+            Some(10),
+        );
+    }
+
     #[test]
     fn test_scrub_breadcrumb_data_http_not_scrubbed() {
         let mut breadcrumb: Annotated<Breadcrumb> = Annotated::from_json(

relay-server/src/metrics_extraction/snapshots/relay_server__metrics_extraction__event__tests__extract_span_metrics_mobile.snap

@@ -325,6 +325,11 @@ expression: "(&event.value().unwrap().spans, metrics)"
                 resource_render_blocking_status: ~,
                 server_address: ~,
                 http_response_status_code: ~,
+                ai_input_messages: ~,
+                ai_completion_tokens_used: ~,
+                ai_prompt_tokens_used: ~,
+                ai_total_tokens_used: ~,
+                ai_responses: ~,
                 thread_name: ~,
                 transaction: ~,
                 ui_component_name: ~,
@@ -399,6 +404,11 @@ expression: "(&event.value().unwrap().spans, metrics)"
                 resource_render_blocking_status: ~,
                 server_address: ~,
                 http_response_status_code: ~,
+                ai_input_messages: ~,
+                ai_completion_tokens_used: ~,
+                ai_prompt_tokens_used: ~,
+                ai_total_tokens_used: ~,
+                ai_responses: ~,
                 thread_name: ~,
                 transaction: ~,
                 ui_component_name: ~,