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/7] 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/7] 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/7] 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]:

From b8a260e0d9b2ae755839866847c0cd648cfb1ea8 Mon Sep 17 00:00:00 2001
From: Yufeng He <40085740+he-yufeng@users.noreply.github.com>
Date: Wed, 28 Oct 2026 00:26:39 +0800
Subject: [PATCH 4/7] Python: fix(core): restore aliased nested state values

---
 .../core/agent_framework/_sessions.py         | 35 +++++++++++++----
 .../packages/core/tests/core/test_sessions.py | 39 +++++++++++++++++++
 2 files changed, 66 insertions(+), 8 deletions(-)

diff --git a/python/packages/core/agent_framework/_sessions.py b/python/packages/core/agent_framework/_sessions.py
index 7ee2e7ee811..9adac362a0b 100644
--- a/python/packages/core/agent_framework/_sessions.py
+++ b/python/packages/core/agent_framework/_sessions.py
@@ -266,7 +266,7 @@ def decode_from_dict(payload: Mapping[str, Any]) -> Any:
 
         return decode_from_dict
 
-    from pydantic import BaseModel, ValidationError
+    from pydantic import AliasChoices, BaseModel, ValidationError
 
     if issubclass(cls, BaseModel):
 
@@ -282,15 +282,34 @@ def decode_pydantic(payload: Mapping[str, Any]) -> Any:
                 # 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:
+                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_fields else value
+                    key: _deserialize_value(value, path=f"{cls.__name__}.{key}") if key in failed_keys else value
                     for key, value in raw.items()
                 })
 
diff --git a/python/packages/core/tests/core/test_sessions.py b/python/packages/core/tests/core/test_sessions.py
index 23161b07790..f37b73fe3ae 100644
--- a/python/packages/core/tests/core/test_sessions.py
+++ b/python/packages/core/tests/core/test_sessions.py
@@ -1230,6 +1230,45 @@ class PydanticFileMessageState(BaseModel):
         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")

From f6078b7ffb372e1ef2adb60e8167c81b06671fe9 Mon Sep 17 00:00:00 2001
From: Yufeng He <40085740+he-yufeng@users.noreply.github.com>
Date: Mon, 2 Nov 2026 21:41:21 +0800
Subject: [PATCH 5/7] Python: fix(core): restore only the failing nested value
 in the default Pydantic state decoder

The retry keyed on the first segment of the validation-error location, so a failure pydantic reported at `child.message` restored the whole `child`. `_deserialize_value` then walked the siblings too, and a sibling that legitimately holds a tag-shaped dict, an ordinary `dict[str, Any]` whose `type` merely happens to name a registered type, was handed to that type's codec. It either came back as a framework object or raised `Failed to deserialize registered state type`, which failed the whole restore because of one broken field.

The complete error location is now walked instead. Each segment is resolved onto the raw payload through the field name and its aliases and through container indexes, so a failure two levels down behind `Field(alias=...)` inside a list or a keyed mapping repairs exactly that value and leaves its siblings byte-identical. The class governing each level comes from the annotations: a nested model's own fields resolve the keys beneath it, while `dict[str, Item]` and `list[Item]` treat their keys or positions as application data and hand the item class down to the values.

Where a location would reach inside a payload owned by a registered codec, that codec still restores the value as a whole, since only it knows the shape. A deeper location supersedes a shallower one on the same path, so blaming both `child` and `child.message` repairs `child.message` alone.

Three regressions pin it, each pairing a `Message` field with raw tag-shaped application data: a nested model, an aliased field behind a list index, and an aliased field behind a keyed container. All three fail without the fix, on `Failed to deserialize registered state type 'message' at ...data`, which is the sibling being swept up.

200 tests pass in packages/core/tests/core/test_sessions.py and 8983 in packages/core/tests. `poe pyright` reports the same errors as the base on the file changed, and ruff check and format are clean.
---
 .../core/agent_framework/_sessions.py         | 223 +++++++++++++++---
 .../packages/core/tests/core/test_sessions.py | 111 +++++++++
 2 files changed, 298 insertions(+), 36 deletions(-)

diff --git a/python/packages/core/agent_framework/_sessions.py b/python/packages/core/agent_framework/_sessions.py
index 9adac362a0b..0313096a1e3 100644
--- a/python/packages/core/agent_framework/_sessions.py
+++ b/python/packages/core/agent_framework/_sessions.py
@@ -33,7 +33,8 @@
 from dataclasses import dataclass
 from functools import partial
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypeAlias, TypeVar, cast
+from types import UnionType
+from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypeAlias, TypeVar, Union, cast, get_args, get_origin
 
 import msgspec
 from typing_extensions import TypedDict
@@ -257,6 +258,180 @@ def encode_pydantic(value: Any) -> Mapping[str, Any]:
     )
 
 
