Python: fix(core): deep-restore nested registered values in the default Pydantic state decoder

python/packages/core/agent_framework/_sessions.py

@@ -266,12 +266,33 @@ def decode_from_dict(payload: Mapping[str, Any]) -> Any:
 
         return decode_from_dict
 
-    from pydantic import BaseModel
+    from pydantic import 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_fields = {
+                    error["loc"][0]
+                    for error in exc.errors()
+                    if error.get("loc") and isinstance(error["loc"][0], str)
+                }
+                if not failed_fields:
+                    raise
+                return cls.model_validate({
+                    key: _deserialize_value(value, path=f"{cls.__name__}.{key}") if key in failed_fields else value
+                    for key, value in raw.items()
+                })
 
         return decode_pydantic
 

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

@@ -903,6 +903,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]:
@@ -1148,6 +1209,26 @@ 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"
+
     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")