How to make tools store custom variables in Agent's state when using create_deep_agent()?

[Sai Akhilesh Ande (akhil189)]

I have a workflow with 3 tools (SQL tool, Retriever Tool, Report Tool).

Retriever Tool - retrieves context from a vectorstore (List[Document])
SQL Tool - creates a valid SQL query
Report Tool - accepts a SQL query and creates a csv report

Based on user query, the agent has to call one or more tools. The workflow is specified in skills. I need my tools to

1. Store the context in agent's state and
2. Return a ToolMessage to the main agent

I need them in agent's state because:

1. I need the context for highlighting (grounding) on the UI and also logging.
2. If there a filepath, I want to provide a download option on the UI.

from langchain.tools import ToolRuntime, tool
from langgraph.types import Command
from langchain_core.messages import ToolMessage
from langchain.agents import AgentState, create_agent
from typing import Annotated, List
from langchain_core.documents import Document


@tool
def generate_report(self, sql_query: str, runtime: ToolRuntime) -> Command:
    """Takes in a SQL query and generates a csv file with postprocessing

    Args:
        sql_query: A valid SQL query
        runtime: LangChain ToolRuntime providing the tool_call_id for ToolMessage.

    Returns:
        Command updating the agent state with the saved document and a ToolMessage.
    """

    # <Logic to create the csv file>
    path = ...

    return Command(
        update={
            "reports": [path],
            "messages": [
                ToolMessage(
                    content=f"report successfully generated",
                    tool_call_id=runtime.tool_call_id,
                )
            ],
        }
    )

@tool
def context_retriever(query: str, runtime: ToolRuntime) -> Command:
    """Retrieves context from a vectorstore

    Args:
        query: The query to find relevant context.
    """

    # logic to retrieve Documents from a vectorstore
    docs = ...
    formatted = "\n".join([doc.content for doc in docs])

    return Command(
        update={
            "agent_context": docs,
            "messages": [
                ToolMessage(
                    content=formatted,
                    tool_call_id=runtime.tool_call_id,
                )
            ],
        }
    )

class CustomState(AgentState):
    """Extended state that tracks context alongside messages."""

    agent_context: List[Document] # docs from vectorstore
    reports = List[str] # filepaths

### with create_agent()
agent = create_agent(
    model=...,
    system_prompt=...,
    tools=[...],
    state_schema=CustomState
)

### with deep_agent()
agent = create_deep_agent(
    model=...,
    system_prompt=...,
    tools=[...],
    state_schema=CustomState # not accepted
)

How to achieve the same with create_deep_agent()?