Resolves https://github.com/langchain-ai/langchain/issues/35585
This would previously raise KeyError:
```python
from typing import Annotated
from langchain_core.tools import tool
from langchain.agents import create_agent
from typing_extensions import NotRequired
from langgraph.prebuilt import InjectedState
from langchain.agents import AgentState
class CustomAgentState(AgentState):
city: NotRequired[str]
@tool
def get_weather(city: Annotated[str | None, InjectedState("city")] = None) -> str:
"""Get weather for a given city."""
if city is None:
city = "Boston"
return f"It's always sunny in {city}!"
agent = create_agent(
model="claude-sonnet-4-6",
tools=[get_weather],
system_prompt="You are a helpful assistant",
state_schema=CustomAgentState,
)
input_message = {
"role": "user",
"content": "What's the weather?",
}
result = agent.invoke({"messages": [input_message]})
for m in result["messages"]:
m.pretty_print()
```
---------
Co-authored-by: Sydney Runkle <sydneymarierunkle@gmail.com>