How to pass a parent's graph state into a React agent (a subgraph)?

[miroslavtushev]

Can't wrap my head around this. If I call create_react_agent it will create a React agent as a subgraph and compile it. Once you invoke it from a node from the outer graph and pass a state created in the outer graph - why can't that state be updated by that React agent or its tools later? Am I missing something here?

class SampleState(MessagesState):
    some_attr: str

haiku = ChatBedrockConverse(model="us.anthropic.claude-3-5-haiku-20241022-v1:0", temperature=0, region_name='us-west-2')

@tool
def update_attr(s: str, tool_call_id: Annotated[str, InjectedToolCallId]):
    """Run this to update the attribute"""
    return Command(
        update={"messages": [ToolMessage(f"Attribute has been updated", tool_call_id=tool_call_id)], "some_attr" : s} # updating some_attr
    )

agent = create_react_agent(
    haiku,
    tools=[update_attr]
)

def state_update(state: SampleState):
    return {"some_attr" : "from outer graph"}

def node(state: SampleState):
    res = agent.invoke(state) #invoking here with the original state. Why does update_attr not apply the correct update to some_attr?
    return {"messages" : res["messages"]}

builder = StateGraph(SampleState)
builder.add_node("node", node)
builder.add_node("state_update", state_update)
builder.add_edge(START, "state_update")
builder.add_edge("state_update", "node")
builder.add_edge("node", END)

checkpointer = InMemorySaver()

graph = builder.compile(checkpointer=checkpointer)

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

event = graph.invoke(
    {"messages": [{"role": "user", "content": "Please update attribute with the following: risk123"}]},
    config=config,
    # stream_mode='updates'
)
for message in event['messages']:
    message.pretty_print()

================================ Human Message =================================

Please update attribute with the following: risk123
================================== Ai Message ==================================

[{'type': 'text', 'text': 'I\'ll help you update the attribute with the value "risk123" using the `update_attr` function.'}, {'type': 'tool_use', 'name': 'update_attr', 'input': {'s': 'risk123'}, 'id': 'tooluse_PonaNmUfTnyBRnC_7DONqg'}]
Tool Calls:
  update_attr (tooluse_PonaNmUfTnyBRnC_7DONqg)
 Call ID: tooluse_PonaNmUfTnyBRnC_7DONqg
  Args:
    s: risk123
================================= Tool Message =================================
Name: update_attr

Attribute has been updated
================================== Ai Message ==================================

The attribute has been successfully updated with the value "risk123". Is there anything else I can help you with?

Printing the state shows some_attr=="from outer graph"

graph.get_state(config, subgraphs=True)

