diff --git a/python/packages/core/agent_framework/_sessions.py b/python/packages/core/agent_framework/_sessions.py index 7ea242078a..9adac362a0 100644 --- a/python/packages/core/agent_framework/_sessions.py +++ b/python/packages/core/agent_framework/_sessions.py @@ -266,12 +266,52 @@ def decode_from_dict(payload: Mapping[str, Any]) -> Any: return decode_from_dict - from pydantic import BaseModel + from pydantic import AliasChoices, BaseModel, ValidationError 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"}) + raw = {key: value for key, value in payload.items() if key != "type"} + try: + return cls.model_validate(raw) + except ValidationError as exc: + # The encoder's payload items went through _serialize_value, + # so a nested registered value (e.g. Message) sits in the + # payload as its tagged dict. Restore only the field(s) pydantic + # blames; fields that already validate raw keep their payload + # untouched, so a field that legitimately holds a tag-shaped + # dict (e.g. plain dict[str, Any] data) is never converted or + # rejected by the retry. + failed_keys: set[str] = set() + for error in exc.errors(): + location = error.get("loc") + if not location or not isinstance(location[0], str): + continue + field_name = location[0] + if field_name in raw: + failed_keys.add(field_name) + continue + field = cls.model_fields.get(field_name) + if field is None: + continue + validation_alias = field.validation_alias + aliases = [field.alias, field.serialization_alias] + if isinstance(validation_alias, str): + aliases.insert(0, validation_alias) + elif isinstance(validation_alias, AliasChoices): + aliases = [*validation_alias.choices, *aliases] + failed_key = next( + (alias for alias in aliases if isinstance(alias, str) and alias in raw), + None, + ) + if failed_key is not None: + failed_keys.add(failed_key) + if not failed_keys: + raise + return cls.model_validate({ + key: _deserialize_value(value, path=f"{cls.__name__}.{key}") if key in failed_keys else value + for key, value in raw.items() + }) return decode_pydantic diff --git a/python/packages/core/tests/core/test_sessions.py b/python/packages/core/tests/core/test_sessions.py index e778f400e2..f37b73fe3a 100644 --- a/python/packages/core/tests/core/test_sessions.py +++ b/python/packages/core/tests/core/test_sessions.py @@ -904,6 +904,67 @@ class PydanticState(BaseModel): assert isinstance(restored.state["value"], PydanticState) assert restored.state["value"].value == "ok" + def test_registered_pydantic_state_restores_nested_registered_values(self) -> None: + from pydantic import BaseModel, ConfigDict + + class PydanticMessageState(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True) + + message: Message + + register_state_type(PydanticMessageState, type_id="pydantic_message_state_test") + session = AgentSession(session_id="nested-message") + session.state["value"] = PydanticMessageState(message=Message(role="user", contents=["persisted"])) + + restored = AgentSession.from_dict(session.to_dict()) + + assert isinstance(restored.state["value"], PydanticMessageState) + assert isinstance(restored.state["value"].message, Message) + assert restored.state["value"].message.text == "persisted" + + def test_registered_pydantic_state_keeps_tag_shaped_dict_fields_raw(self) -> None: + from pydantic import BaseModel + + class PydanticDictState(BaseModel): + data: dict[str, Any] + + register_state_type(PydanticDictState, type_id="pydantic_dict_state_test") + session = AgentSession(session_id="dict-tag-shaped") + # Ordinary dict data whose "type" happens to name a registered state + # type is data, not an encoded nested value. + session.state["value"] = PydanticDictState(data={"type": "message", "note": "plain data"}) + + restored = AgentSession.from_dict(session.to_dict()) + + assert isinstance(restored.state["value"], PydanticDictState) + assert restored.state["value"].data == {"type": "message", "note": "plain data"} + + def test_registered_pydantic_state_restores_only_failed_fields(self) -> None: + """A mixed model restores the broken nested field and keeps valid tag-shaped data raw.""" + from pydantic import BaseModel, ConfigDict + + class PydanticMixedState(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True) + + message: Message + data: dict[str, Any] + + register_state_type(PydanticMixedState, type_id="pydantic_mixed_state_test") + session = AgentSession(session_id="mixed-restore") + # `message` fails raw validation (its payload is the tagged dict), while + # `data` validates raw; the retry must restore only `message`. + session.state["value"] = PydanticMixedState( + message=Message(role="user", contents=["persisted"]), + data={"type": "message", "note": "plain data"}, + ) + + restored = AgentSession.from_dict(session.to_dict()) + + assert isinstance(restored.state["value"], PydanticMixedState) + assert isinstance(restored.state["value"].message, Message) + assert restored.state["value"].message.text == "persisted" + assert restored.state["value"].data == {"type": "message", "note": "plain data"} + def test_conflicting_type_identifier_is_rejected(self) -> None: class FirstState: def to_dict(self) -> dict[str, Any]: @@ -1149,6 +1210,65 @@ async def test_round_trips_session_across_store_instances(self, tmp_path: Path) assert files[0].suffix == ".json" assert json.loads(files[0].read_bytes())["version"] == "1.0" + async def test_round_trips_registered_pydantic_state_with_nested_message(self, tmp_path: Path) -> None: + from pydantic import BaseModel, ConfigDict + + class PydanticFileMessageState(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True) + + message: Message + + register_state_type(PydanticFileMessageState, type_id="pydantic_file_message_state_test") + session = AgentSession(session_id="file-nested-message") + session.state["value"] = PydanticFileMessageState(message=Message(role="user", contents=["persisted"])) + + await FileSessionStore(tmp_path).set("tenant_file-nested-message", session) + restored = await FileSessionStore(tmp_path).get("tenant_file-nested-message") + + assert restored is not None + assert isinstance(restored.state["value"], PydanticFileMessageState) + assert isinstance(restored.state["value"].message, Message) + assert restored.state["value"].message.text == "persisted" + + @pytest.mark.parametrize("serialization_format", ["json", "msgpack"]) + async def test_round_trips_registered_pydantic_state_with_aliased_nested_message( + self, tmp_path: Path, serialization_format: Literal["json", "msgpack"] + ) -> None: + from pydantic import BaseModel, ConfigDict, Field + + class PydanticAliasedMessageState(BaseModel): + model_config = ConfigDict( + arbitrary_types_allowed=True, + serialize_by_alias=True, + loc_by_alias=False, + ) + + last_message: Message = Field(alias="lastMessage") + + register_state_type( + PydanticAliasedMessageState, + type_id=f"pydantic_aliased_message_state_{serialization_format}_test", + ) + session = AgentSession(session_id=f"aliased-{serialization_format}") + session.state["value"] = PydanticAliasedMessageState(lastMessage=Message(role="user", contents=["persisted"])) + + await FileSessionStore(tmp_path, serialization_format=serialization_format).set( + "tenant_aliased-message", session + ) + restored = await FileSessionStore(tmp_path, serialization_format=serialization_format).get( + "tenant_aliased-message" + ) + + assert restored is not None + assert isinstance(restored.state["value"], PydanticAliasedMessageState) + assert isinstance(restored.state["value"].last_message, Message) + assert restored.state["value"].last_message.text == "persisted" + files = await asyncio.to_thread(lambda: list(tmp_path.iterdir())) + assert len(files) == 1 + serialized = await asyncio.to_thread(files[0].read_bytes) + assert b"lastMessage" in serialized + assert b"last_message" not in serialized + async def test_round_trips_binary_messagepack_session(self, tmp_path: Path) -> None: store = FileSessionStore(tmp_path, serialization_format="msgpack") session = AgentSession(session_id="binary-session")