State Loss in Hierarchical Multi-Agent System with Deep Agents and Custom AgentState

OSS Product Help LangGraph
5

Hi @Pradeep ran into the same problem.

Custom Middleware to Surface State Schema

By default, create_deep_agent only exposes context_schema, which, as mentioned above, is immutable. I discovered a workaround by creating a lightweight middleware that extends AgentMiddleware and sets state_schema on itself. When create_agent (called internally by create_deep_agent) builds the final schema, it merges state_schema from all middleware so your custom fields appear as top-level state channels.

Here’s a minimal example:

class StateSchemaMiddleware(AgentMiddleware):
    tools = ()

    def __init__(self, schema):
        super().__init__()
        self.state_schema = schema

my_agent = create_deep_agent(
    model=llm,
    system_prompt=prompt,
    context_schema=MyState,
    middleware=[StateSchemaMiddleware(MyState)],
    tools=[...],
)

I have found this makes the additional fields accessible as part of the agent’s state, even when using create_deep_agent.