diff --git a/python/packages/core/agent_framework/_sessions.py b/python/packages/core/agent_framework/_sessions.py index 7ea242078a..da1ddd1b80 100644 --- a/python/packages/core/agent_framework/_sessions.py +++ b/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 diff --git a/python/packages/core/tests/core/test_sessions.py b/python/packages/core/tests/core/test_sessions.py index 23eee40810..f345248bd1 100644 --- a/python/packages/core/tests/core/test_sessions.py +++ b/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]: