Skip to main content

Custom Agent State: Reading and Writing Beyond Messages

Context, from Article 7, is fixed for an entire call — perfect for "who is this user," useless for anything that needs to change during a conversation. Suppose Julie says "I'm replying to Jane" early in a conversation, and three messages later asks Aria to "send it" — Aria needs to remember which email "it" refers to, and that information didn't exist when the conversation started. It came from the conversation itself.

This article closes that gap with custom agent state: a way to give tools a shared scratchpad they can both read from and write to as a conversation unfolds. And it'll clear up something that's been slightly simplified until now — what "memory" from Article 4 actually is, underneath.

🟡 Skill level: Intermediate.

Quick Reference​

When to use this: Whenever a tool needs to save information during a conversation that another tool (or a later turn) needs to read — beyond just the message history itself.

Basic syntax:

from langchain.agents import AgentState
from langgraph.types import Command
from langchain.messages import ToolMessage

class MyState(AgentState):
    current_recipient: str

@tool
def set_recipient(name: str, runtime: ToolRuntime) -> Command:
    """Set the current email recipient."""
    return Command(update={
        "current_recipient": name,
        "messages": [ToolMessage("Recipient set", tool_call_id=runtime.tool_call_id)],
    })

Common patterns:

  • Custom state extends AgentState with your own named fields
  • Tools write to state by returning a Command(update={...})
  • Tools read state through runtime.state["field_name"]
  • State requires a checkpointer + thread_id to persist across separate invoke() calls — same mechanism as memory

Gotchas:

  • ⚠️ State and memory aren't separate systems — the conversation history you persisted in Article 4 is itself just one field (messages) inside the same state object you're now extending.
  • ⚠️ Reading a state field that hasn't been set yet raises a KeyError — always handle this gracefully.

See also: Memory and Threads: Agents That Remember, Runtime Context: Injecting User-Specific Data

What You Need to Know First​

What We'll Cover in This Article​

  • What state actually is, and how it relates to the memory you already learned about
  • How to define a custom state schema
  • How to write to state from inside a tool
  • How to read state from inside a tool

What We'll Explain Along the Way​

  • AgentState, the base class everything builds on
  • Command, the mechanism tools use to update state

What Is State, Really? (And How Does It Relate to Memory?)​

Here's something worth clearing up directly: in Article 4, we talked about "memory" as if it were its own separate system. It isn't, exactly. What's actually being saved and reloaded by a checkpointer is something called state — and the conversation history (the messages list) is just one field inside that state. "Memory" was really just a friendly name for "the messages field of state, persisted across calls."

Once you see it this way, custom state becomes a small, natural extension: instead of state only containing messages, we add our own additional fields to the same object — and they get saved and reloaded by the exact same checkpointer, under the exact same thread_id, automatically.

Diagram: "Memory" was always just the messages field of a larger state object. Custom state means adding more fields to that same object — and they're saved and reloaded together, by the same checkpointer.

This also clears up the difference from context (Article 7): context is fixed, set once by your code at the start of a call, and read-only from a tool's perspective. State is mutable — tools can actively write to it, and those changes persist for the rest of the conversation, the same way the message history does.

Defining a Custom State Schema​

Let's give Aria a way to remember who the current email recipient is, partway through a conversation. We extend AgentState, the base state class every agent already uses, with our own field:

# Purpose: Extend the agent's state with a custom field
# Context: Adds to the same state object that already holds "messages"
# Input: N/A — this defines a shape, not a value yet
# Output: A reusable state class with one extra field beyond the defaults

from langchain.agents import AgentState

class RecipientState(AgentState):
    current_recipient: str

Just like memory in Article 4, custom state needs a checkpointer to actually persist across separate invoke() calls:

# Purpose: Configure an agent to use the custom state schema, with persistence
# Context: state_schema works alongside checkpointer, not instead of it
# Input: N/A
# Output: An agent instance aware of the extra current_recipient field

from langchain.agents import create_agent
from langgraph.checkpoint.memory import InMemorySaver

agent = create_agent(
    model="gpt-5-nano",
    state_schema=RecipientState,
    checkpointer=InMemorySaver(),
)

Writing to State From a Tool​

To update state, a tool returns a special object called Command, instead of just returning a plain value. This tells LangGraph: "update these specific state fields, then continue."

