From 254df281e1a7e8358f6e49de3b78a015414dea59 Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Date: Sat, 10 Oct 2026 21:34:44 +0800 Subject: [PATCH 1/3] Python: fix(core): deep-restore nested registered values in the default Pydantic state decoder The default encoder hands nested values to _serialize_value, which tags them for the registry, but the default Pydantic decoder validated the raw payload, so a nested Message arrived as a tagged dict and restoration of a registered state model failed on AgentSession.from_dict and on a fresh FileSessionStore read. Fixes #9246 --- .../core/agent_framework/_sessions.py | 11 ++++- .../packages/core/tests/core/test_sessions.py | 42 +++++++++++++++++++ 2 files changed, 52 insertions(+), 1 deletion(-) diff --git a/python/packages/core/agent_framework/_sessions.py b/python/packages/core/agent_framework/_sessions.py index 73bfbaa603b..11f6f368229 100644 --- a/python/packages/core/agent_framework/_sessions.py +++ b/python/packages/core/agent_framework/_sessions.py @@ -271,7 +271,16 @@ 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 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 nested values the same way before validating, + # otherwise Pydantic receives the tagged dict where the field + # expects the reconstructed instance. + return cls.model_validate({ + key: _deserialize_value(value, path=f"{cls.__name__}.{key}") + 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 1cf5d5adcfc..78c3317e3d5 100644 --- a/python/packages/core/tests/core/test_sessions.py +++ b/python/packages/core/tests/core/test_sessions.py @@ -901,6 +901,26 @@ 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_conflicting_type_identifier_is_rejected(self) -> None: class FirstState: def to_dict(self) -> dict[str, Any]: @@ -1146,6 +1166,28 @@ 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") From 6cbe9bd9b0865c37646929b0f537e4e21da36b00 Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Date: Tue, 13 Oct 2026 00:21:27 +0800 Subject: [PATCH 2/3] Python: fix(core): validate pydantic payloads raw before restoring values Deep-restoring every payload field up front hijacks a plain dict field whose "type" happens to name a registered state type: the tagged shape is converted before model_validate and then rejected. Validate the raw payload first and only retry with restored values when that fails, and cover a tag-shaped dict field staying raw. Also formats the two earlier nested-message cases that missed ruff format. --- .../core/agent_framework/_sessions.py | 24 ++++++++++-------- .../packages/core/tests/core/test_sessions.py | 25 ++++++++++++++----- 2 files changed, 32 insertions(+), 17 deletions(-) diff --git a/python/packages/core/agent_framework/_sessions.py b/python/packages/core/agent_framework/_sessions.py index 11f6f368229..b54427e7051 100644 --- a/python/packages/core/agent_framework/_sessions.py +++ b/python/packages/core/agent_framework/_sessions.py @@ -266,21 +266,23 @@ 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: - # 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 nested values the same way before validating, - # otherwise Pydantic receives the tagged dict where the field - # expects the reconstructed instance. - return cls.model_validate({ - key: _deserialize_value(value, path=f"{cls.__name__}.{key}") - 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: + # 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 nested values and retry; + # a field that legitimately holds a tag-shaped dict already + # validated raw above and never reaches this path. + return cls.model_validate({ + key: _deserialize_value(value, path=f"{cls.__name__}.{key}") 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 78c3317e3d5..2100ec0478e 100644 --- a/python/packages/core/tests/core/test_sessions.py +++ b/python/packages/core/tests/core/test_sessions.py @@ -911,9 +911,7 @@ class PydanticMessageState(BaseModel): 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"]) - ) + session.state["value"] = PydanticMessageState(message=Message(role="user", contents=["persisted"])) restored = AgentSession.from_dict(session.to_dict()) @@ -921,6 +919,23 @@ class PydanticMessageState(BaseModel): 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_conflicting_type_identifier_is_rejected(self) -> None: class FirstState: def to_dict(self) -> dict[str, Any]: @@ -1176,9 +1191,7 @@ class PydanticFileMessageState(BaseModel): 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"]) - ) + 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") From 9881d30f8c96a41e1453e3b05c236843f310297e Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Date: Fri, 16 Oct 2026 20:25:34 +0800 Subject: [PATCH 3/3] Python: fix(core): restore only the pydantic fields that fail raw validation --- .../core/agent_framework/_sessions.py | 20 ++++++++++---- .../packages/core/tests/core/test_sessions.py | 26 +++++++++++++++++++ 2 files changed, 41 insertions(+), 5 deletions(-) diff --git a/python/packages/core/agent_framework/_sessions.py b/python/packages/core/agent_framework/_sessions.py index b54427e7051..46da0817d20 100644 --- a/python/packages/core/agent_framework/_sessions.py +++ b/python/packages/core/agent_framework/_sessions.py @@ -274,14 +274,24 @@ def decode_pydantic(payload: Mapping[str, Any]) -> Any: raw = {key: value for key, value in payload.items() if key != "type"} try: return cls.model_validate(raw) - except ValidationError: + 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 nested values and retry; - # a field that legitimately holds a tag-shaped dict already - # validated raw above and never reaches this path. + # 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}") for key, value in raw.items() + 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 diff --git a/python/packages/core/tests/core/test_sessions.py b/python/packages/core/tests/core/test_sessions.py index 2100ec0478e..14a73e78808 100644 --- a/python/packages/core/tests/core/test_sessions.py +++ b/python/packages/core/tests/core/test_sessions.py @@ -936,6 +936,32 @@ class PydanticDictState(BaseModel): 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]: