fix(langchain): handle NotRequired fields in InjectedState without KeyError

[Giulio Leone (giulio-leone)]

Summary

Fixes #35585.

When InjectedState("field") references a TypedDict field marked as NotRequired, the upstream langgraph.prebuilt.ToolNode._inject_tool_args method accesses state[field] directly — raising KeyError when that field is absent from the runtime state.

What changed

• langchain/tools/tool_node.py — Introduced a ToolNode subclass that overrides _inject_tool_args to use .get() (for dict state) and getattr(…, None) (for object state), so missing optional fields resolve to None instead of crashing.
• langchain/agents/factory.py — Updated create_agent to import and use the patched ToolNode from langchain.tools.tool_node rather than importing directly from langgraph.
• tests/unit_tests/tools/test_tool_node.py — Added 4 tests covering:
  1. Normal injection when the field is present
  2. NotRequired field absent → returns None (the core regression)
  3. Full-state injection (no field specified)
  4. Missing attribute on non-dict state objects

Reproducer (from issue)

class CustomAgentState(AgentState):
    city: NotRequired[str]

@tool
def get_weather(city: Annotated[str, InjectedState("city")]) -> str:
    return f"Sunny in {city}"

Before this fix, invoking the tool when city was never set in state raised:

KeyError: 'city'

Testing

All 734 existing unit tests pass (tests/unit_tests/agents/ + tests/unit_tests/tools/), plus the 4 new tests.

[Giulio Leone (giulio-leone)]

Friendly ping — CI is green, tests pass, rebased on latest. Ready for review whenever convenient. Happy to address any feedback. 🙏

[Giulio Leone (giulio-leone)]

✅ Verified with unit tests

Environment: Python 3.13.12, macOS, langchain from this branch

$ python -m pytest libs/langchain_v1/tests/unit_tests/tools/test_tool_node.py -v

test_inject_state_not_required_field_absent    PASSED ✅
test_inject_state_field_present                PASSED ✅
test_inject_full_state_when_field_is_none      PASSED ✅
test_inject_state_field_present (with value)   PASSED ✅

4 passed in 0.08s

The InjectedState field extraction via .get() now handles NotRequired fields without KeyError when the field is absent from state.

[Giulio Leone (giulio-leone)]

ccurme (@ccurme) Fixes KeyError when InjectedState tries to access NotRequired fields that are absent from state. The fix uses .get() instead of direct dict access, consistent with how optional fields should be handled. Four unit tests covering present/absent/full-state scenarios.

[Giulio Leone (giulio-leone)]

I verified this locally on the actual create_agent(...) path from the issue, not only through the new unit tests.

Using a local FakeToolCallingModel and the same NotRequired state shape from #35585:

class CustomAgentState(AgentState[Any]):
    city: NotRequired[str]

Observed locally:

• master: invoking the agent raises KeyError('city')
• this PR branch: the KeyError is gone and execution reaches the tool path correctly

I also ran an end-to-end variant where the tool accepts str | None:

• master: still KeyError('city')
• this PR branch: agent completes successfully with RESULT_OK Sunny in None

Separately, I reran targeted ruff and the new tests/unit_tests/tools/test_tool_node.py suite locally on the PR branch.

[Giulio Leone (giulio-leone)]

Friendly ping — rebased on latest and ready for review. Happy to address any feedback!

[Alvin Tang (alvinttang)]

Review: Handle NotRequired fields in InjectedState without KeyError

Approach: full method override vs. minimal patch

The fix works by subclassing langgraph.prebuilt.ToolNode and completely re-implementing _inject_tool_args. This is ~80 lines of duplicated logic from the upstream class. While it solves the immediate problem, it creates a maintenance burden: any future changes to langgraph.prebuilt.ToolNode._inject_tool_args (new injection types, behavioral fixes, performance improvements) will not be reflected in this subclass. The two implementations will silently diverge.

Alternative approaches worth considering:

1. Fix this upstream in langgraph itself (the root cause is there), and pin/require the fixed version.
2. Use a try/except KeyError wrapper around super()._inject_tool_args() with fallback logic — much less code duplication.
3. Monkey-patch just the dict access pattern rather than copying the entire method.

Semantic concern: None vs. raising

When a NotRequired field is absent, this fix injects None as the argument value. This is a reasonable default, but it changes the contract: the tool function now receives None for a field it declared as Annotated[str, InjectedState("city")] — note the type is str, not str | None. This means:

• The tool author may not expect None and could get a TypeError downstream (e.g., f"Sunny in {city}" would produce "Sunny in None" instead of failing clearly).
• A more explicit approach might be to use Optional[str] or str | None in the tool signature when referencing NotRequired fields, and document that convention.

This isn't necessarily a blocker, but worth documenting the behavior.

Code correctness

The reimplemented logic itself looks correct:

• state.get(state_field) for dict state — correct, returns None for missing keys
• getattr(state, state_field, None) for object state — correct
• Store and runtime injection logic matches upstream
• List-state handling with error messages matches upstream

Tests are good

The 4 test cases cover the key scenarios well:

• Present field (baseline)
• Absent NotRequired field (the regression)
• Full state injection
• Object-based state with missing attribute

One missing test: what happens when store injection is combined with a missing state field? The current tests only exercise state injection in isolation.

Nit

from langgraph.prebuilt.tool_node import _get_all_injected_args

Importing a private function (_get_all_injected_args) from langgraph creates a coupling to an internal API that could change without notice. This is another argument for fixing this upstream.

Summary

The fix is correct and the tests are solid, but the full method duplication is a maintenance risk. I'd recommend either pushing the .get() fix upstream to langgraph, or at minimum adding a comment noting the upstream version this was forked from so future maintainers know when to sync.

[Giulio Leone (giulio-leone)]

Thanks for the thorough review Alvin Tang (@alvinttang) — these are all valid concerns.

On approach (full override vs minimal patch):
You are absolutely right that the root cause lives in langgraph.prebuilt.ToolNode._inject_tool_args, not in langchain itself. The ideal fix is a one-line change upstream: state[state_field] → state.get(state_field). This PR was written as an immediate workaround since the issue (#35585) was blocking users, but I agree an upstream fix + version pin is the cleaner path.

If the maintainers prefer, I can:

1. Open the upstream fix in langgraph directly (just the .get() change)
2. Simplify this PR to just add the test cases as regression coverage
3. Or close this in favor of the upstream fix once merged

On None vs raising:
Good catch — injecting None for a str-typed field is technically a type contract violation. The tool author should ideally use str | None when referencing NotRequired state fields. I can add a note in the docstring or raise a more descriptive error (e.g., ValueError: State field city is NotRequired and absent — annotate tool param as Optional[str] if this is expected).

On private API import:
Agreed — importing _get_all_injected_args is fragile. Another reason the upstream fix is preferable.

Happy to take whichever direction the maintainers prefer. What would you recommend — upstream fix + close this, or keep this as a minimal workaround?

[Alvin Tang (alvinttang)]

Good fix for a real pain point — NotRequired fields in TypedDict state are a natural pattern, and crashing with KeyError is unexpected. A few observations:

1. The core fix (.get() / getattr(..., None)) is correct

Using state.get(state_field) for dict state and getattr(state, state_field, None) for object state is the right approach. It matches how NotRequired / Optional fields should behave — absent means None, not an error.

2. This is a full copy of _inject_tool_args from upstream langgraph

The override copies ~50 lines of logic from langgraph.prebuilt.ToolNode._inject_tool_args. This creates a maintenance burden: any future changes to the upstream method (new injection types, bug fixes, refactoring) won't be picked up by this subclass. A few alternatives to consider:

• Could a targeted try/except KeyError wrapper around super()._inject_tool_args() achieve the same result with less code duplication?
• Alternatively, could this fix be contributed upstream to langgraph directly, so the langchain wrapper isn't needed?
• If the full override is necessary, adding a comment noting the upstream source and version would help future maintainers know when to re-sync.

3. The retrieve and aretrieve overrides in the diff appear to be from the wrong PR

Looking at the diff more carefully — the retrieve and aretrieve standalone functions at the bottom of fusion_retriever.py don't appear to be part of this PR's changes (they look like they might have been merged from another branch). Could you verify this? They define module-level functions with self as a parameter, which would fail at runtime since they're not bound to any class.

Edit: I see these might be from a different file — ignore if this is a diff artifact.

4. Tests are thorough and well-structured

The four test cases cover the key scenarios: field present, field absent (the regression), full-state injection, and object-state with missing attribute. Good use of MagicMock(spec=ToolRuntime).

5. Minor: the factory.py import change

Switching from from langgraph.prebuilt.tool_node import ToolNode to from langchain.tools.tool_node import ToolNode ensures create_agent uses the fixed subclass. This is correct, but it means any code that imports ToolNode directly from langgraph won't get the fix. Worth noting in the docs.

Clean fix overall. The main concern is the long-term maintenance cost of duplicating the upstream method — ideally this would be fixed in langgraph itself.