+_RESTORE_LEAF: Any = object()
+_LocationTree: TypeAlias = dict[Any, Any]
+
+
+def _field_alias_names(field: Any) -> list[str]:
+    """Return every serialized key one pydantic field may legitimately appear under."""
+    names: list[str] = []
+    validation_alias = getattr(field, "validation_alias", None)
+    if isinstance(validation_alias, str):
+        names.append(validation_alias)
+    else:
+        names.extend(choice for choice in getattr(validation_alias, "choices", ()) if isinstance(choice, str))
+    for alias in (getattr(field, "alias", None), getattr(field, "serialization_alias", None)):
+        if isinstance(alias, str) and alias not in names:
+            names.append(alias)
+    return names
+
+
+def _model_fields(value_cls: Any) -> Mapping[str, Any]:
+    fields = cast("Mapping[str, Any] | None", getattr(value_cls, "model_fields", None))
+    return fields if isinstance(fields, Mapping) else {}
+
+
+def _descend_classes(annotation: Any) -> tuple[Any | None, Any | None]:
+    """Name the classes one location segment meets: this value's, then its items'.
+
+    Pydantic reports locations as field names and container indexes, so walking
+    one needs to know what governs each level. A ``Child`` field is a model, so
+    its own fields resolve the keys beneath it. ``dict[str, Item]`` and
+    ``list[Item]`` are containers instead: their keys or positions are
+    application data, and the item class is what resolves names further down.
+    """
+    annotation_origin = getattr(annotation, "__origin__", None)
+    if getattr(annotation, "__metadata__", None) is not None and annotation_origin is not None:
+        return _descend_classes(annotation_origin)
+    origin = get_origin(annotation)
+    if origin is Union or origin is UnionType:
+        for arg in get_args(annotation):
+            if arg is None or arg is type(None):
+                continue
+            keyed_cls, item_cls = _descend_classes(arg)
+            if keyed_cls is not None or item_cls is not None:
+                return keyed_cls, item_cls
+        return None, None
+    if origin is not None:
+        # A container defers to its item type; the container itself is not a model.
+        for arg in get_args(annotation):
+            if arg is None or arg is type(None) or arg in (str, bytes):
+                continue
+            keyed_cls, item_cls = _descend_classes(arg)
+            if keyed_cls is not None or item_cls is not None:
+                return None, item_cls if item_cls is not None else keyed_cls
+        return None, None
+    if annotation is Any or annotation in (str, bytes) or not isinstance(annotation, type):
+        return None, None
+    return annotation, annotation
+
+
+def _location_tree(errors: Sequence[Mapping[str, Any]]) -> _LocationTree | None:
+    """Collapse validation-error locations into a tree of segments to restore.
+
+    A deeper location supersedes a shallower one on the same path. When pydantic
+    blames both ``child`` and ``child.message``, only ``child.message`` is
+    repaired, because restoring ``child`` would again sweep up its siblings.
+    """
+    tree: _LocationTree = {}
+    found = False
+    for error in errors:
+        location = cast("Sequence[Any] | None", error.get("loc"))
+        if location is None or not location:
+            continue
+        node: _LocationTree = tree
+        for segment in location[:-1]:
+            child = node.get(segment)
+            subtree: _LocationTree = cast("_LocationTree", child) if isinstance(child, dict) else {}
+            node[segment] = subtree
+            node = subtree
+        last = location[-1]
+        if not isinstance(node.get(last), dict):
+            node[last] = _RESTORE_LEAF
+        found = True
+    return tree if found else None
+
+
+def _restore_at_loc_tree(
+    value: Any, tree: Any, *, classes: tuple[Any | None, Any | None], path: str
+) -> tuple[Any, bool]:
+    keyed_cls, item_cls = classes
+    if tree is _RESTORE_LEAF:
+        return _deserialize_value(value, path=path), True
+    subtree = cast("_LocationTree | None", tree)
+    if subtree is None or not subtree:
+        return value, False
+    if isinstance(value, Mapping):
+        mapping = cast("Mapping[str, Any]", value)
+        if (
+            keyed_cls is not None
+            and _STATE_CLASS_REGISTRY.get(keyed_cls) is not None
+            and _STATE_TYPE_REGISTRY.get(str(mapping.get("type", ""))) is not None
+        ):
+            # The location reaches inside a payload owned by a registered codec.
+            # Only that codec knows the shape, so restore the value as a whole.
+            return _deserialize_value(mapping, path=path), True
+        return _restore_mapping(mapping, subtree, keyed_cls=keyed_cls, item_cls=item_cls, path=path)
+    if isinstance(value, list):
+        return _restore_list(cast("list[Any]", value), subtree, item_cls=item_cls, path=path)
+    return value, False
+
+
+def _restore_mapping(
+    value: Mapping[str, Any],
+    tree: Mapping[Any, Any],
+    *,
+    keyed_cls: Any,
+    item_cls: Any,
+    path: str,
+) -> tuple[Any, bool]:
+    fields = _model_fields(keyed_cls)
+    out = dict(value)
+    changed = False
+    for segment, sub in tree.items():
+        field_name, field = next(
+            (
+                (name, candidate)
+                for name, candidate in fields.items()
+                if name == segment or segment in _field_alias_names(candidate)
+            ),
+            (None, None),
+        )
+        candidates = [segment]
+        if field_name is not None:
+            candidates.append(field_name)
+            candidates.extend(_field_alias_names(field))
+        raw_key = next((key for key in candidates if key in value), None)
+        if raw_key is None:
+            continue
+        restored, key_changed = _restore_at_loc_tree(
+            out[raw_key],
+            sub,
+            classes=_descend_classes(getattr(field, "annotation", None)) if field is not None else (item_cls, item_cls),
+            path=f"{path}.{raw_key}",
+        )
+        if key_changed:
+            out[raw_key] = restored
+            changed = True
+    return (out, True) if changed else (value, False)
+
+
+def _restore_list(value: Sequence[Any], tree: Mapping[Any, Any], *, item_cls: Any, path: str) -> tuple[Any, bool]:
+    out = list(value)
+    changed = False
+    for segment, sub in tree.items():
+        if not isinstance(segment, int) or not 0 <= segment < len(out):
+            continue
+        restored, item_changed = _restore_at_loc_tree(
+            out[segment], sub, classes=(item_cls, item_cls), path=f"{path}[{segment}]"
+        )
+        if item_changed:
+            out[segment] = restored
+            changed = True
+    return (out, True) if changed else (value, False)
+
+
+def _restore_failed_values(
+    cls: type[Any], raw: Mapping[str, Any], errors: Sequence[Mapping[str, Any]], *, path: str
+) -> Mapping[str, Any] | None:
+    """Repair only the values pydantic blamed, or ``None`` when none can be located."""
+    tree = _location_tree(errors)
+    if tree is None:
+        return None
+    restored, changed = _restore_at_loc_tree(raw, tree, classes=(cls, cls), path=path)
+    return cast(Mapping[str, Any], restored) if changed else None
+
+
 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)):
