feat(schemas): add feature state metadata envelope

src/flagsmith_schemas/api.py

@@ -12,6 +12,8 @@
 When updating this module, ensure that the changes are backwards compatible.
 """
 
+from typing import Any
+
 from flag_engine.engine import ContextValue
 from flag_engine.segments.types import ConditionOperator, RuleType
 from typing_extensions import NotRequired, TypedDict
@@ -57,6 +59,24 @@ class FeatureSegment(TypedDict):
     """The priority of this segment feature override. Lower numbers indicate stronger priority. If null or not set, the weakest priority is assumed."""
 
 
+class ExperimentMetadata(TypedDict):
+    """Represents the experiment a feature state is part of. Only present while the experiment is running."""
+
+    id: int
+    """Unique identifier for the experiment in Core."""
+    name: str
+    """Name of the experiment."""
+    in_experiment: bool
+    """Whether this evaluation counts towards the experiment, i.e. whether the identity was bucketed by the experiment itself."""
+
+
+class FeatureStateMetadata(TypedDict, extra_items=Any, total=False):  # type: ignore[call-arg]  # TODO https://github.com/python/mypy/issues/18176
+    """Additional, non-evaluation data about a feature state."""
+
+    experiment: ExperimentMetadata
+    """The experiment this feature state is part of, if any."""
+
+
 class FeatureState(TypedDict):
     """Used to define the state of a feature for an environment, segment overrides, and identity overrides."""
 
@@ -72,6 +92,8 @@ class FeatureState(TypedDict):
     """Segment override data, if this feature state is for a segment override."""
     multivariate_feature_state_values: list[MultivariateFeatureStateValue]
     """List of multivariate feature state values, if this feature state is for a multivariate feature."""
+    metadata: NotRequired[FeatureStateMetadata]
+    """Additional, non-evaluation data about this feature state. Absent if there is none."""
 
 
 class Trait(TypedDict):
@@ -161,6 +183,8 @@ class V1Flag(TypedDict):
     """Variant key for the feature state."""
     reason: NotRequired[str | None]
     """Why and how this feature state is resolved."""
+    metadata: NotRequired[FeatureStateMetadata]
+    """Additional, non-evaluation data about this feature state. Absent if there is none."""
 
 
 ### Root request schemas below. ###

src/flagsmith_schemas/dynamodb.py

@@ -7,7 +7,7 @@
 Otherwise, they serve as documentation for the structure of the data stored in DynamoDB.
 """
 
-from typing import Annotated, Literal
+from typing import Annotated, Any, Literal
 
 from flag_engine.segments.types import ConditionOperator, RuleType
 from typing_extensions import NotRequired, TypedDict
