Python: fix(core): deep-restore nested registered values in the default Pydantic state decoder by he-yufeng · Pull Request #9250 · microsoft/agent-framework

Yufeng He (he-yufeng) · 2026-10-10T13:35:31Z

Motivation & Context

A Pydantic application-state model with arbitrary_types_allowed=True and a Message field, registered via register_state_type with the default codecs, persisted fine but failed to restore through both AgentSession.from_dict(session.to_dict()) and a fresh FileSessionStore read: the nested Message arrived at model_validate as its tagged dict because the default decoder never deep-restored payload items.

Description & Review Guide

What are the major changes? The default Pydantic decoder validates raw input first and retries only the values implicated by validation errors. Complete error locations are walked through nested models, aliases, mapping values, and container indexes. Pydantic's synthetic union-branch labels select model and container branches instead of being treated as serialized keys. Arbitrary-type is-instance labels and discriminated-union values also resolve through independent candidate validation. Each union branch restores from the original value and contributes repairs only when that branch validates independently; a candidate codec failure is isolated to that candidate. Nullable unions unwrap a single non-None alternative before resolving field metadata. Mapping values use their value annotation, and fixed tuples select annotations by index. Registered-codec payloads remain whole-value boundaries.
What is the impact of these changes? Nested registered values round-trip through JSON and MessagePack without converting ordinary tag-shaped application dictionaries or leaking repairs between union branches. Custom codecs keep the raw-payload contract, valid raw siblings stay byte-identical, unregistered models retain the implicit-registration compatibility path, and unknown type tags still pass through as raw dicts.
What do you want reviewers to focus on? The raw-first retry boundary, independent union candidates and fallback, schema-defined branch labels, per-candidate codec failure isolation, nullable and container union normalization, value-versus-key annotation selection, fixed tuple indexes, and preservation of valid application data beside each restored value.

Regression evidence includes nested Message, tag-shaped dictionary, mixed-model, aliased, union-branch isolation and fallback, arbitrary-type union, discriminated union, candidate codec failure, nullable union, list and dict container union, integer-keyed mapping, and fixed-tuple cases. The new regressions run through fresh JSON and MessagePack stores and failed before the corresponding changes. After merging current main, the complete packages/core/tests run on the final head exits successfully with 9471 passed, 22 skipped, and 2 xfailed. Ruff check and formatting pass. Pyright, mypy, pyrefly, ty, and zuban all pass for core.

Related Issue

Fixes #9246

AI Assistance

No material AI assistance was used.
This is an AI-assisted contribution. I reviewed, understood, and verified all submitted content and accept responsibility for it.

AI assistance details: implementation, tests, and this PR description. The tests listed above were actually run locally and pass.

Contribution Checklist

The code builds clean without any errors or warnings
All unit tests pass, and I have added new tests where possible
The PR follows the Contribution Guidelines
This PR links to an agreed issue with no competing open PR, or the Related Issue section documents a trivial-change or repository-automation exception.
This is not a breaking change. If it is a breaking change, add the breaking change label (or add "[BREAKING]" to the title prefix, before or after any language prefix). A workflow keeps the label and title prefix in sync automatically.

…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

Copilot

🟢 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.

Eduard van Valkenburg (eavanvalkenburg) · 2026-10-12T07:49:46Z

+            # 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({


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.

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.

Eduard van Valkenburg (eavanvalkenburg) · 2026-10-12T07:49:49Z

Yufeng He (@he-yufeng) Please update the PR body to use the repository's current PR template at .github/pull_request_template.md, including its required headings and fields.

…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.

Yufeng He (he-yufeng) · 2026-10-12T23:11:37Z

Eduard van Valkenburg (@eavanvalkenburg) done, the body now follows the current template including the AI Assistance section with the assisted areas identified.

Eduard van Valkenburg (eavanvalkenburg) · 2026-10-13T11:06:25Z

+                # 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({


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.

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).

…idation

Eduard van Valkenburg (eavanvalkenburg)

Yufeng He (@he-yufeng) Please address the inline alias-restoration finding before merge.

Eduard van Valkenburg (eavanvalkenburg) · 2026-10-26T15:23:10Z

+                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


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.

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.

Eduard van Valkenburg (eavanvalkenburg) · 2026-11-02T11:25:49Z

+                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


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.

