feat: support state updates from `wrap_model_call` with command(s) (#35033)
Alternative to https://github.com/langchain-ai/langchain/pull/35024.
Paving the way for summarization in `wrap_model_call` (which requires
state updates).
---
Add `ExtendedModelResponse` dataclass that allows `wrap_model_call`
middleware to return a `Command` alongside the model response for
additional state updates.
```py
@dataclass
class ExtendedModelResponse(Generic[ResponseT]):
model_response: ModelResponse[ResponseT]
command: Command
```
## Motivation
Previously, `wrap_model_call` middleware could only return a
`ModelResponse` or `AIMessage` — there was no way to inject additional
state updates (e.g. custom state fields) from the model call middleware
layer. `ExtendedModelResponse` fills this gap by accepting an optional
`Command`.
This feature is needed by the summarization middleware, which needs to
track summarization trigger points calculated during `wrap_model_call`.
## Why `Command` instead of a plain `state_update` dict?
We chose `Command` rather than the raw `state_update: dict` approach
from the earlier iteration because `Command` is the established
LangGraph primitive for state updates from nodes. Using `Command` means:
- State updates flow through the graph's reducers (e.g. `add_messages`)
rather than being merged as raw dicts. This makes messages updates
additive alongside the model response instead of replacing them.
- Consistency with `wrap_tool_call`, which already returns `Command`.
- Future-proof: as `Command` gains new capabilities (e.g. `goto`,
`send`), middleware can leverage them without API changes.
## Why keep `model_response` separate instead of using `Command`
directly?
The model node needs to distinguish the model's actual response
(messages + structured output) from supplementary middleware state
updates. If middleware returned only a `Command`, there would be no
clean way to extract the `ModelResponse` for structured output handling,
response validation, and the core model-to-tools routing logic. Keeping
`model_response` explicit preserves a clear boundary between "what the
model said" and "what middleware wants to update."
Also, in order to avoid breaking, the `handler` passed to
`wrap_tool_call` needs to always return a `ModelResponse`. There's no
easy way to preserve this if we pump it into a `Command`.
One nice thing about having this `ExtendedModelResponse` structure is
that it's extensible if we want to add more metadata in the future.
## Composition
When multiple middleware layers return `ExtendedModelResponse`, their
commands compose naturally:
- **Inner commands propagate outward:** At composition boundaries,
`ExtendedModelResponse` is unwrapped to its underlying `ModelResponse`
so outer middleware always sees a plain `ModelResponse` from
`handler()`. The inner command is captured and accumulated.
- **Commands are applied through reducers:** Each `Command` becomes a
separate state update applied through the graph's reducers. For
messages, this means they're additive (via `add_messages`), not
replacing.
- **Outer wins on conflicts:** For non-reducer state fields, commands
are applied inner-first then outer, so the outermost middleware's value
takes precedence on conflicting keys.
- **Retry-safe:** When outer middleware retries by calling `handler()`
again, accumulated inner commands are cleared and re-collected from the
fresh call.
```python
class Outer(AgentMiddleware):
def wrap_model_call(self, request, handler):
response = handler(request) # sees ModelResponse, not ExtendedModelResponse
return ExtendedModelResponse(
model_response=response,
command=Command(update={"outer_key": "val"}),
)
class Inner(AgentMiddleware):
def wrap_model_call(self, request, handler):
response = handler(request)
return ExtendedModelResponse(
model_response=response,
command=Command(update={"inner_key": "val"}),
)
# Final state merges both commands: {"inner_key": "val", "outer_key": "val"}
```
## Backwards compatibility
Fully backwards compatible. The `ModelCallResult` type alias is widened
from `ModelResponse | AIMessage` to `ModelResponse | AIMessage |
ExtendedModelResponse`, but existing middleware returning
`ModelResponse` or `AIMessage` continues to work identically.
## Internals
- `model_node` / `amodel_node` now return `list[Command]` instead of
`dict[str, Any]`
- `_build_commands` converts the model response + accumulated middleware
commands into a list of `Command` objects for LangGraph
- `_ComposedExtendedModelResponse` is the internal type that accumulates
commands across layers during composition S
Sydney Runkle committed
8767a462ca3e93873c1002eb3ee2bf669cee9974
Parent: 273d282
Committed by GitHub <[email protected]>
on 2/6/2026, 12:28:04 PM