`middleware` and Custom `state_schema` are mutually exclusive in `create_agent()`

[jmoreno11]

Checked other resources

• This is a feature request, not a bug report or usage question.
• I added a clear and descriptive title that summarizes the feature request.
• I used the GitHub search to find a similar feature request and didn't find it.
• I checked the LangChain documentation and API reference to see if this feature already exists.
• This is not related to the langchain-community package.

Feature Description

When using create_agent(), the middleware parameter and state_schema parameter cannot be used together. This creates a significant limitation for RAG applications that need both:

1. Custom state - to store retrieved documents for citation purposes
2. Middleware - for message summarization to manage context window

The assertion at line 1150 in langchain/agents/react_agent.py enforces this:

assert state_schema is None  # noqa: S101

Use Case

Building a RAG agent that needs to:

• Retrieve documents from a vector store
• Provide citations by accessing the retrieved documents after agent execution
• Manage conversation history with summarization to fit within context windows

Proposed Solution

1. Design Rationale: Why are middleware and state_schema mutually exclusive? Is there a technical limitation or just an implementation detail?

2. Feature Request: Could middleware support custom state schemas? For example:

   agent = create_agent(
       model=llm,
       tools=tools,
       middleware=[SummarizationMiddleware(model=llm)],
       state_schema=CustomState,  # Should work together
       context_schema=Context,
   )

Alternatives Considered

1. Tool Message Artifacts (as suggested in docs)

@tool(response_format="content_and_artifact")
def retrieve_context(query: str):
    retrieved_docs = vector_store.similarity_search(query, k=2)
    serialized = "\n\n".join(...)
    return serialized, retrieved_docs

Issue: When multiple tools exist, you must navigate the message history to find the correct ToolMessage with the retrieval artifacts. This becomes fragile and complex.

2. Custom State via Command

def search_knowledge_base(
    state: Annotated[RAGState, InjectedState],
    tool_call_id: Annotated[str, InjectedToolCallId],
) -> Command:
    docs = await retrieve_documents(...)
    return Command(update={
        "retrieved_documents": docs,
        "messages": [ToolMessage(content, tool_call_id=tool_call_id)]
    })

Issue: This requires state_schema=RAGState in create_agent(), which is incompatible with middleware=[SummarizationMiddleware(...)].

3. Post-Model Hook

Could potentially attach documents to message metadata, but still doesn't solve the core limitation.

Additional Context

This limitation affects any RAG application that needs:

• Clean access to retrieved documents for citations
• Conversation history management
• Multiple tools (where artifact navigation becomes complex)

Would appreciate guidance on:

• The reasoning behind this design decision
• Any plans to support both features together
• Recommended patterns for this common RAG use case

[Rohanvasudev1]

Hey could I help with this feature? I already gave it a go on my pull request. Would be happy to collaborate too

[jmoreno11]

I also found this in the docs for state extension: https://docs.langchain.com/oss/python/langchain/middleware#state-extension

which I can apply to my use case:

# utils.py

from langchain.agents.middleware import AgentMiddleware, AgentState
from langchain_core.documents import Document


class RAGState(AgentState):
    """Extended state for RAG agent with retrieved documents."""

    retrieved_documents: list[Document]


class CitationMiddleware(AgentMiddleware):
    """Citation middleware that extends state with RAGState schema."""

    state_schema = RAGState

# tools.py
...

@tool
async def search_knowledge_base(
    query: str,
    config: RunnableConfig,
    tool_call_id: Annotated[str, InjectedToolCallId],
) -> Command:
    """Search the knowledge base.

    Args:
        query: Optimized query to be used for retrieving documents
        state: Agent state
        config: RunnableConfig to configure the retriever (top_k, score_threshold)
        tool_call_id: Tool call ID

    Returns:
        Updated state with retrieved documents and tool message
    """
    ...
    return Command(
        update={
            "retrieved_documents": top_k_docs,
            # update the message history
            "messages": [ToolMessage(content, tool_call_id=tool_call_id)],
        }
    )

# agent.py
...
agent = create_agent(
            name=self.name,
            model=self.llm,
            tools=self.tools,
            prompt=self.system_prompt,
            middleware=[CitationMiddleware(), SummarizationMiddleware(model=self.llm)],
            context_schema=Context,
            checkpointer=self.checkpointer,
        )

I still think it's a bit of an overkill to have a CustomMiddleware object just to extend the state.

I think better and a more unified documentation around State Management is needed here, as well as maybe an OutputSchema (!= Structured Output if we store things like filenames)

[eyurtsev]

For now we recommend using custom middleware this purpose. We'll publish documentation with migration notes. Should be able to feed this into an llm / claude code to get your code migrated for you.

[jmoreno11]

Hey thanks Eugene Yurtsev (@eyurtsev) for the response. Please also note that when following this approach, there is a risk of the extended state keys becoming outdated, example:

# agent with custom middleware that adds a state key `retrieved_documents`
>>> response = await agent.ainvoke({"messages": [query_1]}, config={"configurable": {"thread_id": "xxx", ...}}, context=context)

{'messages': [...], 'retrieved_documents': [Doc1, Doc2]} <--- relevant docs after tool use

>>> response = await agent.ainvoke({"messages": [follow_up_query]}, config={"configurable": {"thread_id": "xxx", ...}}, context=context)

{'messages': [...], 'retrieved_documents': [Doc1, Doc2]} <--- no tool use!! outdated docs!!!

Any suggestions with this?

[eyurtsev]

Jose Moreno Ortega (@jmoreno11) this is working as expected due to how langgraph persistence works. It will persist state updates for the entire conversation thread.

If you want to reset the state, you can add an explicit before agent hook that sets retrieved documents to an empty list or you can start a conversation on a new thread (what you want depends on your use case)

[eyurtsev]

Agent interface also allows passing state_schema directly now to make it more convenient to use.

Closing this issue (as there's no bug to fix).