# Purpose: A tool that records the current recipient in state
# Context: Once Julie says who she's replying to, this gets remembered
# Input: A recipient's name
# Output: A Command updating state, plus a confirmation message

from langchain.tools import tool, ToolRuntime
from langgraph.types import Command
from langchain.messages import ToolMessage

@tool
def set_current_recipient(name: str, runtime: ToolRuntime) -> Command:
    """Record who the current email recipient is, once known."""
    return Command(update={
        "current_recipient": name,
        "messages": [
            ToolMessage(f"Recipient set to {name}", tool_call_id=runtime.tool_call_id)
        ],
    })

Two things worth noting here. First, Command(update={...}) can update multiple state fields at once — here, both current_recipient and messages (we still need to add a ToolMessage confirming what happened, the same as any tool result). Second, runtime.tool_call_id connects this result back to the specific tool call that triggered it — required for the message history to stay coherent.

Reading State From a Tool​

Now let's add a tool that reads the recipient back out, later in the conversation:

# Purpose: A tool that reads the current recipient back out of state
# Context: Lets a later tool call use information set earlier in the conversation
# Input: None directly — reads from runtime.state
# Output: The current recipient's name, or a clear fallback if unset

from langchain.tools import tool, ToolRuntime

@tool
def get_current_recipient(runtime: ToolRuntime) -> str:
    """Get the current email recipient, if one has been set."""
    try:
        return runtime.state["current_recipient"]
    except KeyError:
        return "No recipient has been set yet."

The try/except here matters: if set_current_recipient was never called in this conversation, current_recipient won't exist in state at all, and accessing it directly would raise a KeyError. Always handle this gracefully rather than letting it crash.

Watching It Work Across a Conversation​

Let's wire both tools into Aria and watch state actually persist across two separate messages in the same thread:

# Purpose: Confirm state set in one message is readable in a later, separate call
# Context: Same thread_id pattern as memory in Article 4 — state persists the same way
# Input: Two messages — one that sets the recipient, one that reads it back
# Output: The second response correctly recalls what was set in the first

from langchain.agents import create_agent
from langgraph.checkpoint.memory import InMemorySaver
from langchain.messages import HumanMessage

agent = create_agent(
    model="gpt-5-nano",
    tools=[set_current_recipient, get_current_recipient],
    state_schema=RecipientState,
    checkpointer=InMemorySaver(),
)

config = {"configurable": {"thread_id": "1"}}

# First message: establish who we're replying to
first_message = HumanMessage(content="I'm going to reply to Jane about coffee.")
response_one = agent.invoke({"messages": [first_message]}, config)
print(response_one["messages"][-1].content)

# Second message, separate call, same thread: ask who the recipient is
second_message = HumanMessage(content="Who am I currently replying to?")
response_two = agent.invoke({"messages": [second_message]}, config)
print(response_two["messages"][-1].content)

The second response correctly says Jane — not because it was repeated in the message, but because current_recipient was written to state during the first call and reloaded, alongside the message history, when the second call started. This is the exact same persistence mechanism from Article 4, just carrying more than messages now.

Common Misconceptions​

❌ Misconception: State and memory are two separate systems​

Reality: Memory, as covered in Article 4, is just the messages field of state. Custom state means adding your own fields to that same object — they're saved and reloaded together, by the same checkpointer, under the same thread_id.

Why this matters: Once you see state and memory as the same underlying mechanism, a lot of behavior becomes predictable — anything you add to state behaves exactly like the message history already did, because it's part of the same persisted object.

❌ Misconception: A Command update happens instantly to other tools running at the same time​

Reality: A state update from Command(update={...}) takes effect after the tool call completes, as part of that step in the graph — not as some kind of live, instantly-shared variable available mid-execution to other simultaneously-running tools.

Why this matters: If you're picturing state like a global variable multiple things can read and write to at the exact same instant, that's not quite the model — updates are applied in discrete steps as the agent processes the conversation, not continuously.

Troubleshooting Common Issues​

Problem: KeyError when reading a custom state field​

Symptoms: A tool reading runtime.state["some_field"] crashes with a KeyError.

Common Causes:

  1. The field was never set yet in this conversation/thread (most common — this is expected, not a bug, if you haven't handled it)
  2. The field name has a typo, not matching the state class exactly

Diagnostic Steps:

# Step 1: Always wrap state reads in a try/except, or check first
try:
    value = runtime.state["current_recipient"]
except KeyError:
    value = "Not set yet"

# Step 2: Double check the field name matches the state class exactly
class RecipientState(AgentState):
    current_recipient: str  # must match runtime.state["current_recipient"] exactly

Solution: Always handle the case where a state field hasn't been set yet — it's a normal part of working with state, not an error condition to avoid entirely.

Prevention: Treat every custom state field as optional from a reading tool's perspective, even if you expect it to usually be set by the time it's read.

Problem: State doesn't persist between separate invoke() calls​

Symptoms: A field set in one call appears to reset or disappear in the next call.

Common Causes:

  1. No checkpointer was configured on create_agent (most common — exactly the same requirement as memory in Article 4)
  2. Different thread_id values were used across the calls that were supposed to be connected

Solution: Confirm checkpointer=InMemorySaver() (or another checkpointer) is set, and that the same thread_id is passed in the config for every call that should share state — this is identical to the memory troubleshooting from Article 4, because it's the same underlying mechanism.

Check Your Understanding​

Quick Quiz​

  1. What is "memory" from Article 4, in terms of what you learned in this article?

    Show Answer

    Memory is just the messages field of agent state, persisted across calls by a checkpointer. Custom state means adding your own additional fields to that same state object, which get saved and reloaded the exact same way.

  2. How does a tool update state?

    Show Answer

    By returning a Command(update={...}) object, specifying which state fields to update (and typically including an updated messages list with a ToolMessage confirming what happened).

  3. Why does get_current_recipient use a try/except instead of reading runtime.state["current_recipient"] directly?

    Show Answer

    Because if set_current_recipient was never called earlier in the conversation, that key won't exist in state at all, and a direct read would raise a KeyError. Handling this gracefully is expected, not optional.

Hands-On Exercise​

Challenge: Add a draft_count: int field to RecipientState that increments by one every time a reply is drafted, and a tool to read the current count.

Show Solution
from langchain.agents import AgentState
from langchain.tools import tool, ToolRuntime
from langgraph.types import Command
from langchain.messages import ToolMessage

class RecipientState(AgentState):
    current_recipient: str
    draft_count: int

@tool
def record_draft(runtime: ToolRuntime) -> Command:
    """Record that a new draft reply was created, incrementing the count."""
    try:
        current_count = runtime.state["draft_count"]
    except KeyError:
        current_count = 0

    return Command(update={
        "draft_count": current_count + 1,
        "messages": [
            ToolMessage(f"Draft count now: {current_count + 1}", tool_call_id=runtime.tool_call_id)
        ],
    })

@tool
def get_draft_count(runtime: ToolRuntime) -> str:
    """Get how many drafts have been created so far in this conversation."""
    try:
        return str(runtime.state["draft_count"])
    except KeyError:
        return "0"

Explanation: Reading the current value before incrementing it (with the same try/except safety pattern) is necessary because a tool can't simply "add one" to a field that might not exist yet — it has to handle the unset case the same way any other state read does.

Summary: Key Takeaways​

  • "Memory" from Article 4 is really just the messages field of agent state — not a separate system
  • Custom state extends AgentState with your own fields, persisted the same way as messages, by the same checkpointer
  • Tools write to state by returning Command(update={...})
  • Tools read state through runtime.state["field_name"], and must handle the case where a field hasn't been set yet
  • State requires a checkpointer + thread_id, exactly like memory — because it's the same persistence mechanism
  • Aria can now carry information forward during a conversation, beyond just the literal messages exchanged

Version Information​

Tested with:

  • Python: >=3.10, <4.0
  • langchain: >=1.1.3 (latest stable as of writing: 1.3.4)
  • langgraph: >=1.0.3 — Command comes from langgraph.types, the same package InMemorySaver came from in Article 4

Known issues:

  • None specific to this article's functionality at the time of writing.

What's Next?​

You now understand the relationship between memory and state, and how to extend state with your own fields that tools can read and write during a conversation.

The natural next step is SQL Agents: Querying Databases in Natural Language — Aria is about to gain access to a real structured data source for the first time, querying it using everything you've learned about tools so far.

References​