From 90051ef7f6441552fa731e7bf6feb8cf3aa0ef42 Mon Sep 17 00:00:00 2001 From: "pydanty[bot]" <251103656+pydanty[bot]@users.noreply.github.com> Date: Fri, 9 Oct 2026 11:26:09 -0600 Subject: [PATCH] Fix `merge_json_schema_defs` corrupting branch schemas on transitive def-name collisions (#9735) Co-authored-by: pydanty Co-authored-by: Douwe Maan --- pydantic_ai_slim/pydantic_ai/_utils.py | 59 +++++----- tests/test_agent_output_schemas.py | 120 +++++++++++++++++++- tests/test_utils.py | 146 ++++++++++++++++++++++++- 3 files changed, 292 insertions(+), 33 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/_utils.py b/pydantic_ai_slim/pydantic_ai/_utils.py index 132fd3cccd..b8689d4b10 100644 --- a/pydantic_ai_slim/pydantic_ai/_utils.py +++ b/pydantic_ai_slim/pydantic_ai/_utils.py @@ -988,48 +988,45 @@ def merge_json_schema_defs(schemas: list[dict[str, Any]]) -> tuple[list[dict[str rewritten_schemas.append(schema) continue - schema = schema.copy() - defs = schema.pop('$defs', None) - schema_name_mapping: dict[str, str] = {} + schema = copy.deepcopy(schema) + defs = schema.pop('$defs') + renames: dict[str, str] = {} + new_names: list[str] = [] - # Process definitions and build mapping + # New names are kept; same-named defs with a different body get a unique name. Both point + # `all_defs` at this schema's own copy of the def, whose refs are rewritten in place below. for name, def_schema in defs.items(): if name not in all_defs: all_defs[name] = def_schema - schema_name_mapping[name] = name + new_names.append(name) elif def_schema != all_defs[name]: - # Different def with same name — assign a unique name - schema_name_mapping[name] = _unique_def_name(name, schema, all_defs) - all_defs[schema_name_mapping[name]] = def_schema - # else: structurally equal — handled below + renames[name] = _unique_def_name(name, schema, all_defs) + all_defs[renames[name]] = def_schema - # Defs that are structurally equal (same dict) may still be semantically - # different if they contain $refs that point to defs that were renamed in - # this schema. E.g. both schemas have Wrapper={$ref Inner}, but their - # Inner defs differ, so Schema B's Inner was renamed to Inner_1. The shared - # Wrapper is not actually equal — Schema B needs its own copy with updated refs. - # Loop until stable, since creating a copy can trigger further copies - # in defs that reference it (transitive chains). + # A def equal to an earlier schema's def ("shared") can reuse it only if none of the defs it + # transitively references were renamed; otherwise its `$ref`s resolve differently and it needs + # its own copy. Iterate to a fixpoint before keeping any shared def, so the result does not + # depend on the order of `$defs`. + shared = [name for name in defs if name not in renames and name not in new_names] changed = True while changed: changed = False - for name, def_schema in defs.items(): - if name not in schema_name_mapping: - updated = copy.deepcopy(def_schema) - _update_mapped_json_schema_refs(updated, schema_name_mapping) - if updated != def_schema: - schema_name_mapping[name] = _unique_def_name(name, schema, all_defs) - all_defs[schema_name_mapping[name]] = updated - changed = True - else: - schema_name_mapping[name] = name + for name in shared: + if name in renames: + continue + updated = copy.deepcopy(defs[name]) + _update_mapped_json_schema_refs(updated, renames) + if updated != defs[name]: + renames[name] = _unique_def_name(name, schema, all_defs) + all_defs[renames[name]] = defs[name] + changed = True - # Update refs inside definitions so internal cross-references - # (e.g. Outer referencing Inner which was renamed to Inner_1) are corrected. - for new_name in schema_name_mapping.values(): - _update_mapped_json_schema_refs(all_defs[new_name], schema_name_mapping) + # Rewrite the refs of every def this schema contributes. Shared defs left under their + # name only reference unrenamed defs, so the earlier schema's copy is already correct. + for name in [*new_names, *renames]: + _update_mapped_json_schema_refs(defs[name], renames) - _update_mapped_json_schema_refs(schema, schema_name_mapping) + _update_mapped_json_schema_refs(schema, renames) rewritten_schemas.append(schema) return rewritten_schemas, all_defs diff --git a/tests/test_agent_output_schemas.py b/tests/test_agent_output_schemas.py index 3173d5f29a..769557d4a2 100644 --- a/tests/test_agent_output_schemas.py +++ b/tests/test_agent_output_schemas.py @@ -3,7 +3,7 @@ import json from typing import Annotated, Any, Optional import pytest -from pydantic import AfterValidator, BaseModel, Field +from pydantic import AfterValidator, BaseModel, Field, create_model from typing_extensions import TypeAliasType from pydantic_ai import ( @@ -98,6 +98,124 @@ async def test_auto_output_json_schema(): ) +@pytest.mark.parametrize( + ('names', 'expected'), + [ + pytest.param( + ('AOuter', 'Middle', 'ZLeaf'), + snapshot( + { + 'anyOf': [ + { + 'properties': {'middle': {'$ref': '#/$defs/Middle'}}, + 'required': ['middle'], + 'title': 'AOuter', + 'type': 'object', + }, + { + 'properties': {'middle': {'$ref': '#/$defs/AOuter_Middle_1'}}, + 'required': ['middle'], + 'title': 'AOuter', + 'type': 'object', + }, + ], + '$defs': { + 'Middle': { + 'properties': {'leaf': {'$ref': '#/$defs/ZLeaf'}}, + 'required': ['leaf'], + 'title': 'Middle', + 'type': 'object', + }, + 'ZLeaf': { + 'properties': {'value': {'title': 'Value', 'type': 'string'}}, + 'required': ['value'], + 'title': 'ZLeaf', + 'type': 'object', + }, + 'AOuter_ZLeaf_1': { + 'properties': {'value': {'title': 'Value', 'type': 'integer'}}, + 'required': ['value'], + 'title': 'ZLeaf', + 'type': 'object', + }, + 'AOuter_Middle_1': { + 'properties': {'leaf': {'$ref': '#/$defs/AOuter_ZLeaf_1'}}, + 'required': ['leaf'], + 'title': 'Middle', + 'type': 'object', + }, + }, + } + ), + id='sorted-chain', + ), + pytest.param( + ('A', 'C', 'B'), + snapshot( + { + 'anyOf': [ + { + 'properties': {'middle': {'$ref': '#/$defs/C'}}, + 'required': ['middle'], + 'title': 'A', + 'type': 'object', + }, + { + 'properties': {'middle': {'$ref': '#/$defs/A_C_1'}}, + 'required': ['middle'], + 'title': 'A', + 'type': 'object', + }, + ], + '$defs': { + 'B': { + 'properties': {'value': {'title': 'Value', 'type': 'string'}}, + 'required': ['value'], + 'title': 'B', + 'type': 'object', + }, + 'C': { + 'properties': {'leaf': {'$ref': '#/$defs/B'}}, + 'required': ['leaf'], + 'title': 'C', + 'type': 'object', + }, + 'A_B_1': { + 'properties': {'value': {'title': 'Value', 'type': 'integer'}}, + 'required': ['value'], + 'title': 'B', + 'type': 'object', + }, + 'A_C_1': { + 'properties': {'leaf': {'$ref': '#/$defs/A_B_1'}}, + 'required': ['leaf'], + 'title': 'C', + 'type': 'object', + }, + }, + } + ), + id='out-of-order-chain', + ), + ], +) +def test_output_json_schema_transitive_collision_preserves_branches( + names: tuple[str, str, str], expected: dict[str, Any] +): + """Same-named nested defs with different bodies keep per-branch fidelity.""" + outer_name, middle_name, leaf_name = names + first_leaf = create_model(leaf_name, value=(str, ...)) + first_middle = create_model(middle_name, leaf=(first_leaf, ...)) + first = create_model(outer_name, middle=(first_middle, ...)) + second_leaf = create_model(leaf_name, value=(int, ...)) + second_middle = create_model(middle_name, leaf=(second_leaf, ...)) + second = create_model(outer_name, middle=(second_middle, ...)) + + schema = Agent(TestModel(), output_type=[first, second]).output_json_schema() + + assert schema == expected + + async def test_tool_output_json_schema(): # one output agent = Agent( diff --git a/tests/test_utils.py b/tests/test_utils.py index 69a5a6e92b..46447fabaa 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -3,8 +3,10 @@ from __future__ import annotations as _annotations import asyncio import contextlib import contextvars +import copy import functools import importlib.util +import itertools import os import sys import threading @@ -13,7 +15,7 @@ from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field from importlib.metadata import distributions from types import ModuleType -from typing import Any +from typing import Any, Literal, cast import anyio import pytest @@ -30,6 +32,7 @@ from pydantic_ai._utils import ( get_first_param_type, group_by_temporal, is_async_callable, + is_str_dict, merge_json_schema_defs, replace_no_init, run_in_executor, @@ -1133,6 +1136,147 @@ def test_merge_json_schema_defs_structurally_equal_with_different_ref_targets(): ) +def _chain_schema( + names: tuple[str, str, str, str], + value_type: str, + *, + defs_layout: Literal['sorted', 'chain'], + title: str, +) -> dict[str, Any]: + d1, d2, d3, leaf = names + chain_defs = { + d1: {'type': 'object', 'properties': {'next': {'$ref': f'#/$defs/{d2}'}}}, + d2: {'type': 'object', 'properties': {'next': {'$ref': f'#/$defs/{d3}'}}}, + d3: {'type': 'object', 'properties': {'next': {'$ref': f'#/$defs/{leaf}'}}}, + leaf: {'type': 'object', 'properties': {'value': {'type': value_type}}}, + } + if defs_layout == 'sorted': + chain_defs = {name: chain_defs[name] for name in sorted(chain_defs)} + + return { + '$defs': chain_defs, + 'title': title, + 'type': 'object', + 'properties': {'root': {'$ref': f'#/$defs/{d1}'}}, + } + + +def _assert_all_json_schema_refs_resolve(value: Any, defs: dict[str, dict[str, Any]]) -> None: + if is_str_dict(value): + if ref := value.get('$ref'): + assert str(ref).removeprefix('#/$defs/') in defs + for nested in value.values(): + _assert_all_json_schema_refs_resolve(nested, defs) + elif isinstance(value, list): + for nested in cast(list[Any], value): + _assert_all_json_schema_refs_resolve(nested, defs) + + +def _resolve_branch_leaf_value_type(branch: dict[str, Any], defs: dict[str, dict[str, Any]]) -> str: + """Follow `$ref`s from a merged branch down to its leaf model's `value` field.""" + node: dict[str, Any] = branch + while True: + props: dict[str, Any] = node.get('properties', {}) + next_props = [prop for prop in props.values() if '$ref' in prop] + if not next_props: + return str(node['properties']['value']['type']) + node = defs[str(next_props[0]['$ref']).removeprefix('#/$defs/')] + + +@pytest.mark.parametrize('names', list(itertools.permutations(('A', 'B', 'C', 'D')))) +@pytest.mark.parametrize('defs_layout', ['sorted', 'chain']) +def test_merge_json_schema_defs_transitive_rename_all_orderings( + names: tuple[str, str, str, str], defs_layout: Literal['sorted', 'chain'] +) -> None: + """Every 4-level chain ordering keeps branch types; too many cases to run through `Agent.output_json_schema()`.""" + schemas = [ + _chain_schema(names, 'string', defs_layout=defs_layout, title='StringRoot'), + _chain_schema(names, 'integer', defs_layout=defs_layout, title='IntegerRoot'), + ] + original_schemas = copy.deepcopy(schemas) + + rewritten_schemas, all_defs = merge_json_schema_defs(schemas) + + assert _resolve_branch_leaf_value_type(rewritten_schemas[0], all_defs) == 'string' + assert _resolve_branch_leaf_value_type(rewritten_schemas[1], all_defs) == 'integer' + _assert_all_json_schema_refs_resolve(rewritten_schemas, all_defs) + _assert_all_json_schema_refs_resolve(all_defs, all_defs) + assert schemas == original_schemas + + +def test_merge_json_schema_defs_transitive_rename_snapshot(): + """Pins the names for one previously broken ordering; the Agent-level test covers the user-facing path.""" + names = ('A', 'C', 'B', 'D') + schemas = [ + _chain_schema(names, 'string', defs_layout='sorted', title='StringRoot'), + _chain_schema(names, 'integer', defs_layout='sorted', title='IntegerRoot'), + ] + + rewritten_schemas, all_defs = merge_json_schema_defs(schemas) + + assert (all_defs, rewritten_schemas) == snapshot( + ( + { + 'A': {'type': 'object', 'properties': {'next': {'$ref': '#/$defs/C'}}}, + 'B': {'type': 'object', 'properties': {'next': {'$ref': '#/$defs/D'}}}, + 'C': {'type': 'object', 'properties': {'next': {'$ref': '#/$defs/B'}}}, + 'D': {'type': 'object', 'properties': {'value': {'type': 'string'}}}, + 'IntegerRoot_D_1': {'type': 'object', 'properties': {'value': {'type': 'integer'}}}, + 'IntegerRoot_B_1': {'type': 'object', 'properties': {'next': {'$ref': '#/$defs/IntegerRoot_D_1'}}}, + 'IntegerRoot_C_1': {'type': 'object', 'properties': {'next': {'$ref': '#/$defs/IntegerRoot_B_1'}}}, + 'IntegerRoot_A_1': {'type': 'object', 'properties': {'next': {'$ref': '#/$defs/IntegerRoot_C_1'}}}, + }, + [ + {'title': 'StringRoot', 'type': 'object', 'properties': {'root': {'$ref': '#/$defs/A'}}}, + {'title': 'IntegerRoot', 'type': 'object', 'properties': {'root': {'$ref': '#/$defs/IntegerRoot_A_1'}}}, + ], + ) + ) + + +def test_merge_json_schema_defs_shared_chain_not_renamed_for_unrelated_collision(): + """Shared defs that do not reach a renamed def are reused, not copied; a merge-level detail of the def names.""" + shared_defs = { + 'SharedRoot': {'type': 'object', 'properties': {'leaf': {'$ref': '#/$defs/SharedLeaf'}}}, + 'SharedLeaf': {'type': 'object', 'properties': {'value': {'type': 'boolean'}}}, + } + schemas = [ + { + '$defs': { + **shared_defs, + 'Collision': {'type': 'object', 'properties': {'value': {'type': 'string'}}}, + }, + 'title': 'First', + 'properties': {'shared': {'$ref': '#/$defs/SharedRoot'}}, + }, + { + '$defs': { + **shared_defs, + 'Collision': {'type': 'object', 'properties': {'value': {'type': 'integer'}}}, + }, + 'title': 'Second', + 'properties': {'shared': {'$ref': '#/$defs/SharedRoot'}}, + }, + ] + + rewritten_schemas, all_defs = merge_json_schema_defs(schemas) + + assert (all_defs, rewritten_schemas) == snapshot( + ( + { + 'SharedRoot': {'type': 'object', 'properties': {'leaf': {'$ref': '#/$defs/SharedLeaf'}}}, + 'SharedLeaf': {'type': 'object', 'properties': {'value': {'type': 'boolean'}}}, + 'Collision': {'type': 'object', 'properties': {'value': {'type': 'string'}}}, + 'Second_Collision_1': {'type': 'object', 'properties': {'value': {'type': 'integer'}}}, + }, + [ + {'title': 'First', 'properties': {'shared': {'$ref': '#/$defs/SharedRoot'}}}, + {'title': 'Second', 'properties': {'shared': {'$ref': '#/$defs/SharedRoot'}}}, + ], + ) + ) + + def test_strip_markdown_fences(): assert strip_markdown_fences('{"foo": "bar"}') == '{"foo": "bar"}' assert strip_markdown_fences('```json\n{"foo": "bar"}\n```') == '{"foo": "bar"}'