Python: fix(sessions): validate AliasChoices state fields when restoring a session snapshot

python/packages/core/agent_framework/_sessions.py

@@ -257,6 +257,27 @@ def encode_pydantic(value: Any) -> Mapping[str, Any]:
     )
 
 
+def _pydantic_validation_key(name: str, field_info: Any) -> str:
+    """Return the name Pydantic accepts for one model field stored under ``name``.
+
+    A string validation alias (which ``alias=`` also installs) or a string model alias is the name
+    Pydantic accepts, so it is used unchanged. ``AliasChoices`` declares several accepted aliases and
+    validates on any one of them, so the first string choice is used; a field keyed by its own name
+    only validates when the model also allows population by name, which declaring an alias does not
+    require. Nested ``AliasPath`` aliases keep the previous field-name fallback.
+    """
+    validation_alias = getattr(field_info, "validation_alias", None)
+    if isinstance(validation_alias, str):
+        return validation_alias
+    choices = getattr(validation_alias, "choices", None)
+    if isinstance(choices, Sequence) and not isinstance(choices, str):
+        for choice in cast("Sequence[Any]", choices):
+            if isinstance(choice, str):
+                return choice
+    alias = getattr(field_info, "alias", None)
+    return alias if isinstance(alias, str) else name
+
+
 def _default_state_decoder(cls: type[Any]) -> StateDecoder:
     """Create the default decoder for one explicitly registered state type."""
     if callable(getattr(cls, "from_dict", None)):
@@ -271,7 +292,13 @@ def decode_from_dict(payload: Mapping[str, Any]) -> Any:
     if issubclass(cls, BaseModel):
 
         def decode_pydantic(payload: Mapping[str, Any]) -> Any:
-            return cls.model_validate({key: value for key, value in payload.items() if key != "type"})
+            # The default encoder dumps field names, so map each one onto a name this model accepts;
+            # an alias-only field rejects its own field name and failed to restore from a snapshot.
+            return cls.model_validate({
+                _pydantic_validation_key(key, cls.model_fields[key]) if key in cls.model_fields else key: value
+                for key, value in payload.items()
+                if key != "type"
+            })
 
         return decode_pydantic
 

python/packages/core/tests/core/test_sessions.py

@@ -910,6 +910,22 @@ class PydanticState(BaseModel):
         assert isinstance(restored.state["value"], PydanticState)
         assert restored.state["value"].value == "ok"
 
+    def test_explicit_pydantic_registration_round_trips_an_alias_choices_field(self) -> None:
+        from pydantic import AliasChoices, BaseModel
+        from pydantic import Field as PydanticField
+
+        class AliasChoiceState(BaseModel):
+            note: str = PydanticField(validation_alias=AliasChoices("note_text", "memo"))
+
+        register_state_type(AliasChoiceState, type_id="alias_choices_state_test")
+        session = AgentSession(session_id="alias-choices")
+        session.state["value"] = AliasChoiceState(note_text="ok")
+
+        restored = AgentSession.from_dict(session.to_dict())
+
+        assert isinstance(restored.state["value"], AliasChoiceState)
+        assert restored.state["value"].note == "ok"
+
     def test_conflicting_type_identifier_is_rejected(self) -> None:
         class FirstState:
             def to_dict(self) -> dict[str, Any]: