Do not inject args in RunnableCallable if arg already exists (#3185)

Co-authored-by: Eugene Yurtsev <eyurtsev@gmail.com>

libs/langgraph/langgraph/prebuilt/tool_node.py

@@ -27,7 +27,6 @@ from langchain_core.runnables.config import (
     get_config_list,
     get_executor_for_config,
 )
-from langchain_core.runnables.utils import Input
 from langchain_core.tools import BaseTool, InjectedToolArg
 from langchain_core.tools import tool as create_tool
 from langchain_core.tools.base import get_all_basemodel_annotations
@@ -210,7 +209,7 @@ class ToolNode(RunnableCallable):
         ],
         config: RunnableConfig,
         *,
-        store: BaseStore,
+        store: Optional[BaseStore],
     ) -> Any:
         tool_calls, input_type = self._parse_input(input, store)
         config_list = get_config_list(config, len(tool_calls))
@@ -220,12 +219,14 @@ class ToolNode(RunnableCallable):
                 *executor.map(self._run_one, tool_calls, input_types, config_list)
             ]
 
-        # preserve existing behavior for non-command tool outputs for backwards compatibility
+        # preserve existing behavior for non-command tool outputs for backwards
+        # compatibility
         if not any(isinstance(output, Command) for output in outputs):
             # TypedDict, pydantic, dataclass, etc. should all be able to load from dict
             return outputs if input_type == "list" else {self.messages_key: outputs}
 
-        # LangGraph will automatically handle list of Command and non-command node updates
+        # LangGraph will automatically handle list of Command and non-command node
+        # updates
         combined_outputs: list[
             Command | list[ToolMessage] | dict[str, list[ToolMessage]]
         ] = []
@@ -238,20 +239,6 @@ class ToolNode(RunnableCallable):
                 )
         return combined_outputs
 
-    def invoke(
-        self, input: Input, config: Optional[RunnableConfig] = None, **kwargs: Any
-    ) -> Any:
-        if "store" not in kwargs:
-            kwargs["store"] = None
-        return super().invoke(input, config, **kwargs)
-
-    async def ainvoke(
-        self, input: Input, config: Optional[RunnableConfig] = None, **kwargs: Any
-    ) -> Any:
-        if "store" not in kwargs:
-            kwargs["store"] = None
-        return await super().ainvoke(input, config, **kwargs)
-
     async def _afunc(
         self,
         input: Union[
@@ -261,7 +248,7 @@ class ToolNode(RunnableCallable):
         ],
         config: RunnableConfig,
         *,
-        store: BaseStore,
+        store: Optional[BaseStore],
     ) -> Any:
         tool_calls, input_type = self._parse_input(input, store)
         outputs = await asyncio.gather(
@@ -404,7 +391,7 @@ class ToolNode(RunnableCallable):
             dict[str, Any],
             BaseModel,
         ],
-        store: BaseStore,
+        store: Optional[BaseStore],
     ) -> Tuple[list[ToolCall], Literal["list", "dict"]]:
         if isinstance(input, list):
             input_type = "list"
@@ -484,7 +471,9 @@ class ToolNode(RunnableCallable):
         }
         return tool_call
 
-    def _inject_store(self, tool_call: ToolCall, store: BaseStore) -> ToolCall:
+    def _inject_store(
+        self, tool_call: ToolCall, store: Optional[BaseStore]
+    ) -> ToolCall:
         store_arg = self.tool_to_store_arg[tool_call["name"]]
         if not store_arg:
             return tool_call
@@ -509,7 +498,7 @@ class ToolNode(RunnableCallable):
             dict[str, Any],
             BaseModel,
         ],
-        store: BaseStore,
+        store: Optional[BaseStore],
     ) -> ToolCall:
         if tool_call["name"] not in self.tools_by_name:
             return tool_call

libs/langgraph/langgraph/utils/runnable.py

@@ -14,6 +14,7 @@ from typing import (
     Iterator,
     Optional,
     Sequence,
+    Tuple,
     Union,
     cast,
 )
@@ -66,9 +67,10 @@ class StrEnum(str, enum.Enum):
 # Special type to denote any type is accepted
 ANY_TYPE = object()
 
-
 ASYNCIO_ACCEPTS_CONTEXT = sys.version_info >= (3, 11)
 
