fix(langchain): ordered schema resolution — list replaces set so state_schema wins (#37223)

## 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
This commit is contained in:
Sydney Runkle
2026-05-06 14:38:04 -04:00
committed by GitHub
parent 5161c4ecbe
commit 294f332531
2 changed files with 39 additions and 10 deletions

View File

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

View File

@@ -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