@@ -266,7 +441,7 @@ def decode_from_dict(payload: Mapping[str, Any]) -> Any:
 
         return decode_from_dict
 
-    from pydantic import AliasChoices, BaseModel, ValidationError
+    from pydantic import BaseModel, ValidationError
 
     if issubclass(cls, BaseModel):
 
@@ -277,41 +452,17 @@ def decode_pydantic(payload: Mapping[str, Any]) -> Any:
             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:
+                # payload as its tagged dict. Restore only the values pydantic
+                # actually blames: the complete error location is walked, so a
+                # failure deep inside a nested model or a container repairs that
+                # one value and leaves every sibling untouched. A sibling that
+                # legitimately holds a tag-shaped dict (e.g. plain
+                # dict[str, Any] data) therefore keeps its payload instead of
+                # being converted or rejected by the retry.
+                restored = _restore_failed_values(cls, raw, exc.errors(), path=cls.__name__)
+                if restored is None:
                     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 cls.model_validate(restored)
 
         return decode_pydantic
 
diff --git a/python/packages/core/tests/core/test_sessions.py b/python/packages/core/tests/core/test_sessions.py
index f37b73fe3ae..3867fabf130 100644
--- a/python/packages/core/tests/core/test_sessions.py
+++ b/python/packages/core/tests/core/test_sessions.py
@@ -965,6 +965,117 @@ class PydanticMixedState(BaseModel):
         assert restored.state["value"].message.text == "persisted"
         assert restored.state["value"].data == {"type": "message", "note": "plain data"}
 
+    def test_registered_pydantic_state_restores_only_the_failing_nested_value(self) -> None:
+        """A nested failure restores that nested value only, not its tag-shaped sibling."""
+        from pydantic import BaseModel, ConfigDict
+
+        class NestedMixed(BaseModel):
+            model_config = ConfigDict(arbitrary_types_allowed=True)
+
+            message: Message
+            data: dict[str, Any]
+
+        class WithNestedMixed(BaseModel):
+            child: NestedMixed
+
+        register_state_type(WithNestedMixed, type_id="pydantic_nested_mixed_state_test")
+        session = AgentSession(session_id="nested-mixed-restore")
+        # Pydantic blames `child.message` only. Restoring the whole `child`
+        # would also run the codec over `child.data`, whose `type` merely
+        # happens to name a registered type, and either convert it or fail.
+        session.state["value"] = WithNestedMixed(
+            child=NestedMixed(
+                message=Message(role="user", contents=["persisted"]),
+                data={"type": "message", "note": "plain data"},
+            )
+        )
+
+        restored = AgentSession.from_dict(session.to_dict())
+
+        assert isinstance(restored.state["value"], WithNestedMixed)
+        assert isinstance(restored.state["value"].child.message, Message)
+        assert restored.state["value"].child.message.text == "persisted"
+        assert restored.state["value"].child.data == {"type": "message", "note": "plain data"}
+
+    def test_registered_pydantic_state_restores_nested_value_behind_alias_and_index(self) -> None:
+        """The failing value is found through a container index and a serialized alias."""
+        from pydantic import BaseModel, ConfigDict, Field
+
+        class AliasedItem(BaseModel):
+            model_config = ConfigDict(
+                arbitrary_types_allowed=True,
+                serialize_by_alias=True,
+                loc_by_alias=False,
+            )
+
+            last_message: Message = Field(alias="lastMessage")
+            data: dict[str, Any]
+
+        class AliasedWrapper(BaseModel):
+            model_config = ConfigDict(serialize_by_alias=True, loc_by_alias=False)
+
+            items: list[AliasedItem]
+
+        register_state_type(AliasedWrapper, type_id="pydantic_aliased_wrapper_state_test")
+        session = AgentSession(session_id="aliased-wrapper-restore")
+        # The error loc reports the field name `last_message` while the payload
+        # holds the alias `lastMessage`, two levels down past a list index.
+        session.state["value"] = AliasedWrapper(
+            items=[
+                AliasedItem(
+                    lastMessage=Message(role="user", contents=["persisted"]),
+                    data={"type": "message", "note": "plain data"},
+                )
+            ]
+        )
+
+        restored = AgentSession.from_dict(session.to_dict())
+
+        assert isinstance(restored.state["value"], AliasedWrapper)
+        item = restored.state["value"].items[0]
+        assert isinstance(item.last_message, Message)
+        assert item.last_message.text == "persisted"
+        assert item.data == {"type": "message", "note": "plain data"}
+
+    def test_registered_pydantic_state_restores_value_inside_a_keyed_container(self) -> None:
+        """A keyed container resolves aliases on its value type, not against the key."""
+        from pydantic import BaseModel, ConfigDict, Field
+
+        class KeyedItem(BaseModel):
+            model_config = ConfigDict(
+                arbitrary_types_allowed=True,
+                serialize_by_alias=True,
+                loc_by_alias=False,
+            )
+
+            last_message: Message = Field(alias="lastMessage")
+            data: dict[str, Any]
+
+        class KeyedWrapper(BaseModel):
+            model_config = ConfigDict(serialize_by_alias=True, loc_by_alias=False)
+
+            by_key: dict[str, KeyedItem]
+
+        register_state_type(KeyedWrapper, type_id="pydantic_keyed_wrapper_state_test")
+        session = AgentSession(session_id="keyed-wrapper-restore")
+        # The key is application data; the class governing the value is what
+        # resolves `last_message` onto the serialized `lastMessage`.
+        session.state["value"] = KeyedWrapper(
+            by_key={
+                "first": KeyedItem(
+                    lastMessage=Message(role="user", contents=["persisted"]),
+                    data={"type": "message", "note": "plain data"},
+                )
+            }
+        )
+
+        restored = AgentSession.from_dict(session.to_dict())
+
+        item = restored.state["value"].by_key["first"]
+        assert isinstance(item.last_message, Message)
+        assert item.last_message.text == "persisted"
+        assert item.data == {"type": "message", "note": "plain data"}
+
     def test_conflicting_type_identifier_is_rejected(self) -> None:
         class FirstState:
             def to_dict(self) -> dict[str, Any]:

