fix: preserve additional top-level properties on AirbyteStateMessage

airbyte_cdk/models/airbyte_protocol_serializers.py

@@ -1,5 +1,5 @@
 # Copyright (c) 2024 Airbyte, Inc., all rights reserved.
-from typing import Any, Dict
+from typing import Any, Dict, cast
 
 from serpyco_rs import CustomType, Serializer
 
@@ -28,10 +28,61 @@ def get_json_schema(self) -> Dict[str, Any]:
         return {"type": "object"}
 
 
-def custom_type_resolver(t: type) -> CustomType[AirbyteStateBlob, Dict[str, Any]] | None:
+def _blob_custom_type_resolver(
+    t: type,
+) -> CustomType[AirbyteStateBlob, Dict[str, Any]] | None:
     return AirbyteStateBlobType() if t is AirbyteStateBlob else None
 
 
+class AirbyteStateMessageType(CustomType[AirbyteStateMessage, Dict[str, Any]]):
+    _KNOWN_PROPERTIES = {
+        "type",
+        "stream",
+        "global",
+        "data",
+        "sourceStats",
+        "destinationStats",
+    }
+
+    def __init__(self) -> None:
+        self._inner = Serializer(
+            AirbyteStateMessage,
+            omit_none=True,
+            custom_type_resolver=_blob_custom_type_resolver,
+        )
+
+    def serialize(self, value: AirbyteStateMessage) -> Dict[str, Any]:
+        result = cast(Dict[str, Any], self._inner.dump(value))
+        result.update(
+            {
+                key: property_value
+                for key, property_value in getattr(value, "additional_properties", {}).items()
+                if key not in self._KNOWN_PROPERTIES
+            }
+        )
+        return result
+
+    def deserialize(self, value: Dict[str, Any]) -> AirbyteStateMessage:
+        message = self._inner.load(value)
+        additional_properties = {
+            key: property_value
+            for key, property_value in value.items()
+            if key not in self._KNOWN_PROPERTIES
+        }
+        if additional_properties:
+            message.additional_properties = additional_properties  # type: ignore[attr-defined]
+        return message
+
+    def get_json_schema(self) -> Dict[str, Any]:
+        return self._inner.get_json_schema()
+
+
+def custom_type_resolver(t: type) -> CustomType[Any, Any] | None:
+    if t is AirbyteStateMessage:
+        return AirbyteStateMessageType()
+    return _blob_custom_type_resolver(t)
+
+
 AirbyteCatalogSerializer = Serializer(AirbyteCatalog, omit_none=True)
 AirbyteStreamSerializer = Serializer(AirbyteStream, omit_none=True)
 AirbyteStreamStateSerializer = Serializer(

unit_tests/test_airbyte_protocol_serializers.py

@@ -0,0 +1,74 @@
+# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
+
+from airbyte_cdk.models import AirbyteMessageSerializer, AirbyteStateMessageSerializer
+
+
+def test_stream_state_message_preserves_additional_properties():
+    message = {
+        "type": "STATE",
+        "state": {
+            "type": "STREAM",
+            "stream": {
+                "stream_descriptor": {"name": "users", "namespace": "public"},
+                "stream_state": {"cursor": "value"},
+            },
+            "id": 42,
+        },
+    }
+
+    loaded = AirbyteMessageSerializer.load(message)
+
+    assert AirbyteMessageSerializer.dump(loaded) == message
+
+
+def test_global_state_message_preserves_additional_properties():
+    message = {
+        "type": "STATE",
+        "state": {
+            "type": "GLOBAL",
+            "global": {
+                "stream_states": [
+                    {
+                        "stream_descriptor": {"name": "users", "namespace": "public"},
+                        "stream_state": {"cursor": "value"},
+                    }
+                ],
+                "shared_state": {"shared": "value"},
+            },
+            "id": 42,
+        },
+    }
+
+    loaded = AirbyteMessageSerializer.load(message)
+
+    assert AirbyteMessageSerializer.dump(loaded) == message
+
+
+def test_direct_state_message_serializer_preserves_additional_properties():
+    message = {
+        "type": "STREAM",
+        "stream": {
+            "stream_descriptor": {"name": "users"},
+            "stream_state": {"cursor": "value"},
+        },
+        "id": 42,
+    }
+
+    loaded = AirbyteStateMessageSerializer.load(message)
+
+    assert AirbyteStateMessageSerializer.dump(loaded) == message
+
+
+def test_state_message_without_additional_properties_does_not_leak_attribute():
+    message = {
+        "type": "STREAM",
+        "stream": {
+            "stream_descriptor": {"name": "users"},
+            "stream_state": {"cursor": "value"},
+        },
+    }
+
+    loaded = AirbyteStateMessageSerializer.load(message)
+
+    assert AirbyteStateMessageSerializer.dump(loaded) == message
+    assert "additional_properties" not in AirbyteStateMessageSerializer.dump(loaded)