StateSnapshot(values={'messages': [HumanMessage(content='Please update attribute with the following: risk123', additional_kwargs={}, response_metadata={}, id='cb19f9b8-70fd-4d63-9fef-40a318d3d69e'), AIMessage(content=[{'type': 'text', 'text': 'I\'ll help you update the attribute with the value "risk123" using the update_attrfunction.'}, {'type': 'tool_use', 'name': 'update_attr', 'input': {'s': 'risk123'}, 'id': 'tooluse_PonaNmUfTnyBRnC_7DONqg'}], additional_kwargs={}, response_metadata={'ResponseMetadata': {'RequestId': '6b857955-fe38-4f97-8f52-1bdf628d9713', 'HTTPStatusCode': 200, 'HTTPHeaders': {'date': 'Fri, 27 Jun 2025 22:33:04 GMT', 'content-type': 'application/json', 'content-length': '383', 'connection': 'keep-alive', 'x-amzn-requestid': '6b857955-fe38-4f97-8f52-1bdf628d9713'}, 'RetryAttempts': 0}, 'stopReason': 'tool_use', 'metrics': {'latencyMs': [1701]}, 'model_name': 'us.anthropic.claude-3-5-haiku-20241022-v1:0'}, id='run--d8863fdd-3f40-4083-94fa-3c3aa0dc1afb-0', tool_calls=[{'name': 'update_attr', 'args': {'s': 'risk123'}, 'id': 'tooluse_PonaNmUfTnyBRnC_7DONqg', 'type': 'tool_call'}], usage_metadata={'input_tokens': 334, 'output_tokens': 78, 'total_tokens': 412, 'input_token_details': {'cache_creation': 0, 'cache_read': 0}}), ToolMessage(content='Attribute has been updated', name='update_attr', id='8d504238-b6d4-4dd7-b6c1-ee094b82f856', tool_call_id='tooluse_PonaNmUfTnyBRnC_7DONqg'), AIMessage(content='The attribute has been successfully updated with the value "risk123". Is there anything else I can help you with?', additional_kwargs={}, response_metadata={'ResponseMetadata': {'RequestId': '4dfcb2e2-c515-4fd7-9cd5-dccc7f257093', 'HTTPStatusCode': 200, 'HTTPHeaders': {'date': 'Fri, 27 Jun 2025 22:33:09 GMT', 'content-type': 'application/json', 'content-length': '409', 'connection': 'keep-alive', 'x-amzn-requestid': '4dfcb2e2-c515-4fd7-9cd5-dccc7f257093'}, 'RetryAttempts': 0}, 'stopReason': 'end_turn', 'metrics': {'latencyMs': [4956]}, 'model_name': 'us.anthropic.claude-3-5-haiku-20241022-v1:0'}, id='run--3a36e053-02e4-4315-99b9-954713fdd6d5-0', usage_metadata={'input_tokens': 428, 'output_tokens': 27, 'total_tokens': 455, 'input_token_details': {'cache_creation': 0, 'cache_read': 0}})], 'some_attr': 'from outer graph'}, next=(), config={'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1f053a6b-3ebe-6856-8002-a360ad4a4792'}}, metadata={'source': 'loop', 'writes': {'node': {'messages': [HumanMessage(content='Please update attribute with the following: risk123', additional_kwargs={}, response_metadata={}, id='cb19f9b8-70fd-4d63-9fef-40a318d3d69e'), AIMessage(content=[{'type': 'text', 'text': 'I\'ll help you update the attribute with the value "risk123" using theupdate_attr function.'}, {'type': 'tool_use', 'name': 'update_attr', 'input': {'s': 'risk123'}, 'id': 'tooluse_PonaNmUfTnyBRnC_7DONqg'}], additional_kwargs={}, response_metadata={'ResponseMetadata': {'RequestId': '6b857955-fe38-4f97-8f52-1bdf628d9713', 'HTTPStatusCode': 200, 'HTTPHeaders': {'date': 'Fri, 27 Jun 2025 22:33:04 GMT', 'content-type': 'application/json', 'content-length': '383', 'connection': 'keep-alive', 'x-amzn-requestid': '6b857955-fe38-4f97-8f52-1bdf628d9713'}, 'RetryAttempts': 0}, 'stopReason': 'tool_use', 'metrics': {'latencyMs': [1701]}, 'model_name': 'us.anthropic.claude-3-5-haiku-20241022-v1:0'}, id='run--d8863fdd-3f40-4083-94fa-3c3aa0dc1afb-0', tool_calls=[{'name': 'update_attr', 'args': {'s': 'risk123'}, 'id': 'tooluse_PonaNmUfTnyBRnC_7DONqg', 'type': 'tool_call'}], usage_metadata={'input_tokens': 334, 'output_tokens': 78, 'total_tokens': 412, 'input_token_details': {'cache_creation': 0, 'cache_read': 0}}), ToolMessage(content='Attribute has been updated', name='update_attr', id='8d504238-b6d4-4dd7-b6c1-ee094b82f856', tool_call_id='tooluse_PonaNmUfTnyBRnC_7DONqg'), AIMessage(content='The attribute has been successfully updated with the value "risk123". Is there anything else I can help you with?', additional_kwargs={}, response_metadata={'ResponseMetadata': {'RequestId': '4dfcb2e2-c515-4fd7-9cd5-dccc7f257093', 'HTTPStatusCode': 200, 'HTTPHeaders': {'date': 'Fri, 27 Jun 2025 22:33:09 GMT', 'content-type': 'application/json', 'content-length': '409', 'connection': 'keep-alive', 'x-amzn-requestid': '4dfcb2e2-c515-4fd7-9cd5-dccc7f257093'}, 'RetryAttempts': 0}, 'stopReason': 'end_turn', 'metrics': {'latencyMs': [4956]}, 'model_name': 'us.anthropic.claude-3-5-haiku-20241022-v1:0'}, id='run--3a36e053-02e4-4315-99b9-954713fdd6d5-0', usage_metadata={'input_tokens': 428, 'output_tokens': 27, 'total_tokens': 455, 'input_token_details': {'cache_creation': 0, 'cache_read': 0}})]}}, 'step': 2, 'parents': {}, 'thread_id': '1'}, created_at='2025-06-27T22:33:09.674391+00:00', parent_config={'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1f053a6a-fed3-6943-8001-1c9d145094ed'}}, tasks=(), interrupts=())

[miroslavtushev]

Update: after some pondering and reading through the docs I was able to fix this by providing a state_schema argument to create_react_agent, like so:

agent = create_react_agent(
    haiku,
    tools=[update_attr],
    state_schema=SampleState
)