From fd2d7241e7acf15bdb6b530d518ee0bb1170aa73 Mon Sep 17 00:00:00 2001
From: Yufeng He <40085740+he-yufeng@users.noreply.github.com>
Date: Wed, 4 Nov 2026 21:44:49 +0800
Subject: [PATCH 6/7] Python: fix(core): restore nested state through unions
 and containers

---
 .../core/agent_framework/_sessions.py         |  84 ++++++++----
 .../packages/core/tests/core/test_sessions.py | 121 ++++++++++++++++++
 2 files changed, 177 insertions(+), 28 deletions(-)

diff --git a/python/packages/core/agent_framework/_sessions.py b/python/packages/core/agent_framework/_sessions.py
index 0313096a1e3..2ac3379721e 100644
--- a/python/packages/core/agent_framework/_sessions.py
+++ b/python/packages/core/agent_framework/_sessions.py
@@ -282,40 +282,41 @@ def _model_fields(value_cls: Any) -> Mapping[str, Any]:
 
 
 def _descend_classes(annotation: Any) -> tuple[Any | None, Any | None]:
-    """Name the classes one location segment meets: this value's, then its items'.
+    """Return the model class and item annotation for the next location level.
 
-    Pydantic reports locations as field names and container indexes, so walking
-    one needs to know what governs each level. A ``Child`` field is a model, so
-    its own fields resolve the keys beneath it. ``dict[str, Item]`` and
-    ``list[Item]`` are containers instead: their keys or positions are
-    application data, and the item class is what resolves names further down.
+    Union annotations stay intact until a synthetic branch label identifies the
+    branch. Mappings defer to their value annotation, while fixed tuples keep
+    their complete annotation so a list index can select the element type.
     """
     annotation_origin = getattr(annotation, "__origin__", None)
     if getattr(annotation, "__metadata__", None) is not None and annotation_origin is not None:
         return _descend_classes(annotation_origin)
     origin = get_origin(annotation)
     if origin is Union or origin is UnionType:
-        for arg in get_args(annotation):
-            if arg is None or arg is type(None):
-                continue
-            keyed_cls, item_cls = _descend_classes(arg)
-            if keyed_cls is not None or item_cls is not None:
-                return keyed_cls, item_cls
-        return None, None
+        return annotation, annotation
+    if origin in (dict, Mapping):
+        args = get_args(annotation)
+        return None, args[1] if len(args) == 2 else None
+    if origin is tuple:
+        return None, annotation
     if origin is not None:
-        # A container defers to its item type; the container itself is not a model.
-        for arg in get_args(annotation):
-            if arg is None or arg is type(None) or arg in (str, bytes):
-                continue
-            keyed_cls, item_cls = _descend_classes(arg)
-            if keyed_cls is not None or item_cls is not None:
-                return None, item_cls if item_cls is not None else keyed_cls
-        return None, None
+        args = get_args(annotation)
+        return None, args[0] if args else None
     if annotation is Any or annotation in (str, bytes) or not isinstance(annotation, type):
         return None, None
     return annotation, annotation
 
 
+def _list_item_annotation(annotation: Any, index: int) -> Any:
+    """Select a fixed tuple's element annotation, or the variadic item type."""
+    if get_origin(annotation) is not tuple:
+        return annotation
+    args = get_args(annotation)
+    if len(args) == 2 and args[1] is Ellipsis:
+        return args[0]
+    return args[index] if index < len(args) else None
+
+
 def _location_tree(errors: Sequence[Mapping[str, Any]]) -> _LocationTree | None:
     """Collapse validation-error locations into a tree of segments to restore.
 
@@ -345,7 +346,7 @@ def _location_tree(errors: Sequence[Mapping[str, Any]]) -> _LocationTree | None:
 def _restore_at_loc_tree(
     value: Any, tree: Any, *, classes: tuple[Any | None, Any | None], path: str
 ) -> tuple[Any, bool]:
-    keyed_cls, item_cls = classes
+    keyed_cls, item_annotation = classes
     if tree is _RESTORE_LEAF:
         return _deserialize_value(value, path=path), True
     subtree = cast("_LocationTree | None", tree)
@@ -361,9 +362,9 @@ def _restore_at_loc_tree(
             # The location reaches inside a payload owned by a registered codec.
             # Only that codec knows the shape, so restore the value as a whole.
             return _deserialize_value(mapping, path=path), True
-        return _restore_mapping(mapping, subtree, keyed_cls=keyed_cls, item_cls=item_cls, path=path)
+        return _restore_mapping(mapping, subtree, keyed_cls=keyed_cls, item_annotation=item_annotation, path=path)
     if isinstance(value, list):
-        return _restore_list(cast("list[Any]", value), subtree, item_cls=item_cls, path=path)
+        return _restore_list(cast("list[Any]", value), subtree, item_annotation=item_annotation, path=path)
     return value, False
 
 
@@ -372,13 +373,31 @@ def _restore_mapping(
     tree: Mapping[Any, Any],
     *,
     keyed_cls: Any,
-    item_cls: Any,
+    item_annotation: Any,
     path: str,
 ) -> tuple[Any, bool]:
     fields = _model_fields(keyed_cls)
+    union_args = get_args(keyed_cls) if get_origin(keyed_cls) in (Union, UnionType) else ()
     out = dict(value)
     changed = False
     for segment, sub in tree.items():
+        branch_cls = next(
+            (arg for arg in union_args if isinstance(arg, type) and arg.__name__ == segment),
+            None,
+        )
+        if branch_cls is not None:
+            # A synthetic union branch label describes the payload's type; it
+            # is not a serialized key. Select that branch without consuming a key.
+            restored, branch_changed = _restore_at_loc_tree(
+                out,
+                sub,
+                classes=_descend_classes(branch_cls),
+                path=f"{path}.{segment}",
+            )
+            if branch_changed:
+                out = dict(cast("Mapping[str, Any]", restored))
+                changed = True
+            continue
         field_name, field = next(
             (
                 (name, candidate)
@@ -397,7 +416,11 @@ def _restore_mapping(
         restored, key_changed = _restore_at_loc_tree(
             out[raw_key],
             sub,
-            classes=_descend_classes(getattr(field, "annotation", None)) if field is not None else (item_cls, item_cls),
+            classes=(
+                _descend_classes(getattr(field, "annotation", None))
+                if field is not None
+                else (item_annotation, item_annotation)
+            ),
             path=f"{path}.{raw_key}",
         )
         if key_changed:
@@ -406,14 +429,19 @@ def _restore_mapping(
     return (out, True) if changed else (value, False)
 
 
-def _restore_list(value: Sequence[Any], tree: Mapping[Any, Any], *, item_cls: Any, path: str) -> tuple[Any, bool]:
+def _restore_list(
+    value: Sequence[Any], tree: Mapping[Any, Any], *, item_annotation: Any, path: str
+) -> tuple[Any, bool]:
     out = list(value)
     changed = False
     for segment, sub in tree.items():
         if not isinstance(segment, int) or not 0 <= segment < len(out):
             continue
         restored, item_changed = _restore_at_loc_tree(
-            out[segment], sub, classes=(item_cls, item_cls), path=f"{path}[{segment}]"
+            out[segment],
+            sub,
+            classes=_descend_classes(_list_item_annotation(item_annotation, segment)),
+            path=f"{path}[{segment}]",
         )
         if item_changed:
             out[segment] = restored
diff --git a/python/packages/core/tests/core/test_sessions.py b/python/packages/core/tests/core/test_sessions.py
index 3867fabf130..4b648d8eace 100644
--- a/python/packages/core/tests/core/test_sessions.py
+++ b/python/packages/core/tests/core/test_sessions.py
@@ -1380,6 +1380,127 @@ class PydanticAliasedMessageState(BaseModel):
         assert b"lastMessage" in serialized
         assert b"last_message" not in serialized
 
+    @pytest.mark.parametrize("serialization_format", ["json", "msgpack"])
+    async def test_round_trips_nested_message_behind_union_branch_label(
+        self, tmp_path: Path, serialization_format: Literal["json", "msgpack"]
+    ) -> None:
+        from pydantic import BaseModel, ConfigDict
+
+        class Item(BaseModel):
+            model_config = ConfigDict(arbitrary_types_allowed=True)
+
+            message: Message
+            data: dict[str, Any]
+
+        class Other(BaseModel):
+            other: str
+
+        class UnionWrapper(BaseModel):
+            child: Item | Other
+
+        register_state_type(UnionWrapper, type_id=f"pydantic_union_branch_{serialization_format}_test")
+        session = AgentSession(session_id=f"union-branch-{serialization_format}")
+        session.state["value"] = UnionWrapper(
+            child=Item(
+                message=Message(role="user", contents=["persisted"]),
+                data={"type": "message", "note": "plain data"},
+            )
+        )
+
+        await FileSessionStore(tmp_path, serialization_format=serialization_format).set("tenant_union-branch", session)
+        restored = await FileSessionStore(tmp_path, serialization_format=serialization_format).get(
+            "tenant_union-branch"
+        )
+
+        assert restored is not None
+        child = restored.state["value"].child
+        assert isinstance(child.message, Message)
+        assert child.message.text == "persisted"
+        assert child.data == {"type": "message", "note": "plain data"}
+
+    @pytest.mark.parametrize("serialization_format", ["json", "msgpack"])
+    async def test_round_trips_nested_message_in_integer_keyed_mapping(
+        self, tmp_path: Path, serialization_format: Literal["json", "msgpack"]
+    ) -> None:
+        from pydantic import BaseModel, ConfigDict, Field
+
+        class Item(BaseModel):
+            model_config = ConfigDict(
+                arbitrary_types_allowed=True,
+                serialize_by_alias=True,
+                loc_by_alias=False,
+            )
+
+            last_message: Message = Field(alias="lastMessage")
+            data: dict[str, Any]
+
+        class MappingWrapper(BaseModel):
+            by_key: dict[int, Item]
+
+        register_state_type(MappingWrapper, type_id=f"pydantic_integer_mapping_{serialization_format}_test")
+        session = AgentSession(session_id=f"integer-mapping-{serialization_format}")
+        session.state["value"] = MappingWrapper(
+            by_key={
+                7: Item(
+                    lastMessage=Message(role="user", contents=["persisted"]),
+                    data={"type": "message", "note": "plain data"},
+                )
+            }
+        )
+
+        await FileSessionStore(tmp_path, serialization_format=serialization_format).set(
+            "tenant_integer-mapping", session
+        )
+        restored = await FileSessionStore(tmp_path, serialization_format=serialization_format).get(
+            "tenant_integer-mapping"
+        )
+
+        assert restored is not None
+        item = restored.state["value"].by_key[7]
+        assert isinstance(item.last_message, Message)
+        assert item.last_message.text == "persisted"
+        assert item.data == {"type": "message", "note": "plain data"}
+
+    @pytest.mark.parametrize("serialization_format", ["json", "msgpack"])
+    async def test_round_trips_nested_message_at_fixed_tuple_index(
+        self, tmp_path: Path, serialization_format: Literal["json", "msgpack"]
+    ) -> None:
+        from pydantic import BaseModel, ConfigDict, Field
+
+        class Item(BaseModel):
+            model_config = ConfigDict(
+                arbitrary_types_allowed=True,
+                serialize_by_alias=True,
+                loc_by_alias=False,
+            )
+
+            last_message: Message = Field(alias="lastMessage")
+            data: dict[str, Any]
+
+        class TupleWrapper(BaseModel):
+            pair: tuple[int, Item]
+
+        register_state_type(TupleWrapper, type_id=f"pydantic_fixed_tuple_{serialization_format}_test")
+        session = AgentSession(session_id=f"fixed-tuple-{serialization_format}")
+        session.state["value"] = TupleWrapper(
+            pair=(
+                7,
+                Item(
+                    lastMessage=Message(role="user", contents=["persisted"]),
+                    data={"type": "message", "note": "plain data"},
+                ),
+            )
+        )
+
+        await FileSessionStore(tmp_path, serialization_format=serialization_format).set("tenant_fixed-tuple", session)
+        restored = await FileSessionStore(tmp_path, serialization_format=serialization_format).get("tenant_fixed-tuple")
+
+        assert restored is not None
+        item = restored.state["value"].pair[1]
+        assert isinstance(item.last_message, Message)
+        assert item.last_message.text == "persisted"
+        assert item.data == {"type": "message", "note": "plain data"}
+
     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 875a03bca79b84fac528aa0f79b59e80a1da2a37 Mon Sep 17 00:00:00 2001
From: Yufeng He <40085740+he-yufeng@users.noreply.github.com>
Date: Fri, 6 Nov 2026 10:52:15 +0800
Subject: [PATCH 7/7] Python: fix(core): isolate union state restoration
 branches

---
 .../core/agent_framework/_sessions.py         |  74 +++++--
 .../packages/core/tests/core/test_sessions.py | 187 ++++++++++++++++++
 2 files changed, 241 insertions(+), 20 deletions(-)

diff --git a/python/packages/core/agent_framework/_sessions.py b/python/packages/core/agent_framework/_sessions.py
index 2ac3379721e..2eb6482c969 100644
--- a/python/packages/core/agent_framework/_sessions.py
+++ b/python/packages/core/agent_framework/_sessions.py
@@ -281,11 +281,31 @@ def _model_fields(value_cls: Any) -> Mapping[str, Any]:
     return fields if isinstance(fields, Mapping) else {}
 
 
+def _union_branch_label(annotation: Any) -> str:
+    """Format a type the way pydantic labels its synthetic union branches."""
+    annotation_origin = getattr(annotation, "__origin__", None)
+    if getattr(annotation, "__metadata__", None) is not None and annotation_origin is not None:
+        return _union_branch_label(annotation_origin)
+    origin = get_origin(annotation)
+    if origin in (Union, UnionType):
+        return f"Union[{','.join(_union_branch_label(arg) for arg in get_args(annotation))}]"
+    if origin is not None:
+        origin_name = getattr(origin, "__name__", str(origin).removeprefix("typing."))
+        args = ",".join(_union_branch_label(arg) for arg in get_args(annotation))
+        return f"{origin_name}[{args}]"
+    if annotation is Ellipsis:
+        return "..."
+    if isinstance(annotation, str):
+        return repr(annotation)
+    return getattr(annotation, "__name__", str(annotation))
+
+
 def _descend_classes(annotation: Any) -> tuple[Any | None, Any | None]:
     """Return the model class and item annotation for the next location level.
 