{"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.

Eduard van Valkenburg (eavanvalkenburg) · 2026-11-03T12:07:16Z

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 child.Item.message—and the walker treats Item as a serialized dictionary key, so restoration is skipped. Second, the generic-argument loop uses the first argument for every container, which selects int instead of Item for dict[int, Item] and for the Item position in tuple[int, Item]; aliased nested Message values then remain plain dictionaries. Please normalize union-branch segments, select mapping value annotations and tuple annotations by index, and add JSON and MessagePack round-trip regressions for these cases.

…ted-message

Yufeng He (he-yufeng) · 2026-11-04T14:05:22Z

Eduard van Valkenburg (@eavanvalkenburg) Implemented in fd2d7241e.

Union annotations now stay intact until a synthetic branch label such as Item selects the matching model, so that segment is no longer looked up as a serialized key. Container tracking now uses a mapping's value annotation and selects fixed tuple annotations by the actual index, including variadic tuple items.

Added JSON and MessagePack round trips for all three cases: child.Item.message, dict[int, Item], and tuple[int, Item]. Every case also retains a raw tag-shaped application-data sibling. All six failed before this change and pass now.

After merging current main and syncing the repository's full test environment, the complete packages/core/tests run exits successfully. Ruff and formatting pass. Pyright on the changed module still reports only the pre-existing _run_identity_scope unused-function finding outside this diff.

Eduard van Valkenburg (eavanvalkenburg) · 2026-11-05T13:22:26Z

+        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(


Yufeng He (@he-yufeng) Each union branch must start from the original value — The repair loop applies changes from every failed union branch to one shared copy. For A | B, trying B after A can rewrite fields that were valid for A, and final validation can succeed with silently changed application data. Please try each branch against a fresh copy and keep only the successful branch's repairs.

Implemented in 875a03bca. Union branches now restore independently from the original value, and only a branch that validates with its own repairs can supply them. The regressions cover both a successful first branch and a fallback after the first branch still fails validation. Both previously allowed one branch's Message restoration to contaminate a sibling's raw dict; JSON and MessagePack now preserve the intended branch and raw data. Full core: 9361 passed, 20 skipped, 2 xfailed.

Eduard van Valkenburg (eavanvalkenburg) · 2026-11-05T13:22:28Z

+    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:


Yufeng He (@he-yufeng) Item | None can lose the nested field metadata — Pydantic can report child.last_message without an Item branch segment. This code then keeps the union instead of Item's field metadata, so lastMessage is not mapped back to last_message and the nested Message stays a dictionary. Please unwrap a nullable union with one non-None alternative before resolving fields.

Implemented in 875a03bca. A nullable union with exactly one non-None alternative now unwraps to that alternative before resolving field metadata. The new aliased Item | None regressions restore lastMessage as a Message and preserve raw tag-shaped sibling data through both JSON and MessagePack. Core pyright with the Python 3.11 target reports 0 errors.

Eduard van Valkenburg (eavanvalkenburg) · 2026-11-05T13:22:29Z

+        restored, key_changed = _restore_at_loc_tree(
+            out[raw_key],
+            sub,
+            classes=(


Yufeng He (@he-yufeng) Nested models inside mapping values are not reached — For values such as list[Item] or dict[str, Item], this path passes the container annotation without descending to Item. Alias metadata is lost, so nested Message values remain dictionaries after restore. Please normalize the mapping value annotation before walking it.

Implemented in 875a03bca. Synthetic union labels now resolve generic container annotations as well as classes, so dict[str, Item] and list[Item] branches descend into Item before alias resolution. New JSON and MessagePack cases cover both container kinds and keep raw sibling data unchanged. Full core: 9361 passed, 20 skipped, 2 xfailed; Ruff syntax and Python 3.11 pyright pass.

Eduard van Valkenburg (eavanvalkenburg) · 2026-11-06T09:31:47Z

Yufeng He (@he-yufeng), I found two material union-restore gaps on the current head:

python/packages/core/agent_framework/_sessions.py:375: Pydantic can label union branches as is-instance[Message] or by discriminator values rather than the class name. Those branches are skipped, so tagged dictionaries remain unrestored. Please support schema-defined branch labels and add arbitrary-type/discriminated-union coverage.
python/packages/core/agent_framework/_sessions.py:380: if an unsuccessful union candidate's codec raises before validation, restoration aborts instead of trying the next valid branch. Please isolate codec failures per candidate and add JSON and MessagePack regressions.

…d-message

Yufeng He (he-yufeng) · 2026-11-06T12:08:34Z

Eduard van Valkenburg (@eavanvalkenburg) Implemented both follow-ups in 13ddfbc3c and merged current main in f70d63879.

Union recovery no longer depends only on class-name labels. For is-instance[Message], discriminator values, and other schema-defined labels, each candidate restores from the original value and must validate independently before supplying repairs. Candidate validation preserves arbitrary-type support without applying a model config to BaseModel branches. Codec failures are scoped to the candidate that raised, so a later valid branch still gets its chance.

Added JSON and MessagePack round trips for an arbitrary-type union, an aliased nested Message in a discriminated union, and a failing first-candidate codec beside a valid fallback with raw application data. All six regressions fail on the previous head and pass now. The container fixture redeclaration that broke the four typing checkers is also removed.

On the merged head, the complete core suite passes with 9471 passed, 22 skipped, and 2 xfailed. Ruff, formatting, Pyright, mypy, pyrefly, ty, and zuban all pass for core.

Copilot AI balanced review requested due to automatic review settings October 10, 2026 13:35

Yufeng He (he-yufeng) temporarily deployed to github-app-auth October 10, 2026 13:35 — with GitHub Actions Inactive

agent-framework-automation Bot added the python Usage: [Issues, PRs], Target: Python label Oct 10, 2026

Copilot started reviewing on behalf of Yufeng He (he-yufeng) October 10, 2026 13:36

View session

Copilot AI reviewed Oct 10, 2026

View reviewed changes

Eduard van Valkenburg (eavanvalkenburg) reviewed Oct 12, 2026

View reviewed changes

Yufeng He (he-yufeng) temporarily deployed to github-app-auth October 12, 2026 16:21 — with GitHub Actions Inactive

Yufeng He (he-yufeng) temporarily deployed to github-app-auth October 12, 2026 23:11 — with GitHub Actions Inactive

Eduard van Valkenburg (eavanvalkenburg) reviewed Oct 13, 2026

View reviewed changes

Python: fix(core): restore only the pydantic fields that fail raw val…

9881d30

…idation

Yufeng He (he-yufeng) temporarily deployed to github-app-auth October 16, 2026 12:26 — with GitHub Actions Inactive

Merge branch 'main' into fix/state-codec-nested-message

9662713

Yufeng He (he-yufeng) temporarily deployed to github-app-auth October 23, 2026 13:06 — with GitHub Actions Inactive

Eduard van Valkenburg (eavanvalkenburg) requested changes Oct 26, 2026

View reviewed changes

Yufeng He (he-yufeng) added 2 commits October 28, 2026 00:15

Merge remote-tracking branch 'origin/main' into HEAD

4808752

Python: fix(core): restore aliased nested state values

b8a260e

Yufeng He (he-yufeng) temporarily deployed to github-app-auth October 27, 2026 16:26 — with GitHub Actions Inactive

Yufeng He (he-yufeng) temporarily deployed to github-app-auth October 27, 2026 16:27 — with GitHub Actions Inactive

Eduard van Valkenburg (eavanvalkenburg) reviewed Nov 2, 2026

View reviewed changes

Yufeng He (he-yufeng) temporarily deployed to github-app-auth November 2, 2026 13:42 — with GitHub Actions Inactive

Eduard van Valkenburg (eavanvalkenburg) mentioned this pull request Nov 3, 2026

Python: [Bug]: Default registered Pydantic state cannot restore a Message field #9246

Open

3 tasks done

Yufeng He (he-yufeng) added 2 commits November 4, 2026 21:44

Python: fix(core): restore nested state through unions and containers

fd2d724

Merge remote-tracking branch 'upstream/main' into fix/state-codec-nes…

7b8c6c4

…ted-message

Yufeng He (he-yufeng) temporarily deployed to github-app-auth November 4, 2026 14:04 — with GitHub Actions Inactive

Yufeng He (he-yufeng) temporarily deployed to github-app-auth November 4, 2026 14:05 — with GitHub Actions Inactive

Yufeng He (he-yufeng) marked this pull request as ready for review November 5, 2026 10:28

Yufeng He (he-yufeng) requested review from SergeyMenshykh, Darren Cohen (dargilco), Jose Alvarez (jpalvarezl), Roger Barreto (rogerbarreto), sophia-ramsey and westey (westey-m) as code owners November 5, 2026 10:28

Yufeng He (he-yufeng) temporarily deployed to github-app-auth November 5, 2026 10:29 — with GitHub Actions Inactive

Eduard van Valkenburg (eavanvalkenburg) reviewed Nov 5, 2026

View reviewed changes

Python: fix(core): isolate union state restoration branches

875a03b

Yufeng He (he-yufeng) temporarily deployed to github-app-auth November 6, 2026 02:52 — with GitHub Actions Inactive

Yufeng He (he-yufeng) deployed to github-app-auth November 6, 2026 02:55 — with GitHub Actions Active

Yufeng He (he-yufeng) added 2 commits November 6, 2026 19:58

Python: fix(core): restore schema-labeled union branches

13ddfbc

Merge remote-tracking branch 'origin/main' into fix/state-codec-neste…

f70d638

…d-message

Yufeng He (he-yufeng) deployed to github-app-auth November 6, 2026 12:08 — with GitHub Actions Active

Yufeng He (he-yufeng) mentioned this pull request Nov 6, 2026

Python: Prefer injected DocumentDB clients over connection settings #9999

Open

6 of 7 tasks

Conversation

Yufeng He (he-yufeng) commented Oct 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation & Context

Description & Review Guide

Related Issue

AI Assistance

Contribution Checklist

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟢 Approval recommended

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Eduard van Valkenburg (eavanvalkenburg) commented Oct 12, 2026

Uh oh!

Yufeng He (he-yufeng) commented Oct 12, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Eduard van Valkenburg (eavanvalkenburg) left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Eduard van Valkenburg (eavanvalkenburg) commented Nov 3, 2026

Uh oh!

Yufeng He (he-yufeng) commented Nov 4, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Eduard van Valkenburg (eavanvalkenburg) commented Nov 6, 2026

Uh oh!

Yufeng He (he-yufeng) commented Nov 6, 2026

Uh oh!

This branch was successfully deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Yufeng He (he-yufeng) commented Oct 10, 2026 •

edited

Loading