@@ -95,6 +95,8 @@ class FeatureState(TypedDict):
 
     Total `percentage_allocation` sum of the child multivariate feature state values must be less or equal to 100.
     """
+    metadata: NotRequired[dict[str, Any]]
+    """Additional, non-evaluation data about this feature state, written by Core and served as-is by the SDK API."""
 
 
 class Trait(TypedDict):

tests/integration/flagsmith_schemas/conftest.py

@@ -0,0 +1,15 @@
+from typing import Any
+
+import pytest
+
+
+@pytest.fixture()
+def feature_state_metadata() -> dict[str, Any]:
+    return {
+        "experiment": {
+            "id": 42,
+            "name": "New checkout CTA",
+            "in_experiment": True,
+        },
+        "future_key": {"nested": ["anything"]},
+    }

tests/integration/flagsmith_schemas/test_api.py

@@ -1,6 +1,9 @@
-from pydantic import TypeAdapter
+from typing import Any
 
-from flagsmith_schemas.api import FeatureState
+import pytest
+from pydantic import TypeAdapter, ValidationError
+
+from flagsmith_schemas.api import FeatureState, V1Flag
 
 
 def test_feature_state__featurestate_uuid__expected_json_schema() -> None:
@@ -16,3 +19,88 @@ def test_feature_state__featurestate_uuid__expected_json_schema() -> None:
         "title": "Featurestate Uuid",
         "type": "string",
     }
+
+
+@pytest.mark.parametrize(
+    "schema_type",
+    [FeatureState, V1Flag],
+)
+def test_type_adapter__metadata__preserved_verbatim(
+    schema_type: type[FeatureState] | type[V1Flag],
+    feature_state_metadata: dict[str, Any],
+) -> None:
+    # Given
+    type_adapter = TypeAdapter(schema_type)
+    data = {
+        "feature": {"id": 1, "name": "feature", "type": "STANDARD"},
+        "enabled": True,
+        "feature_state_value": "value",
+        "featurestate_uuid": "652d8931-37d9-438e-9825-f525b9e83077",
+        "feature_segment": None,
+        "multivariate_feature_state_values": [],
+        "metadata": feature_state_metadata,
+    }
+
+    # When
+    document = type_adapter.validate_python(data)
+
+    # Then
+    # Known keys are validated, unknown keys are preserved as-is
+    assert document["metadata"] == feature_state_metadata
+
+
+@pytest.mark.parametrize(
+    "schema_type",
+    [FeatureState, V1Flag],
+)
+def test_type_adapter__no_metadata__key_absent(
+    schema_type: type[FeatureState] | type[V1Flag],
+) -> None:
+    # Given
+    type_adapter = TypeAdapter(schema_type)
+    data = {
+        "feature": {"id": 1, "name": "feature", "type": "STANDARD"},
+        "enabled": True,
+        "feature_state_value": "value",
+        "featurestate_uuid": "652d8931-37d9-438e-9825-f525b9e83077",
+        "feature_segment": None,
+        "multivariate_feature_state_values": [],
+    }
+
+    # When
+    document = type_adapter.validate_python(data)
+
+    # Then
+    assert "metadata" not in document
+
+
+@pytest.mark.parametrize(
+    ("invalid_experiment_field", "invalid_value"),
+    [
+        ("id", "not-an-int"),
+        ("in_experiment", "not-a-bool"),
+    ],
+)
+def test_type_adapter__invalid_experiment_metadata__raises_expected(
+    invalid_experiment_field: str,
+    invalid_value: Any,
+    feature_state_metadata: dict[str, Any],
+) -> None:
+    # Given
+    type_adapter = TypeAdapter(V1Flag)
+    feature_state_metadata["experiment"][invalid_experiment_field] = invalid_value
+    data = {
+        "feature": {"id": 1, "name": "feature", "type": "STANDARD"},
+        "enabled": True,
+        "feature_state_value": "value",
+        "metadata": feature_state_metadata,
+    }
+
+    # When
+    with pytest.raises(ValidationError) as exc_info:
+        type_adapter.validate_python(data)
+
+    # Then
+    assert [error["loc"] for error in exc_info.value.errors()] == [
+        ("metadata", "experiment", invalid_experiment_field)
+    ]

tests/integration/flagsmith_schemas/test_dynamodb.py

@@ -2,7 +2,7 @@
 from decimal import Decimal
 from importlib import reload
 from sys import modules
-from typing import TypeVar
+from typing import Any, TypeVar
 
 import pytest
 import simplejson as json
@@ -16,6 +16,7 @@
     EnvironmentV2IdentityOverride,
     EnvironmentV2Meta,
     EnvironmentV2MetaCompressed,
+    FeatureState,
     Identity,
     MultivariateFeatureOption,
 )
@@ -1083,6 +1084,79 @@ def test_type_adapter__multivariate_feature_option_with_key__key_preserved() ->
     assert document == {"id": Decimal("1"), "value": "control", "key": "control"}
 
 
+@pytest.mark.parametrize(
+    "metadata",
+    [
+        pytest.param(None, id="absent"),
+        pytest.param("feature_state_metadata", id="present"),
+    ],
+)
+def test_type_adapter__feature_state_metadata__preserved_verbatim(
+    metadata: str | None,
+    request: pytest.FixtureRequest,
+) -> None:
+    # Given
+    type_adapter = TypeAdapter(FeatureState)
+    data: dict[str, Any] = {
+        "feature": {"id": 1, "name": "feature", "type": "STANDARD"},
+        "enabled": True,
+        "feature_state_value": "value",
+    }
+    if metadata:
+        data["metadata"] = request.getfixturevalue(metadata)
+
+    # When
+    document = type_adapter.validate_python(data)
+
+    # Then
+    if metadata:
+        assert document["metadata"] == data["metadata"]
+    else:
+        assert "metadata" not in document
+
+
+def test_type_adapter__compressed_environment_feature_state_metadata__metadata_preserved(
+    feature_state_metadata: dict[str, Any],
+) -> None:
+    # Given
+    type_adapter = TypeAdapter(EnvironmentCompressed)
+    python_data = {
+        "id": 1,
+        "api_key": "envkey",
+        "compressed": True,
+        "project": {
+            "id": 1,
+            "name": "Project",
+            "organisation": {
+                "id": 1,
+                "name": "Org",
+                "feature_analytics": False,
+                "stop_serving_flags": False,
+                "persist_trait_data": True,
+            },
+            "segments": [],
+            "hide_disabled_flags": False,
+        },
+        "feature_states": [
+            {
+                "feature": {"id": 1, "name": "feature", "type": "STANDARD"},
+                "enabled": True,
+                "feature_state_value": "value",
+                "django_id": 1,
+                "multivariate_feature_state_values": [],
+                "metadata": feature_state_metadata,
+            }
+        ],
+    }
+
+    # When
+    document = type_adapter.validate_python(python_data)
+
+    # Then
+    feature_states = json.loads(gzip.decompress(bytes(document["feature_states"])))
+    assert feature_states[0]["metadata"] == feature_state_metadata
+
+
 def test_import__no_pydantic__expected_annotations(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None: