fix(prebuilt): handle injected NotRequired keys (#7392)

Resolves langchain-ai/langchain#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>

Author: ccurme

libs/prebuilt/langgraph/prebuilt/tool_node.py

@@ -614,6 +614,7 @@ def my_tool(
     store: str | None
     runtime: str | None
     all_injected_keys: set[str]
+    _optional_state_args: set[str]
 
 
 class ToolNode(RunnableCallable):
@@ -1333,7 +1334,7 @@ def _inject_tool_args(
             return tool_call
 
         tool_call_copy: ToolCall = copy(tool_call)
-        injected_args = {}
+        injected_args: dict[str, Any] = {}
 
         # Inject state
         if injected.state:
@@ -1361,14 +1362,20 @@ def _inject_tool_args(
             # Extract state values
             if isinstance(state, dict):
                 for tool_arg, state_field in injected.state.items():
-                    injected_args[tool_arg] = (
-                        state[state_field] if state_field else state
-                    )
+                    if not state_field:
+                        injected_args[tool_arg] = state
+                    elif state_field in state:
+                        injected_args[tool_arg] = state[state_field]
+                    elif tool_arg not in injected._optional_state_args:
+                        raise KeyError(state_field)
             else:
                 for tool_arg, state_field in injected.state.items():
-                    injected_args[tool_arg] = (
-                        getattr(state, state_field) if state_field else state
-                    )
+                    if not state_field:
+                        injected_args[tool_arg] = state
+                    elif hasattr(state, state_field):
+                        injected_args[tool_arg] = getattr(state, state_field)
+                    elif tool_arg not in injected._optional_state_args:
+                        raise AttributeError(state_field)
 
         # Inject store
         if injected.store:
@@ -1859,6 +1866,7 @@ def _get_all_injected_args(tool: BaseTool) -> _InjectedArgs:
     store_arg: str | None = None
     runtime_arg: str | None = None
     all_injected_keys: set[str] = set()
+    _optional_state_args: set[str] = set()
 
     for name, type_ in all_annotations.items():
         # Track all InjectedToolArg-annotated params (including custom subclasses)
@@ -1873,6 +1881,9 @@ def _get_all_injected_args(tool: BaseTool) -> _InjectedArgs:
         if state_inj := _get_injection_from_type(type_, InjectedState):
             if isinstance(state_inj, InjectedState) and state_inj.field:
                 state_args[name] = state_inj.field
+                field_info = full_schema.model_fields.get(name)
+                if field_info and not field_info.is_required():
+                    _optional_state_args.add(name)
             else:
                 state_args[name] = None
 
@@ -1889,4 +1900,5 @@ def _get_all_injected_args(tool: BaseTool) -> _InjectedArgs:
         store=store_arg,
         runtime=runtime_arg,
         all_injected_keys=all_injected_keys,
+        _optional_state_args=_optional_state_args,
     )

libs/prebuilt/tests/test_injected_state_not_required.py

@@ -0,0 +1,284 @@
+"""Test InjectedState with NotRequired state fields.
+
+This tests the fix for https://github.com/langchain-ai/langchain/issues/35585
+
+When using InjectedState(<field>) on a tool parameter, and the referenced field is
+declared as NotRequired in the custom state schema, the ToolNode should gracefully
+handle missing fields by injecting None instead of raising KeyError.
+"""
+
+import sys
+from typing import Annotated
+
+import pytest
+from langchain_core.messages import AIMessage, AnyMessage, HumanMessage, ToolMessage
+from langchain_core.tools import tool
+from langgraph.graph.message import add_messages
+from pydantic import BaseModel, Field
+from typing_extensions import NotRequired
+
+from langgraph.prebuilt import InjectedState, ToolNode, create_react_agent
+from langgraph.prebuilt.chat_agent_executor import AgentState
+
+from .model import FakeToolCallingModel
+
+
+class CustomAgentStateWithNotRequired(AgentState):
+    """Custom state with a NotRequired field (TypedDict style)."""
+
+    city: NotRequired[str]
+
+
+class CustomAgentStatePydanticWithDefault(BaseModel):
+    """Custom state with Optional field and default (Pydantic style)."""
+
+    messages: Annotated[list[AnyMessage], add_messages]
+    remaining_steps: int = Field(default=10)
+    city: str | None = Field(default=None)
+
+
+@tool
+def get_weather(city: Annotated[str | None, InjectedState("city")] = None) -> str:
+    """Get weather for a given city."""
+    if city is None:
+        return "No city provided"
+    return f"It's always sunny in {city}!"
+
+
+def _create_mock_runtime(
+    state: dict | None = None,
+    store=None,
+):
+    """Create a mock Runtime for testing ToolNode directly."""
+    from unittest.mock import Mock
+
+    from langgraph.runtime import Runtime
+
+    mock_runtime = Mock(spec=Runtime)
+    mock_runtime.context = {}
+    return mock_runtime
+
+
+def _create_config_with_runtime(store=None, state=None):
+    """Create a RunnableConfig with mocked runtime for direct ToolNode testing."""
+    from langgraph.prebuilt.tool_node import ToolRuntime
+
+    tool_runtime = ToolRuntime(
+        state=state or {},
+        config={},
+        context={},
+        store=store,
+        stream_writer=None,
+        tool_call_id="test_id",
+    )
+    return {
+        "configurable": {
+            "__pregel_runtime": _create_mock_runtime(),
+            "__tool_runtime__": tool_runtime,
+        }
+    }
+
+
+@pytest.mark.skipif(
+    sys.version_info < (3, 11),
+    reason="InjectedState field extraction from Optional[Annotated[...]] not supported on Python <3.11",
+)
+def test_injected_state_not_required_field_missing_injects_none():
+    """Test that InjectedState with NotRequired field injects None when field is missing.
+
+    This verifies the fix for https://github.com/langchain-ai/langchain/issues/35585
+    """
+    tool_node = ToolNode([get_weather])
+
+    tool_call = {
+        "name": "get_weather",
+        "args": {},
+        "id": "call_1",
+        "type": "tool_call",
+    }
+    ai_msg = AIMessage("Let me check the weather", tool_calls=[tool_call])
+
+    # State WITHOUT the "city" field - should inject None instead of raising KeyError
+    state_without_city: CustomAgentStateWithNotRequired = {
+        "messages": [HumanMessage("What's the weather?"), ai_msg],
+    }
+
+    result = tool_node.invoke(
+        state_without_city,
+        config=_create_config_with_runtime(state=state_without_city),
+    )
+
+    assert len(result["messages"]) == 1
+    tool_msg = result["messages"][0]
+    assert isinstance(tool_msg, ToolMessage)
+    assert "No city provided" in tool_msg.content
+
+
+@pytest.mark.skipif(
+    sys.version_info < (3, 11),
+    reason="InjectedState field extraction from Optional[Annotated[...]] not supported on Python <3.11",
+)
+def test_injected_state_not_required_field_present_works():
+    """Test that InjectedState with NotRequired field works when field IS present."""
+    tool_node = ToolNode([get_weather])
+
+    tool_call = {
+        "name": "get_weather",
+        "args": {},
+        "id": "call_1",
+        "type": "tool_call",
+    }
+    ai_msg = AIMessage("Let me check the weather", tool_calls=[tool_call])
+
+    # State WITH the "city" field - this should work
+    state_with_city: CustomAgentStateWithNotRequired = {
+        "messages": [HumanMessage("What's the weather?"), ai_msg],
+        "city": "San Francisco",
+    }
+
+    result = tool_node.invoke(
+        state_with_city,
+        config=_create_config_with_runtime(state=state_with_city),
+    )
+
+    assert len(result["messages"]) == 1
+    tool_msg = result["messages"][0]
+    assert isinstance(tool_msg, ToolMessage)
+    assert "San Francisco" in tool_msg.content
+
+
+@pytest.mark.skipif(
+    sys.version_info < (3, 11),
+    reason="InjectedState field extraction from Optional[Annotated[...]] not supported on Python <3.11",
+)
+def test_create_react_agent_injected_state_not_required_field_missing():
+    """Test create_react_agent with InjectedState using NotRequired field that is missing.
+
+    This verifies the fix for https://github.com/langchain-ai/langchain/issues/35585
+    """
+    model = FakeToolCallingModel(
+        tool_calls=[
+            [{"name": "get_weather", "args": {}, "id": "call_1"}],
+            [],  # No more tool calls, agent should stop
+        ]
+    )
+
+    agent = create_react_agent(
+        model,
+        tools=[get_weather],
+        state_schema=CustomAgentStateWithNotRequired,
+    )
+
+    # Invoke WITHOUT the city field - should work, injecting None
+    result = agent.invoke(
+        {"messages": [HumanMessage("What's the weather?")]},
+    )
+
+    # Check that the tool was called successfully with None injected
+    messages = result["messages"]
+    tool_messages = [m for m in messages if isinstance(m, ToolMessage)]
+    assert len(tool_messages) == 1
+    assert "No city provided" in tool_messages[0].content
+
+
+@pytest.mark.skipif(
+    sys.version_info < (3, 11),
+    reason="InjectedState field extraction from Optional[Annotated[...]] not supported on Python <3.11",
+)
+def test_create_react_agent_injected_state_not_required_field_present():
+    """Test create_react_agent with InjectedState using NotRequired field that IS present."""
+    model = FakeToolCallingModel(
+        tool_calls=[
+            [{"name": "get_weather", "args": {}, "id": "call_1"}],
+            [],  # No more tool calls, agent should stop
+        ]
+    )
+
+    agent = create_react_agent(
+        model,
+        tools=[get_weather],
+        state_schema=CustomAgentStateWithNotRequired,
+    )
+
+    # Invoke WITH the city field
+    result = agent.invoke(
+        {
+            "messages": [HumanMessage("What's the weather?")],
+            "city": "San Francisco",
+        },
+    )
+
+    # Check that the tool was called successfully
+    messages = result["messages"]
+    tool_messages = [m for m in messages if isinstance(m, ToolMessage)]
+    assert len(tool_messages) == 1
+    assert "San Francisco" in tool_messages[0].content
+
+
+@tool
+def get_weather_optional(city: Annotated[str | None, InjectedState("city")]) -> str:
+    """Get weather for a given city (accepts None)."""
+    if city is None:
+        return "Please provide a city!"
+    return f"It's always sunny in {city}!"
+
+
+def test_pydantic_state_with_default_field_missing_works():
+    """Test that Pydantic state with Optional field and default=None works when field is missing.
+
+    This is the workaround suggested in the issue comments - using Pydantic BaseModel
+    with `city: Optional[str] = Field(default=None)` instead of TypedDict with NotRequired.
+    """
+    model = FakeToolCallingModel(
+        tool_calls=[
+            [{"name": "get_weather_optional", "args": {}, "id": "call_1"}],
+            [],  # No more tool calls, agent should stop
+        ]
+    )
+
+    agent = create_react_agent(
+        model,
+        tools=[get_weather_optional],
+        state_schema=CustomAgentStatePydanticWithDefault,
+    )
+
+    # Invoke WITHOUT the city field - should work because Pydantic provides default
+    result = agent.invoke(
+        {"messages": [HumanMessage("What's the weather?")]},
+    )
+
+    # Check that the tool was called successfully with None
+    messages = result["messages"]
+    tool_messages = [m for m in messages if isinstance(m, ToolMessage)]
+    assert len(tool_messages) == 1
+    assert "Please provide a city!" in tool_messages[0].content
+
+
+def test_pydantic_state_with_default_field_present_works():
+    """Test that Pydantic state with Optional field works when field IS present."""
+    model = FakeToolCallingModel(
+        tool_calls=[
+            [{"name": "get_weather_optional", "args": {}, "id": "call_1"}],
+            [],  # No more tool calls, agent should stop
+        ]
+    )
+
+    agent = create_react_agent(
+        model,
+        tools=[get_weather_optional],
+        state_schema=CustomAgentStatePydanticWithDefault,
+    )
+
+    # Invoke WITH the city field
+    result = agent.invoke(
+        {
+            "messages": [HumanMessage("What's the weather?")],
+            "city": "San Francisco",
+        },
+    )
+
+    # Check that the tool was called successfully
+    messages = result["messages"]
+    tool_messages = [m for m in messages if isinstance(m, ToolMessage)]
+    assert len(tool_messages) == 1
+    assert "San Francisco" in tool_messages[0].content