Repository navigation
Python: fix(core): deep-restore nested registered values in the default Pydantic state decoder - #9250
Yufeng He (he-yufeng) wants to merge 4 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.
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_valuebefore validation, mirroring what the encoder side already does via_serialize_value.to_dict/from_dictdefault path is unchanged; unregistered models keep the implicit-registration compatibility path; unknown type tags still pass through as raw dicts.Regression evidence (run locally, all green): two new tests in
python/packages/core/tests/core/test_sessions.pycover the report,test_registered_pydantic_state_restores_nested_registered_valuesfor theAgentSession.from_dictround-trip andtest_round_trips_registered_pydantic_state_with_nested_messagefor a freshFileSessionStoreread; both are red before the fix and green after. Full filepytest tests/core/test_sessions.py: 189 passed. Sibling suitestest_serializable_mixin.py+test_types.py: 474 passed.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.