Repository navigation
Python: fix(core): deep-restore nested registered values in the default Pydantic state decoder - #9250
Yufeng He (he-yufeng) wants to merge 7 commits into
Conversation
…lt 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 microsoft#9246
GitHub Actions
Inactive
GitHub Actions
Inactive
agent-framework-automation
Bot
added
the
python
There was a problem hiding this comment.
🟢 Approval recommended
The focused decoder change preserves existing codec contracts and is adequately covered by regression tests.
0 open findings
What changed in this PR
Fixes nested registered-value restoration in default Pydantic session-state decoding.
Changes:
- Deep-deserializes Pydantic payload fields before validation.
- Adds direct and file-backed regression coverage.
| File | Description |
|---|---|
python/packages/core/agent_framework/_sessions.py |
Restores nested registered values before Pydantic validation. |
python/packages/core/tests/core/test_sessions.py |
Tests nested Message round trips. |
🧠 Review effort: Balanced
💡 Add a code-review agent skill for context-aware, tailored reviews. Learn more in the docs.
| # 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({ |
There was a problem hiding this comment.
Yufeng He (@he-yufeng) Deep-deserializing every Pydantic payload field can hijack ordinary dict data whose "type" happens to match a registered state type. A model field such as data: dict[str, Any] containing {"type": "message", ...} validates today, but this change converts it to Message before model_validate and then rejects it. Please preserve raw-valid dict fields (for example, validate raw first and only retry with restored values) and add a regression for a registered Pydantic model with such application data.
There was a problem hiding this comment.
Done in 6cbe9bd, exactly along the lines you suggested. decode_pydantic now validates the raw payload first and only retries with _deserialize_value-restored values when raw validation fails, so a data: dict[str, Any] field holding {"type": "message", ...} keeps its dict instead of being converted and rejected. Added test_registered_pydantic_state_keeps_tag_shaped_dict_fields_raw: it fails against the previous decoder (the tagged dict is restored to a Message and validation rejects the whole state) and passes now. The earlier nested-message regressions still pass, so the retry path keeps covering genuinely nested registered values. Full test_sessions.py green (190 passed on 3.14), ruff check and format clean; the commit also formats two earlier test cases that had missed ruff format.
|
Yufeng He (@he-yufeng) Please update the PR body to use the repository's current PR template at |
…lues 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.
GitHub Actions
Inactive
GitHub Actions
Inactive
|
Eduard van Valkenburg (@eavanvalkenburg) done, the body now follows the current template including the AI Assistance section with the assisted areas identified. |
| # 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({ |
There was a problem hiding this comment.
Yufeng He (@he-yufeng) The raw-first validation only protects a tag-shaped application dict when the whole model validates raw. If another field requires restoration, this retry deserializes every field, so a valid data: dict[str, Any] containing {"type": "message", ...} can still be converted or rejected. Please preserve fields that already validate raw and restore only the field(s) that need it, with a mixed-model regression.
There was a problem hiding this comment.
Done in 9881d30. The retry now reads the ValidationError locations and restores only the field(s) pydantic actually blames; every other field keeps its raw payload untouched, so a valid data: dict[str, Any] holding {"type": "message", ...} passes through byte-identical instead of being converted or rejected. Errors pydantic cannot attribute to a field (model-level validators with an empty loc) re-raise immediately rather than falling back to restoring everything.
The mixed-model regression is in: one field holds a nested Message (fails raw, gets restored to a real Message) and a sibling field holds a plain tag-shaped dict (stays the exact raw dict after round-trip).
GitHub Actions
Inactive
GitHub Actions
Inactive
Eduard van Valkenburg (eavanvalkenburg)
left a comment
Eduard van Valkenburg (eavanvalkenburg)
left a comment
There was a problem hiding this comment.
Yufeng He (@he-yufeng) Please address the inline alias-restoration finding before merge.
| if not failed_fields: | ||
| raise | ||
| return cls.model_validate({ | ||
| key: _deserialize_value(value, path=f"{cls.__name__}.{key}") if key in failed_fields else value |
There was a problem hiding this comment.
Yufeng He (@he-yufeng) This retry matches Pydantic error locations directly against serialized keys. With aliases such as last_message serialized as lastMessage, validation can report the field name while the input dictionary contains only the alias, so the tagged value is never selected for restoration and deserialization fails. Please resolve failed-field locations to the actual serialized input key using the model’s alias metadata, and add JSON and MessagePack round-trip regressions for an aliased nested Message field.
There was a problem hiding this comment.
Implemented in b8a260e0d9b2ae755839866847c0cd648cfb1ea8.
The retry now resolves each failed field location to the actual serialized input key. It checks the reported location directly first, then uses the model's validation_alias, alias, AliasChoices, and serialization_alias metadata. Only those resolved keys are restored; every other raw field remains untouched.
The new regression uses serialize_by_alias=True with loc_by_alias=False, so validation reports last_message while the persisted input contains only lastMessage. Both JSON and MessagePack fail before this change and now restore a real nested Message through a fresh FileSessionStore. The tests also inspect the persisted bytes to prove the alias path was exercised.
Final related sweep: 722 passed. Ruff, formatting, and Pyright on the changed module pass.
GitHub Actions
Inactive
GitHub Actions
Inactive
| if not failed_keys: | ||
| raise | ||
| return cls.model_validate({ | ||
| key: _deserialize_value(value, path=f"{cls.__name__}.{key}") if key in failed_keys else value |
There was a problem hiding this comment.
This still restores the entire top-level field when validation fails deeper inside a nested model. If the error is at child.message, a sibling such as child.data: Any can contain an ordinary tag-shaped dictionary and be silently converted into a registered framework type—or make restoration fail—even though that sibling already represented valid application data. Please preserve the complete validation-error location and restore only the failing nested value, resolving aliases and container indexes while retaining registered-codec boundaries. Add a regression with a nested Message field beside raw tag-shaped application data.
There was a problem hiding this comment.
{"body": "Fixed in f6078b7ff. The retry keyed on loc[0], so a failure pydantic reported at child.message restored the whole child, and _deserialize_value then walked the siblings too. A child.data 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 and 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.\n\nThe complete location is now walked instead of its first segment. Each segment resolves 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 rather than letting a key be read as a field name. On the codec boundary, a location that would reach inside a payload owned by a registered type restores that value through its codec as a whole, so a codec value is never rebuilt piecemeal. A deeper location supersedes a shallower one on the same path, so blaming both child and child.message repairs child.message alone.\n\nThree regressions, 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 against b8a260e0d on Failed to deserialize registered state type 'message' at ...data, which is the sibling being swept up.\n\n\npytest packages/core/tests/core/test_sessions.py 200 passed\npytest packages/core/tests 8983 passed, 22 skipped, 2 xfailed\nruff check / ruff format --check clean\n\n\npoe pyright on the whole workspace reports 24 errors, all in packages/declarative from an unresolved powerfx import in my environment and none in the file I changed. Run on the file alone, the head and b8a260e0d report the same single error, so this change adds no typing regressions. Pre-commit passes everything except bandit, which fails identically on files I did not touch with Could not read config file, so that one is my local bandit rather than the change.\n"}
…lt 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.
GitHub Actions
Inactive
|
Yufeng He (@he-yufeng) The deep-restoration walker still misses two supported nested shapes. First, Pydantic can include synthetic union-branch labels in error locations—for example |
Motivation & Context
A Pydantic application-state model with
arbitrary_types_allowed=Trueand aMessagefield, registered viaregister_state_typewith the default codecs, persisted fine but failed to restore through bothAgentSession.from_dict(session.to_dict())and a freshFileSessionStoreread: the nestedMessagearrived atmodel_validateas its tagged dict because the default decoder never deep-restored payload items.Description & Review Guide
_deserialize_value. Failed field locations are resolved to actual serialized input keys using the model's alias metadata, includingalias,validation_alias,AliasChoices, andserialization_alias, so an aliased nestedMessagealso round-trips through JSON and MessagePack.Regression evidence: nested
Message, tag-shaped dictionary, mixed-model, and aliasedlastMessagecases are covered. The final local run oftest_sessions.py,test_serializable_mixin.py, andtest_types.pyreports 722 passed with no skips. Ruff, formatting, and Pyright on the changed module pass.Related Issue
Fixes #9246
AI Assistance
AI assistance details: implementation, tests, and this PR description. The tests listed above were actually run locally and pass.
Contribution Checklist
breaking changelabel (or add "[BREAKING]" to the title prefix, before or after any language prefix). A workflow keeps the label and title prefix in sync automatically.