Isolating Subgraph State When Calling Graphs from Tools

from typing import Annotated
from langchain_core.messages import AIMessage, HumanMessage
from langchain_core.runnables import RunnableConfig
from langchain_core.tools import tool
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import StateGraph
from langgraph.graph.message import add_messages
from langgraph.prebuilt import create_react_agent
from langchain_anthropic import ChatAnthropic
from dotenv import load_dotenv
from typing_extensions import TypedDict

load_dotenv()

# === Graph B (subagent) ===
class GraphBState(TypedDict):
    messages: Annotated[list, add_messages]

def graph_b_node(state: GraphBState):
    return {"messages": [AIMessage(content="[GRAPH_B] Response")]}

checkpointer_b = MemorySaver()  # Own checkpointer
graph_b = StateGraph(GraphBState)
graph_b.add_node("node", graph_b_node)
graph_b.set_entry_point("node")
graph_b.set_finish_point("node")
graph_b = graph_b.compile(checkpointer=checkpointer_b)


# === Tool that calls Graph B ===
@tool
def call_graph_b(query: str, config: RunnableConfig) -> str:
    """Call Graph B subagent."""
    thread_id = config.get("configurable", {}).get("thread_id")
    result = graph_b.invoke(
        {"messages": [HumanMessage(content=query)]},
        {"configurable": {"thread_id": thread_id}}
    )
    return f"Done"


# === Graph A (parent agent) ===
checkpointer_a = MemorySaver()  # Own checkpointer
agent = create_react_agent(
    model=ChatAnthropic(model="claude-sonnet-4-20250514"),
    tools=[call_graph_b],
    checkpointer=checkpointer_a,
)


# === Run ===
thread_id = "conv_12345"
config = {"configurable": {"thread_id": thread_id}}

agent.invoke({"messages": [HumanMessage(content="call graph_b with 'hello'")]}, config)

graph_a_state = agent.get_state(config)
graph_b_state = graph_b.get_state({"configurable": {"thread_id": thread_id}})

print(f"Graph A thread_id: {graph_a_state.config['configurable']['thread_id']}")
print(f"Graph B thread_id: {graph_b_state.config['configurable']['thread_id']}")
print()
print(f"Graph A messages: {len(graph_a_state.values['messages'])}")
print(f"Graph B messages: {len(graph_b_state.values['messages'])}")
print()
print(f"States isolated: {not any('[GRAPH_B]' in str(m.content) for m in graph_a_state.values['messages'])}")