From 294f33253109b92bbccff8b69295e71bc0c3eac9 Mon Sep 17 00:00:00 2001 From: Sydney Runkle <54324534+sydney-runkle@users.noreply.github.com> Date: Wed, 6 May 2026 14:38:04 -0400 Subject: [PATCH] =?UTF-8?q?fix(langchain):=20ordered=20schema=20resolution?= =?UTF-8?q?=20=E2=80=94=20list=20replaces=20set=20so=20state=5Fschema=20wi?= =?UTF-8?q?ns=20(#37223)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Problem `_resolve_schemas` collected middleware state schemas into a `set[type]`, making field-conflict resolution non-deterministic. When two schemas declare the same field with different annotations (e.g. `DeltaChannel` vs `BinaryOperatorAggregate` on `messages`), the winner depended on set iteration order — effectively random across runs. This forced callers to work around it with post-compilation patches like: ```python if not isinstance(agent.channels.get("messages"), DeltaChannel): agent.channels["messages"] = DeltaChannel(...) ``` ## Fix Change `_resolve_schemas` to accept `list[type]` and build it as: ```python state_schemas = [*(m.state_schema for m in middleware), base_state] ``` `base_state` (the caller's explicit `state_schema`, or `AgentState` by default) goes **last** and wins any conflict. Middleware schemas precede it in registration order. This means callers can now pass `state_schema=MyCustomState` and be guaranteed it overrides any conflicting middleware annotation — no post-compilation patching needed. ## Test plan - All 9 existing `test_state_schema.py` tests pass - Updated the two direct `_resolve_schemas({...})` call sites in tests to use lists --- libs/langchain_v1/langchain/agents/factory.py | 18 ++++++++--- .../unit_tests/agents/test_state_schema.py | 31 ++++++++++++++++--- 2 files changed, 39 insertions(+), 10 deletions(-) diff --git a/libs/langchain_v1/langchain/agents/factory.py b/libs/langchain_v1/langchain/agents/factory.py index e4fa8273ddc..b15e9caf1be 100644 --- a/libs/langchain_v1/langchain/agents/factory.py +++ b/libs/langchain_v1/langchain/agents/factory.py @@ -407,8 +407,13 @@ def _get_schema_type_hints(schema: type) -> dict[str, Any]: return get_type_hints(schema, include_extras=True) -def _resolve_schemas(schemas: set[type]) -> tuple[type, type, type]: - """Resolve state, input, and output schemas for the given schemas.""" +def _resolve_schemas(schemas: list[type]) -> tuple[type, type, type]: + """Resolve state, input, and output schemas for the given schemas. + + Schemas are merged in list order; later entries override earlier ones when the + same field is declared by multiple schemas. Duplicates are harmless — a type + that appears more than once is processed at its last position. + """ schema_hints = {schema: _get_schema_type_hints(schema) for schema in schemas} return ( _resolve_schema(schema_hints, "StateSchema", None), @@ -1030,10 +1035,13 @@ def create_agent( ] awrap_model_call_handler = _chain_async_model_call_handlers(async_handlers) - state_schemas: set[type] = {m.state_schema for m in middleware} - # Use provided state_schema if available, otherwise use base AgentState base_state = state_schema if state_schema is not None else AgentState - state_schemas.add(base_state) + # Build an ordered list: middleware schemas first (in registration order), + # base_state last so it wins any field conflict. This lets the caller's + # explicit state_schema override middleware annotations — e.g. passing + # a DeltaChannel-annotated schema wins over BinaryOperatorAggregate from + # AgentState without requiring a post-compilation patch. + state_schemas: list[type] = [*(m.state_schema for m in middleware), base_state] resolved_state_schema, input_schema, output_schema = _resolve_schemas(state_schemas) diff --git a/libs/langchain_v1/tests/unit_tests/agents/test_state_schema.py b/libs/langchain_v1/tests/unit_tests/agents/test_state_schema.py index cbd66d4a0a3..26d8dbe36b8 100644 --- a/libs/langchain_v1/tests/unit_tests/agents/test_state_schema.py +++ b/libs/langchain_v1/tests/unit_tests/agents/test_state_schema.py @@ -6,7 +6,7 @@ AgentState without needing to create custom middleware. from __future__ import annotations -from typing import TYPE_CHECKING, Annotated, Any +from typing import TYPE_CHECKING, Annotated, Any, get_args, get_type_hints from langchain_core.messages import HumanMessage from langchain_core.tools import tool @@ -256,6 +256,27 @@ def test_state_schema_with_private_state_field() -> None: assert len(result["messages"]) == 4 # Human, AI (tool call), Tool result, AI (final) +def test_last_schema_wins_for_conflicting_field() -> None: + """The last schema in the list wins when multiple schemas declare the same field. + + In practice this means state_schema (passed last) overrides middleware annotations, + and a later middleware overrides an earlier one. + """ + + class FirstState(AgentState[Any]): + shared_field: Annotated[str, "first"] + + class MiddleState(AgentState[Any]): + shared_field: Annotated[str, "middle"] + + class LastState(AgentState[Any]): + shared_field: Annotated[str, "last"] + + resolved, _, _ = factory._resolve_schemas([FirstState, MiddleState, LastState]) + hints = get_type_hints(resolved, include_extras=True) + assert get_args(hints["shared_field"])[1] == "last" + + def test_get_schema_type_hints_cache_hits_for_reused_schema() -> None: """Test repeated schema resolution reuses cached type hints for the same schema.""" @@ -266,9 +287,9 @@ def test_get_schema_type_hints_cache_hits_for_reused_schema() -> None: factory._get_schema_type_hints.cache_clear() - factory._resolve_schemas({CachedState}) + factory._resolve_schemas([CachedState]) first_info = factory._get_schema_type_hints.cache_info() - factory._resolve_schemas({CachedState}) + factory._resolve_schemas([CachedState]) second_info = factory._get_schema_type_hints.cache_info() assert first_info.misses == 1 @@ -293,9 +314,9 @@ def test_get_schema_type_hints_cache_accepts_distinct_local_schema_types() -> No schema_a = make_state_schema("LocalStateA") schema_b = make_state_schema("LocalStateB") - factory._resolve_schemas({schema_a, schema_b}) + factory._resolve_schemas([schema_a, schema_b]) first_info = factory._get_schema_type_hints.cache_info() - factory._resolve_schemas({schema_a, schema_b}) + factory._resolve_schemas([schema_a, schema_b]) second_info = factory._get_schema_type_hints.cache_info() assert first_info.misses == 2