mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-17 21:25:46 +02:00
f4aee546ad
## Summary When `ToolNode` receives a bare `[tool_call]` list via the Send API (the dispatch shape `create_agent` will use once langchain-ai/langchain#36960 lands), hydrate `ToolRuntime.state` from the current channel values instead of requiring the dispatcher to inline the full agent state dict into every `Send.arg`. Motivation: the paired langchain PR drops the `ToolCallWithContext` wrapper from `create_agent`'s tool dispatch, which eliminates an O(N²) storage term on `__pregel_tasks` checkpoint writes. Without this companion change there would be no path for the tool node to see the graph state. ## What changed - `libs/prebuilt/langgraph/prebuilt/tool_node.py` — `_extract_state` grows a third branch for list-form input. When the input is a list whose last entry is a `ToolCall` dict, read the current channel values via `CONFIG_KEY_READ` and return them as the state dict. The full new logic is four lines inline in `_extract_state`: ```python read = config.get(CONF, {}).get(CONFIG_KEY_READ) if read is None: return {} # Pregel installs CONFIG_KEY_READ as # `functools.partial(local_read, scratchpad, channels, managed, task)`. channels = read.args[1] return cast("dict[str, Any]", read(list(channels), False)) ``` - No changes to the pregel read machinery (`local_read`, `ChannelRead`). - Only channel values are read; managed values have their own injection path (`ToolRuntime.context`, `InjectedContext`) and were never in the pre-fix inlined state dict, so we don't add them here. - Falls back to `{}` when invoked outside a Pregel context (e.g. direct `ToolNode(...).invoke([tool_call])` from a test harness), which preserves existing `ToolNode` direct-invocation test behavior. - `libs/prebuilt/tests/test_on_tool_call.py` — two new tests covering the list-form hydration path (sync + async). They build a `functools.partial` that matches Pregel's real `CONFIG_KEY_READ` shape and assert `ToolRuntime.state` reflects the current channel values. ## Why it's safe - **Same snapshot semantics as before.** `Send` is emitted at end-of-super-step-N; consumed at start-of-super-step-N+1. Channels at that point reflect every write from super-step N (including the new AIMessage the tool calls originated from). Parallel tool tasks in the tools super-step all read the same values since sibling writes don't land until end-of-super-step. - **Legacy `ToolCallWithContext` path preserved.** External dispatchers that still inline state continue to work unchanged — `_extract_state` checks that branch first. ## Test plan - [x] `make test` in `libs/prebuilt` — **204 pass** - [x] Two new hydration tests (sync + async) green - [x] `make format` / `make lint` / `mypy` clean --------- Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>