-    Union annotations stay intact until a synthetic branch label identifies the
-    branch. Mappings defer to their value annotation, while fixed tuples keep
+    A nullable union with one non-None alternative unwraps to that alternative.
+    Other union annotations stay intact until a synthetic branch label identifies
+    the branch. Mappings defer to their value annotation, while fixed tuples keep
     their complete annotation so a list index can select the element type.
     """
     annotation_origin = getattr(annotation, "__origin__", None)
@@ -293,6 +313,9 @@ def _descend_classes(annotation: Any) -> tuple[Any | None, Any | None]:
         return _descend_classes(annotation_origin)
     origin = get_origin(annotation)
     if origin is Union or origin is UnionType:
+        non_none_args = tuple(arg for arg in get_args(annotation) if arg is not type(None))
+        if len(non_none_args) == 1:
+            return _descend_classes(non_none_args[0])
         return annotation, annotation
     if origin in (dict, Mapping):
         args = get_args(annotation)
@@ -343,6 +366,33 @@ def _location_tree(errors: Sequence[Mapping[str, Any]]) -> _LocationTree | None:
     return tree if found else None
 
 
+def _restore_union(value: Any, tree: Mapping[Any, Any], annotation: Any, path: str) -> tuple[Any, bool]:
+    """Repair one union branch at a time, always starting from the original value."""
+    from pydantic import TypeAdapter, ValidationError
+
+    for segment, sub in tree.items():
+        branch = next(
+            (arg for arg in get_args(annotation) if _union_branch_label(arg) == segment),
+            None,
+        )
+        if branch is None:
+            continue
+        restored, changed = _restore_at_loc_tree(
+            value,
+            sub,
+            classes=_descend_classes(branch),
+            path=f"{path}.{segment}",
+        )
+        if not changed:
+            continue
+        try:
+            TypeAdapter(branch).validate_python(restored)
+        except ValidationError:
+            continue
+        return restored, True
+    return value, False
+
+
 def _restore_at_loc_tree(
     value: Any, tree: Any, *, classes: tuple[Any | None, Any | None], path: str
 ) -> tuple[Any, bool]:
@@ -352,6 +402,8 @@ def _restore_at_loc_tree(
     subtree = cast("_LocationTree | None", tree)
     if subtree is None or not subtree:
         return value, False
+    if get_origin(keyed_cls) in (Union, UnionType):
+        return _restore_union(value, subtree, keyed_cls, path)
     if isinstance(value, Mapping):
         mapping = cast("Mapping[str, Any]", value)
         if (
@@ -377,27 +429,9 @@ def _restore_mapping(
     path: str,
 ) -> tuple[Any, bool]:
     fields = _model_fields(keyed_cls)
-    union_args = get_args(keyed_cls) if get_origin(keyed_cls) in (Union, UnionType) else ()
     out = dict(value)
     changed = False
     for segment, sub in tree.items():
-        branch_cls = next(
-            (arg for arg in union_args if isinstance(arg, type) and arg.__name__ == segment),
-            None,
-        )
-        if branch_cls is not None:
-            # A synthetic union branch label describes the payload's type; it
-            # is not a serialized key. Select that branch without consuming a key.
-            restored, branch_changed = _restore_at_loc_tree(
-                out,
-                sub,
-                classes=_descend_classes(branch_cls),
-                path=f"{path}.{segment}",
-            )
-            if branch_changed:
-                out = dict(cast("Mapping[str, Any]", restored))
-                changed = True
-            continue
         field_name, field = next(
             (
                 (name, candidate)
diff --git a/python/packages/core/tests/core/test_sessions.py b/python/packages/core/tests/core/test_sessions.py
index fc0ea25a372..64abf1fa8ae 100644
--- a/python/packages/core/tests/core/test_sessions.py
+++ b/python/packages/core/tests/core/test_sessions.py
@@ -1417,6 +1417,193 @@ class UnionWrapper(BaseModel):
         assert child.message.text == "persisted"
         assert child.data == {"type": "message", "note": "plain data"}
 
+    @pytest.mark.parametrize("serialization_format", ["json", "msgpack"])
+    async def test_union_branch_repairs_start_from_the_original_value(
+        self, tmp_path: Path, serialization_format: Literal["json", "msgpack"]
+    ) -> None:
+        from pydantic import BaseModel, ConfigDict
+
+        class Item(BaseModel):
+            model_config = ConfigDict(arbitrary_types_allowed=True)
+
+            message: Message
+            data: dict[str, Any]
+
+        class Other(BaseModel):
+            model_config = ConfigDict(arbitrary_types_allowed=True)
+
+            message: Message
+            data: Message
+
+        class UnionWrapper(BaseModel):
+            child: Item | Other
+
+        register_state_type(UnionWrapper, type_id=f"pydantic_union_isolation_{serialization_format}_test")
+        raw_data = Message(role="user", contents=["plain data"]).to_dict()
+        session = AgentSession(session_id=f"union-isolation-{serialization_format}")
+        session.state["value"] = UnionWrapper(
+            child=Item(message=Message(role="user", contents=["persisted"]), data=raw_data)
+        )
+
+        await FileSessionStore(tmp_path, serialization_format=serialization_format).set(
+            "tenant_union-isolation", session
+        )
+        restored = await FileSessionStore(tmp_path, serialization_format=serialization_format).get(
+            "tenant_union-isolation"
+        )
+
+        assert restored is not None
+        child = restored.state["value"].child
+        assert isinstance(child, Item)
+        assert isinstance(child.message, Message)
+        assert child.message.text == "persisted"
+        assert child.data == raw_data
+
+    @pytest.mark.parametrize("serialization_format", ["json", "msgpack"])
+    async def test_failed_union_branch_does_not_contaminate_fallback(
+        self, tmp_path: Path, serialization_format: Literal["json", "msgpack"]
+    ) -> None:
+        from pydantic import BaseModel, ConfigDict
+
+        class Item(BaseModel):
+            model_config = ConfigDict(arbitrary_types_allowed=True)
+
+            message: Message
+            data: Message
+            required: int
+
+        class Other(BaseModel):
+            model_config = ConfigDict(arbitrary_types_allowed=True)
+
+            message: Message
+            data: dict[str, Any]
+
+        class UnionWrapper(BaseModel):
+            child: Item | Other
+
+        register_state_type(UnionWrapper, type_id=f"pydantic_union_fallback_{serialization_format}_test")
+        raw_data = Message(role="user", contents=["plain data"]).to_dict()
+        session = AgentSession(session_id=f"union-fallback-{serialization_format}")
+        session.state["value"] = UnionWrapper(
+            child=Other(message=Message(role="user", contents=["persisted"]), data=raw_data)
+        )
+
+        await FileSessionStore(tmp_path, serialization_format=serialization_format).set(
+            "tenant_union-fallback", session
+        )
+        restored = await FileSessionStore(tmp_path, serialization_format=serialization_format).get(
+            "tenant_union-fallback"
+        )
+
+        assert restored is not None
+        child = restored.state["value"].child
+        assert isinstance(child, Other)
+        assert isinstance(child.message, Message)
+        assert child.message.text == "persisted"
+        assert child.data == raw_data
+
+    @pytest.mark.parametrize("serialization_format", ["json", "msgpack"])
+    async def test_round_trips_aliased_nested_message_inside_nullable_union(
+        self, tmp_path: Path, serialization_format: Literal["json", "msgpack"]
+    ) -> None:
+        from pydantic import BaseModel, ConfigDict, Field
+
+        class Item(BaseModel):
+            model_config = ConfigDict(
+                arbitrary_types_allowed=True,
+                serialize_by_alias=True,
+                loc_by_alias=False,
+            )
+
+            last_message: Message = Field(alias="lastMessage")
+            data: dict[str, Any]
+
+        class OptionalWrapper(BaseModel):
+            child: Item | None
+
+        register_state_type(OptionalWrapper, type_id=f"pydantic_optional_alias_{serialization_format}_test")
+        session = AgentSession(session_id=f"optional-alias-{serialization_format}")
+        session.state["value"] = OptionalWrapper(
+            child=Item(
+                lastMessage=Message(role="user", contents=["persisted"]),
+                data={"type": "message", "note": "plain data"},
+            )
+        )
+
+        await FileSessionStore(tmp_path, serialization_format=serialization_format).set(
+            "tenant_optional-alias", session
+        )
+        restored = await FileSessionStore(tmp_path, serialization_format=serialization_format).get(
+            "tenant_optional-alias"
+        )
+
+        assert restored is not None
+        child = restored.state["value"].child
+        assert isinstance(child.last_message, Message)
+        assert child.last_message.text == "persisted"
+        assert child.data == {"type": "message", "note": "plain data"}
+
+    @pytest.mark.parametrize("serialization_format", ["json", "msgpack"])
+    @pytest.mark.parametrize("container_kind", ["mapping", "sequence"])
+    async def test_round_trips_nested_message_inside_container_union_branch(
+        self,
+        tmp_path: Path,
+        serialization_format: Literal["json", "msgpack"],
+        container_kind: Literal["mapping", "sequence"],
+    ) -> None:
+        from pydantic import BaseModel, ConfigDict, Field
+
+        class Item(BaseModel):
+            model_config = ConfigDict(
+                arbitrary_types_allowed=True,
+                serialize_by_alias=True,
+                loc_by_alias=False,
+            )
+
+            last_message: Message = Field(alias="lastMessage")
+            data: dict[str, Any]
+
+        class Other(BaseModel):
+            other: str
+
+        item = Item(
+            lastMessage=Message(role="user", contents=["persisted"]),
+            data={"type": "message", "note": "plain data"},
+        )
+        if container_kind == "mapping":
+
+            class ContainerUnionWrapper(BaseModel):
+                child: dict[str, Item] | Other
+
+            child_value: Any = {"first": item}
+        else:
+
+            class ContainerUnionWrapper(BaseModel):
+                child: list[Item] | Other
+
+            child_value = [item]
+
+        register_state_type(
+            ContainerUnionWrapper,
+            type_id=f"pydantic_{container_kind}_union_{serialization_format}_test",
+        )
+        session = AgentSession(session_id=f"{container_kind}-union-{serialization_format}")
+        session.state["value"] = ContainerUnionWrapper(child=child_value)
+
+        await FileSessionStore(tmp_path, serialization_format=serialization_format).set(
+            f"tenant_{container_kind}-union", session
+        )
+        restored = await FileSessionStore(tmp_path, serialization_format=serialization_format).get(
+            f"tenant_{container_kind}-union"
+        )
+
+        assert restored is not None
+        child = restored.state["value"].child
+        restored_item = child["first"] if container_kind == "mapping" else child[0]
+        assert isinstance(restored_item.last_message, Message)
+        assert restored_item.last_message.text == "persisted"
+        assert restored_item.data == {"type": "message", "note": "plain data"}
+
     @pytest.mark.parametrize("serialization_format", ["json", "msgpack"])
     async def test_round_trips_nested_message_in_integer_keyed_mapping(
         self, tmp_path: Path, serialization_format: Literal["json", "msgpack"]