+# List of keyword arguments that can be injected at runtime from the config object.
+# A named argument may appear multiple times if it appears with distinct types.
 KWARGS_CONFIG_KEYS: tuple[tuple[str, tuple[Any, ...], str, Any], ...] = (
     (
         sys.intern("writer"),
@@ -77,11 +79,31 @@ KWARGS_CONFIG_KEYS: tuple[tuple[str, tuple[Any, ...], str, Any], ...] = (
         lambda _: None,
     ),
     (
+        # Covers store that is not optional (will raise an error if a store
+        # cannot be injected).
         sys.intern("store"),
-        (BaseStore, "BaseStore", inspect.Parameter.empty),
+        (
+            BaseStore,
+            "BaseStore",
+            inspect.Parameter.empty,
+        ),
         CONFIG_KEY_STORE,
         inspect.Parameter.empty,
     ),
+    (
+        # Covers store that is optional. Will set to None if not found in config.
+        sys.intern("store"),
+        (
+            Optional[BaseStore],
+            # Best effort to catch some forward references.
+            # This will not work for cases like `"Union[None, BaseStore]"`,
+            # we'll need to re-write logic to use get_type_hints()
+            # to resolve forward references.
+            "Optional[BaseStore]",
+        ),
+        CONFIG_KEY_STORE,
+        None,
+    ),
     (
         sys.intern("previous"),
         (ANY_TYPE,),
@@ -149,15 +171,24 @@ class RunnableCallable(Runnable):
         params = inspect.signature(cast(Callable, func or afunc)).parameters
 
         self.func_accepts_config = "config" in params
-        self.func_accepts: dict[str, bool] = {}
-        for kw, typ, _, _ in KWARGS_CONFIG_KEYS:
+        # Mapping from kwarg name to (config key, default value) to be used.
+        # The default value is used if the config key is not found in the config.
+        self.func_accepts: dict[str, Tuple[str, Any]] = {}
+
+        for kw, typ, config_key, default in KWARGS_CONFIG_KEYS:
             p = params.get(kw)
-            if typ == (ANY_TYPE,):
-                self.func_accepts[kw] = p is not None and p.kind in VALID_KINDS
-            else:
-                self.func_accepts[kw] = (
-                    p is not None and p.annotation in typ and p.kind in VALID_KINDS
-                )
+
+            if p is None or p.kind not in VALID_KINDS:
+                # If parameter is not found or is not a valid kind, skip
+                continue
+
+            if typ != (ANY_TYPE,) and p.annotation not in typ:
+                # A specific type is required, but the function annotation does
+                # not match the expected type.
+                continue
+
+            # If the kwarg is accepted by the function, store the default value
+            self.func_accepts[kw] = (config_key, default)
 
     def __repr__(self) -> str:
         repr_args = {
@@ -187,20 +218,22 @@ class RunnableCallable(Runnable):
         if self.func_accepts_config:
             kwargs["config"] = config
         _conf = config[CONF]
-        for kw, _, config_key, default_value in KWARGS_CONFIG_KEYS:
-            if not self.func_accepts[kw]:
+
+        for kw, (config_key, default_value) in self.func_accepts.items():
+            # If the kwarg is already set, use the set value
+            if kw in kwargs:
                 continue
 
             if (
-                default_value is inspect.Parameter.empty
-                and kw not in kwargs
-                and config_key not in _conf
+                # If the kwarg is requested, but isn't in the config AND has no
+                # default value, raise an error
+                config_key not in _conf and default_value is inspect.Parameter.empty
             ):
                 raise ValueError(
                     f"Missing required config key '{config_key}' for '{self.name}'."
                 )
-            elif kwargs.get(kw) is None:
-                kwargs[kw] = _conf.get(config_key, default_value)
+
+            kwargs[kw] = _conf.get(config_key, default_value)
 
         context = copy_context()
         if self.trace:
@@ -244,20 +277,20 @@ class RunnableCallable(Runnable):
         if self.func_accepts_config:
             kwargs["config"] = config
         _conf = config[CONF]
-        for kw, _, config_key, default_value in KWARGS_CONFIG_KEYS:
-            if not self.func_accepts[kw]:
+        for kw, (config_key, default_value) in self.func_accepts.items():
+            # If the kwarg has already been set, use the set value
+            if kw in kwargs:
                 continue
 
             if (
-                default_value is inspect.Parameter.empty
-                and kw not in kwargs
-                and config_key not in _conf
+                # If the kwarg is requested, but isn't in the config AND has no
+                # default value, raise an error
+                config_key not in _conf and default_value is inspect.Parameter.empty
             ):
                 raise ValueError(
                     f"Missing required config key '{config_key}' for '{self.name}'."
                 )
-            elif kwargs.get(kw) is None:
-                kwargs[kw] = _conf.get(config_key, default_value)
+            kwargs[kw] = _conf.get(config_key, default_value)
         context = copy_context()
         if self.trace:
             callback_manager = get_async_callback_manager_for_config(config, self.tags)

libs/langgraph/tests/test_pregel.py

@@ -6220,6 +6220,25 @@ def test_entrypoint_with_return_and_save() -> None:
     assert previous_ == ["hello", "goodbye"]
 
 
+def test_overriding_injectable_args_with_tasks() -> None:
+    """Test overriding injectable args in tasks."""
+    from langgraph.store.memory import InMemoryStore
+
+    @task
+    def foo(store: BaseStore, writer: StreamWriter, value: Any) -> None:
+        assert store is value
+        assert writer is value
+
+    @entrypoint(store=InMemoryStore())
+    def main(inputs, store: BaseStore) -> str:
+        assert store is not None
+        foo(store=None, writer=None, value=None).result()
+        foo(store="hello", writer="hello", value="hello").result()
+        return "OK"
+
+    assert main.invoke({}) == "OK"
+
+
 def test_named_tasks_functional() -> None:
     class Foo:
         def foo(self, value: str) -> dict:

libs/langgraph/tests/test_pregel_async.py

@@ -7446,3 +7446,23 @@ async def test_named_tasks_functional() -> None:
         {"qux": "foo|bar|baz|custom_baz|qux"},
         {"workflow": "foo|bar|baz|custom_baz|qux"},
     ]
+
+
+@NEEDS_CONTEXTVARS
+async def test_overriding_injectable_args_with_async_task() -> None:
+    """Test overriding injectable args in tasks."""
+    from langgraph.store.memory import InMemoryStore
+
+    @task
+    async def foo(store: BaseStore, writer: StreamWriter, value: Any) -> None:
+        assert store is value
+        assert writer is value
+
+    @entrypoint(store=InMemoryStore())
+    async def main(inputs, store: BaseStore) -> str:
+        assert store is not None
+        await foo(store=None, writer=None, value=None)
+        await foo(store="hello", writer="hello", value="hello")
+        return "OK"
+
+    assert await main.ainvoke({}) == "OK"

libs/langgraph/tests/test_runnable.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any
+from typing import Any, Optional
 
 import pytest
 
@@ -43,8 +43,15 @@ def test_runnable_callable_func_accepts():
     expected_writer = {"with_writer": True, "awith_writer": True}
 
     for name, runnable in runnables.items():
-        assert runnable.func_accepts["writer"] == expected_writer.get(name, False)
-        assert runnable.func_accepts["store"] == expected_store.get(name, False)
+        if expected_writer.get(name, False):
+            assert "writer" in runnable.func_accepts
+        else:
+            assert "writer" not in runnable.func_accepts
+
+        if expected_store.get(name, False):
+            assert "store" in runnable.func_accepts
+        else:
+            assert "store" not in runnable.func_accepts
 
 
 async def test_runnable_callable_basic():
@@ -63,3 +70,192 @@ async def test_runnable_callable_basic():
     # Test asynchronous ainvoke
     result_async = await runnable_async.ainvoke("test")
     assert result_async == "test"
+
+
+def test_runnable_callable_injectable_arguments() -> None:
+    """Test injectable arguments for RunnableCallable.
+
+    This test verifies that injectable arguments like BaseStore work correctly.
+    It tests:
+    - Optional store injection
+    - Required store injection
+    - Store injection via config
+    - Store injection override behavior
+    - Store value injection and validation
+    """
+
+    # Test Optional[BaseStore] annotation.
+    def func_optional_store(inputs: Any, store: Optional[BaseStore]) -> str:
+        """Test function that accepts an optional store parameter."""
+        assert store is None
+        return "success"
+
+    assert RunnableCallable(func_optional_store).invoke({"x": "1"}) == "success"
+
+    # Test BaseStore annotation
+    def func_required_store(inputs: Any, store: BaseStore) -> str:
+        """Test function that requires a store parameter."""
+        assert store is None
+        return "success"
+
+    with pytest.raises(ValueError):
+        # Should fail b/c store is not Optional and config is not populated with store.
+        assert RunnableCallable(func_required_store).invoke({}) == "success"
+
+    # Manually provide store
+    assert RunnableCallable(func_required_store).invoke({}, store=None) == "success"
+
+    # Specify a value for store in the config
+    assert (
+        RunnableCallable(func_required_store).invoke(
+            {}, config={"configurable": {"__pregel_store": None}}
+        )
+        == "success"
+    )
+
+    # Specify a value for store in config, but override with None
+    assert (
+        RunnableCallable(func_optional_store).invoke(
+            {"x": "1"},
+            store=None,
+            config={"configurable": {"__pregel_store": "foobar"}},
+        )
+        == "success"
+    )
+
+    # Set of tests where we verify that 'foobar' is injected as the store value.
+    def func_required_store_v2(inputs: Any, store: BaseStore) -> str:
+        """Test function that requires a store parameter and validates its value.
+
+        The store value is expected to be 'foobar' when injected.
+        """
+        assert store == "foobar"
+        return "success"
+
+    assert (
+        RunnableCallable(func_required_store_v2).invoke(
+            {}, config={"configurable": {"__pregel_store": "foobar"}}
+        )
+        == "success"
+    )
+
+    assert RunnableCallable(func_required_store_v2).invoke(
+        # And manual override takes precedence.
+        {},
+        store="foobar",
+        config={"configurable": {"__pregel_store": "barbar"}},
+    )
+
+
+async def test_runnable_callable_injectable_arguments_async() -> None:
+    """Test injectable arguments for async RunnableCallable.
+
+    This test verifies that injectable arguments like BaseStore work correctly
+    in the async context. It tests:
+    - Optional store injection
+    - Required store injection
+    - Store injection via config
+    - Store injection override behavior
+    """
+
+    # Test Optional[BaseStore] annotation.
+    def func_optional_store(inputs: Any, store: Optional[BaseStore]) -> str:
+        """Test function that accepts an optional store parameter."""
+        assert store is None
+        return "success"
+
+    async def afunc_optional_store(inputs: Any, store: Optional[BaseStore]) -> str:
+        """Async version of func_optional_store."""
+        assert store is None
+        return "success"
+
+    assert (
+        await RunnableCallable(
+            func=func_optional_store, afunc=afunc_optional_store
+        ).ainvoke({"x": "1"})
+        == "success"
+    )
+
+    # Test BaseStore annotation
+    def func_required_store(inputs: Any, store: BaseStore) -> str:
+        """Test function that requires a store parameter."""
+        assert store is None
+        return "success"
+
+    async def afunc_required_store(inputs: Any, store: BaseStore) -> str:
+        """Async version of func_required_store."""
+        assert store is None
+        return "success"
+
+    with pytest.raises(ValueError):
+        # Should fail b/c store is not Optional and config is not populated with store.
+        assert (
+            await RunnableCallable(
+                func=func_required_store, afunc=afunc_required_store
+            ).ainvoke({})
+            == "success"
+        )
+
+    # Manually provide store
+    assert (
+        await RunnableCallable(
+            func=func_required_store, afunc=afunc_required_store
+        ).ainvoke({}, store=None)
+        == "success"
+    )
+
+    # Specify a value for store in the config
+    assert (
+        await RunnableCallable(
+            func=func_required_store, afunc=afunc_required_store
+        ).ainvoke({}, config={"configurable": {"__pregel_store": None}})
+        == "success"
+    )
+
+    # Specify a value for store in config, but override with None
+    assert (
+        await RunnableCallable(
+            func=func_optional_store, afunc=afunc_optional_store
+        ).ainvoke(
+            {"x": "1"},
+            store=None,
+            config={"configurable": {"__pregel_store": "foobar"}},
+        )
+        == "success"
+    )
+
+    # Set of tests where we verify that 'foobar' is injected as the store value.
+    def func_required_store_v2(inputs: Any, store: BaseStore) -> str:
+        """Test function that requires a store parameter with specific value.
+
+        The store parameter is expected to be 'foobar' when injected.
+        """
+        assert store == "foobar"
+        return "success"
+
+    async def afunc_required_store_v2(inputs: Any, store: BaseStore) -> str:
+        """Async version of func_required_store_v2.
+
+        The store parameter is expected to be 'foobar' when injected.
+        """
+        assert store == "foobar"
+        return "success"
+
+    assert (
+        await RunnableCallable(
+            func=func_required_store_v2, afunc=afunc_required_store_v2
+        ).ainvoke({}, config={"configurable": {"__pregel_store": "foobar"}})
+        == "success"
+    )
+
+    assert (
+        await RunnableCallable(
+            func=func_required_store_v2, afunc=afunc_required_store_v2
+        ).ainvoke(
+            # And manual override takes precedence.
+            {},
+            store="foobar",
+            config={"configurable": {"__pregel_store": "barbar"}},
+        )
+        == "success"
+    )

libs/langgraph/tests/test_utils.py

@@ -26,7 +26,10 @@ from langgraph.utils.fields import (
     get_enhanced_type_hints,
     get_field_default,
 )
-from langgraph.utils.runnable import is_async_callable, is_async_generator
+from langgraph.utils.runnable import (
+    is_async_callable,
+    is_async_generator,
+)
 
 pytestmark = pytest.mark.anyio