feat(ffe): support number expected flag type

datadog-ffe/benches/eval.rs

@@ -6,7 +6,8 @@ use serde::{Deserialize, Serialize};
 use std::{collections::HashMap, fs, sync::Arc};
 
 use datadog_ffe::rules_based::{
-    get_assignment, Attribute, Configuration, EvaluationContext, Str, UniversalFlagConfig,
+    get_assignment, Attribute, Configuration, EvaluationContext, ExpectedFlagType, FlagType, Str,
+    UniversalFlagConfig,
 };
 
 fn load_configuration_bytes() -> Vec<u8> {
@@ -17,7 +18,7 @@ fn load_configuration_bytes() -> Vec<u8> {
 #[serde(rename_all = "camelCase")]
 struct TestCase {
     flag: String,
-    variation_type: String,
+    variation_type: FlagType,
     default_value: serde_json::Value,
     targeting_key: Str,
     attributes: HashMap<Str, Attribute>,
@@ -72,7 +73,13 @@ fn bench_sdk_test_data_rules_based(b: &mut Bencher) {
     b.iter(|| {
         for (flag_key, context) in black_box(&test_cases) {
             // Evaluate assignment
-            let _assignment = get_assignment(Some(&configuration), flag_key, context, None, now);
+            let _assignment = get_assignment(
+                Some(&configuration),
+                flag_key,
+                context,
+                ExpectedFlagType::Any,
+                now,
+            );
 
             let _ = black_box(_assignment);
         }
@@ -105,7 +112,7 @@ fn bench_single_flag_rules_based(b: &mut Bencher) {
             black_box(Some(&configuration)),
             black_box("kill-switch"),
             black_box(&context),
-            None,
+            ExpectedFlagType::Any,
             now,
         );
         let _ = black_box(_assignment);

datadog-ffe/src/flag_type.rs

@@ -0,0 +1,52 @@
+// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/
+// SPDX-License-Identifier: Apache-2.0
+
+use serde::{Deserialize, Serialize};
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(rename_all = "snake_case")]
+#[repr(u8)]
+pub enum FlagType {
+    #[serde(alias = "BOOLEAN")]
+    Boolean = 1,
+    #[serde(alias = "STRING")]
+    String = 1 << 1,
+    #[serde(alias = "NUMERIC")]
+    Float = 1 << 2,
+    #[serde(alias = "INTEGER")]
+    Integer = 1 << 3,
+    #[serde(alias = "JSON")]
+    Object = 1 << 4,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(rename_all = "snake_case")]
+#[repr(u8)]
+#[non_exhaustive]
+pub enum ExpectedFlagType {
+    Boolean = FlagType::Boolean as u8,
+    String = FlagType::String as u8,
+    Float = FlagType::Float as u8,
+    Integer = FlagType::Integer as u8,
+    Object = FlagType::Object as u8,
+    Number = (FlagType::Integer as u8) | (FlagType::Float as u8),
+    Any = 0xff,
+}
+
+impl From<FlagType> for ExpectedFlagType {
+    fn from(value: FlagType) -> Self {
+        match value {
+            FlagType::String => ExpectedFlagType::String,
+            FlagType::Integer => ExpectedFlagType::Integer,
+            FlagType::Float => ExpectedFlagType::Float,
+            FlagType::Boolean => ExpectedFlagType::Boolean,
+            FlagType::Object => ExpectedFlagType::Object,
+        }
+    }
+}
+
+impl ExpectedFlagType {
+    pub(crate) fn is_compatible(self, ty: FlagType) -> bool {
+        (self as u8) & (ty as u8) != 0
+    }
+}

datadog-ffe/src/lib.rs

@@ -1,4 +1,8 @@
 // Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/
 // SPDX-License-Identifier: Apache-2.0
 
+mod flag_type;
+
 pub mod rules_based;
+
+pub use flag_type::{ExpectedFlagType, FlagType};

datadog-ffe/src/rules_based/error.rs

@@ -3,7 +3,7 @@
 
 use serde::{Deserialize, Serialize};
 
-use crate::rules_based::ufc::VariationType;
+use crate::rules_based::{ExpectedFlagType, FlagType};
 
 /// Enum representing all possible reasons that could result in evaluation returning an error or
 /// default assignment.
@@ -18,9 +18,9 @@ pub enum EvaluationError {
     #[error("invalid flag type (expected: {expected:?}, found: {found:?})")]
     TypeMismatch {
         /// Expected type of the flag.
-        expected: VariationType,
+        expected: ExpectedFlagType,
         /// Actual type of the flag.
-        found: VariationType,
+        found: FlagType,
     },
 
     /// Failed to parse configuration. This should normally never happen and is likely a signal

datadog-ffe/src/rules_based/eval/eval_assignment.rs

@@ -5,11 +5,8 @@ use chrono::{DateTime, Utc};
 
 use crate::rules_based::{
     error::EvaluationError,
-    ufc::{
-        Allocation, Assignment, AssignmentReason, CompiledFlagsConfig, Flag, Shard, Split,
-        VariationType,
-    },
-    Configuration, EvaluationContext, Timestamp,
+    ufc::{Allocation, Assignment, AssignmentReason, CompiledFlagsConfig, Flag, Shard, Split},
+    Configuration, EvaluationContext, ExpectedFlagType, Timestamp,
 };
 
 /// Evaluate the specified feature flag for the given subject and return assigned variation and
@@ -18,7 +15,7 @@ pub fn get_assignment(
     configuration: Option<&Configuration>,
     flag_key: &str,
     subject: &EvaluationContext,
-    expected_type: Option<VariationType>,
+    expected_type: ExpectedFlagType,
     now: DateTime<Utc>,
 ) -> Result<Assignment, EvaluationError> {
     let Some(config) = configuration else {
@@ -37,7 +34,7 @@ impl Configuration {
         &self,
         flag_key: &str,
         context: &EvaluationContext,
-        expected_type: Option<VariationType>,
+        expected_type: ExpectedFlagType,
         now: DateTime<Utc>,
     ) -> Result<Assignment, EvaluationError> {
         let result = self
@@ -72,16 +69,10 @@ impl CompiledFlagsConfig {
         &self,
         flag_key: &str,
         subject: &EvaluationContext,
-        expected_type: Option<VariationType>,
+        expected_type: ExpectedFlagType,
         now: DateTime<Utc>,
     ) -> Result<Assignment, EvaluationError> {
-        let flag = self.get_flag(flag_key)?;
-
-        if let Some(ty) = expected_type {
-            flag.verify_type(ty)?;
-        }
-
-        flag.eval(subject, now)
+        self.get_flag(flag_key)?.eval(subject, expected_type, now)
     }
 
     fn get_flag(&self, flag_key: &str) -> Result<&Flag, EvaluationError> {
@@ -94,22 +85,19 @@ impl CompiledFlagsConfig {
 }
 
 impl Flag {
-    fn verify_type(&self, ty: VariationType) -> Result<(), EvaluationError> {
-        if self.variation_type == ty {
-            Ok(())
-        } else {
-            Err(EvaluationError::TypeMismatch {
-                expected: ty,
-                found: self.variation_type,
-            })
-        }
-    }
-
     fn eval(
         &self,
         subject: &EvaluationContext,
+        expected_type: ExpectedFlagType,
         now: DateTime<Utc>,
     ) -> Result<Assignment, EvaluationError> {
+        if !expected_type.is_compatible(self.variation_type.into()) {
+            return Err(EvaluationError::TypeMismatch {
+                expected: expected_type,
+                found: self.variation_type.into(),
+            });
+        }
+
         let Some((allocation, (split, reason))) = self.allocations.iter().find_map(|allocation| {
             let result = allocation.get_matching_split(subject, now);
             result
@@ -212,15 +200,15 @@ mod tests {
 
     use crate::rules_based::{
         eval::get_assignment,
-        ufc::{AssignmentValue, UniversalFlagConfig, VariationType},
-        Attribute, Configuration, EvaluationContext, Str,
+        ufc::{AssignmentValue, UniversalFlagConfig},
+        Attribute, Configuration, EvaluationContext, FlagType, Str,
     };
 
     #[derive(Debug, Serialize, Deserialize)]
     #[serde(rename_all = "camelCase")]
     struct TestCase {
         flag: String,
-        variation_type: VariationType,
+        variation_type: FlagType,
         default_value: serde_json::Value,
         targeting_key: Str,
         attributes: Arc<HashMap<Str, Attribute>>,
@@ -251,27 +239,31 @@ mod tests {
             let test_cases: Vec<TestCase> = serde_json::from_reader(f).unwrap();
 
             for test_case in test_cases {
-                let default_assignment =
-                    AssignmentValue::from_wire(test_case.variation_type, test_case.default_value)
-                        .unwrap();
+                let default_assignment = AssignmentValue::from_wire(
+                    test_case.variation_type.into(),
+                    test_case.default_value,
+                )
+                .unwrap();
 
                 print!("test subject {:?} ... ", test_case.targeting_key);
                 let subject = EvaluationContext::new(test_case.targeting_key, test_case.attributes);
                 let result = get_assignment(
                     Some(&config),
                     &test_case.flag,
                     &subject,
-                    Some(test_case.variation_type),
+                    test_case.variation_type.into(),
                     now,
                 );
 
                 let result_assingment = result
                     .as_ref()
                     .map(|assignment| &assignment.value)
                     .unwrap_or(&default_assignment);
-                let expected_assignment =
-                    AssignmentValue::from_wire(test_case.variation_type, test_case.result.value)
-                        .unwrap();
+                let expected_assignment = AssignmentValue::from_wire(
+                    test_case.variation_type.into(),
+                    test_case.result.value,
+                )
+                .unwrap();
 
                 assert_eq!(result_assingment, &expected_assignment);
                 println!("ok");

datadog-ffe/src/rules_based/mod.rs

@@ -16,4 +16,6 @@ pub use error::EvaluationError;
 pub use eval::{get_assignment, EvaluationContext};
 pub use str::Str;
 pub use timestamp::{now, Timestamp};
-pub use ufc::{Assignment, AssignmentReason, AssignmentValue, UniversalFlagConfig, VariationType};
+pub use ufc::{Assignment, AssignmentReason, AssignmentValue, UniversalFlagConfig};
+
+pub use crate::{ExpectedFlagType, FlagType};

datadog-ffe/src/rules_based/ufc/assignment.rs

@@ -6,9 +6,7 @@ use std::sync::Arc;
 
 use serde::{Deserialize, Serialize};
 
-use crate::rules_based::Str;
-
-use super::VariationType;
+use crate::rules_based::{ufc::VariationType, FlagType, Str};
 
 /// Reason for assignment evaluation result.
 #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
@@ -267,18 +265,17 @@ impl AssignmentValue {
     ///
     /// # Examples
     /// ```
-    /// # use datadog_ffe::rules_based::AssignmentValue;
-    /// # use datadog_ffe::rules_based::VariationType;
+    /// # use datadog_ffe::rules_based::{AssignmentValue, FlagType};
     /// let value = AssignmentValue::String("example".into());
-    /// assert_eq!(value.variation_type(), VariationType::String);
+    /// assert_eq!(value.variation_type(), FlagType::String);
     /// ```
-    pub fn variation_type(&self) -> VariationType {
+    pub fn variation_type(&self) -> FlagType {
         match self {
-            AssignmentValue::String(_) => VariationType::String,
-            AssignmentValue::Integer(_) => VariationType::Integer,
-            AssignmentValue::Float(_) => VariationType::Numeric,
-            AssignmentValue::Boolean(_) => VariationType::Boolean,
-            AssignmentValue::Json(_) => VariationType::Json,
+            AssignmentValue::String(_) => FlagType::String,
+            AssignmentValue::Integer(_) => FlagType::Integer,
+            AssignmentValue::Float(_) => FlagType::Float,
+            AssignmentValue::Boolean(_) => FlagType::Boolean,
+            AssignmentValue::Json(_) => FlagType::Object,
         }
     }
 

datadog-ffe/src/rules_based/ufc/models.rs

@@ -6,7 +6,7 @@ use std::{collections::HashMap, sync::Arc};
 use regex::Regex;
 use serde::{Deserialize, Serialize};
 
-use crate::rules_based::{EvaluationError, Str, Timestamp};
+use crate::rules_based::{EvaluationError, FlagType, Str, Timestamp};
 
 /// Universal Flag Configuration attributes. This contains the actual flag configuration data.
 #[derive(Debug, Serialize, Deserialize, Clone)]
@@ -88,6 +88,30 @@ pub enum VariationType {
     Json,
 }
 
+impl From<VariationType> for FlagType {
+    fn from(value: VariationType) -> FlagType {
+        match value {
+            VariationType::String => FlagType::String,
+            VariationType::Integer => FlagType::Integer,
+            VariationType::Numeric => FlagType::Float,
+            VariationType::Boolean => FlagType::Boolean,
+            VariationType::Json => FlagType::Object,
+        }
+    }
+}
+
+impl From<FlagType> for VariationType {
+    fn from(value: FlagType) -> VariationType {
+        match value {
+            FlagType::String => VariationType::String,
+            FlagType::Integer => VariationType::Integer,
+            FlagType::Float => VariationType::Numeric,
+            FlagType::Boolean => VariationType::Boolean,
+            FlagType::Object => VariationType::Json,
+        }
+    }
+}
+
 #[derive(Debug, Serialize, Deserialize, Clone)]
 #[serde(rename_all = "camelCase")]
 #[allow(missing_docs)]