How to change create_agent state from middleware during tool execution

OSS Product Help LangChain
1

How to change create_agent state from middleware during tool execution?
I have create_agent instance with CustomState and I want to change state right after tool is executed.

Changing CustomState @before_model isn’t an option because tools can be called in a row and CustomState won’t be changed when 2nd tool is called


from langchain.agents import create_agent
from langchain.agents.middleware import wrap_tool_call, AgentState
from langchain.tools.tool_node import ToolCallRequest
from langchain.messages import ToolMessage
from langgraph.types import Command
from langgraph.runtime import Runtime
from typing import Any, Callable

class CustomState(AgentState):
    user_preferences: dict

@wrap_tool_call
def change_agent_state(
    request: ToolCallRequest,
    handler: Callable[[ToolCallRequest], ToolMessage | Command],
):  
    result = handler(request)
    #change CustomState here
    return result

agent = create_agent(
    model,
    tools=[tool1, tool2],
    middleware=[change_agent_state],
    state_schema=CustomState
)
2

hi @J3zus

AFAIK, you can’t reliably update CustomState by mutating request.state inside wrap_tool_call. Instead, you need to return a Command(update=...) that includes both the tool’s ToolMessage and your state changes.
That Command is what create_agent/LangGraph actually uses to persist state updates between steps.

Simply mutating request.state (or request.runtime.state) does not tell the graph to commit that change.

The only supported way for a tool/tool-middleware to update graph state is to return a Command whose update field describes the new state.

Try sth like this:

from typing import Callable

from langchain.agents.middleware import wrap_tool_call, AgentState
from langchain.tools.tool_node import ToolCallRequest
from langchain.messages import ToolMessage
from langgraph.types import Command


class CustomState(AgentState):
    user_preferences: dict


@wrap_tool_call
def change_agent_state(
    request: ToolCallRequest,
    handler: Callable[[ToolCallRequest], ToolMessage | Command],
) -> ToolMessage | Command:
    # Execute the tool (and any inner middleware)
    result = handler(request)

    # Compute the new preferences you want to store
    prefs = dict(request.state.get("user_preferences") or {})
    prefs["last_tool"] = request.tool_call["name"]

    # If the tool already returned a Command, merge our update
    if isinstance(result, Command):
        update = result.update or {}
        if isinstance(update, dict):
            merged = {**update, "user_preferences": prefs}
            return Command(
                graph=result.graph,
                update=merged,
                resume=result.resume,
                goto=result.goto,
            )
        # Non-dict updates are uncommon here; fall back to original
        return result

    # For a plain ToolMessage, wrap it in a Command that updates both messages and state
    assert isinstance(result, ToolMessage)
    return Command(
        update={
            "messages": [result],           # required so ToolNode sees the tool output
            "user_preferences": prefs,      # your CustomState field
        }
    )

agent = create_agent(
    model,
    tools=[tool1, tool2],
    middleware=[change_agent_state],
    state_schema=CustomState,
)

Now, after each tool call, the underlying graph state will have an updated user_preferences field.

You can access it from the final result (e.g. result["user_preferences"]) or from subsequent middleware hooks (before_model, after_model, etc.).