diff --git a/python/packages/core/agent_framework/_sessions.py b/python/packages/core/agent_framework/_sessions.py index 7ea242078a2..2eb6482c969 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,242 @@ 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 _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. + + 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) + 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: + 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) + return None, args[1] if len(args) == 2 else None + if origin is tuple: + return None, annotation + if origin is not 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. + + 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_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]: + keyed_cls, item_annotation = 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 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 ( + 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_annotation=item_annotation, path=path) + if isinstance(value, list): + return _restore_list(cast("list[Any]", value), subtree, item_annotation=item_annotation, path=path) + return value, False + + +def _restore_mapping( + value: Mapping[str, Any], + tree: Mapping[Any, Any], + *, + keyed_cls: Any, + item_annotation: 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_annotation, item_annotation) + ), + 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_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=_descend_classes(_list_item_annotation(item_annotation, segment)), + 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,12 +503,28 @@ 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 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(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 8558777ce60..64abf1fa8ae 100644 --- a/python/packages/core/tests/core/test_sessions.py +++ b/python/packages/core/tests/core/test_sessions.py @@ -903,6 +903,178 @@ 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_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]: @@ -1148,6 +1320,373 @@ 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 + + @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_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"] + ) -> 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")