fix signature

libs/langgraph/langgraph/graph/state.py

@@ -148,6 +148,15 @@ class _NodeWithConfigWriterStore(Protocol[StateT_contra]):
     ) -> Any: ...
 
 
+class _Invokable(Protocol[StateT_contra]):
+    def invoke(
+        self,
+        input: StateT_contra,
+        config: RunnableConfig | None = None,
+        **kwargs: Any,
+    ) -> Any: ...
+
+
 # TODO: we probably don't want to explicitly support the config / store signatures once
 # we move to adding a context arg. Maybe what we do is we add support for kwargs with param spec
 # this is purely for typing purposes though, so can easily change in the coming weeks.
@@ -160,7 +169,7 @@ StateNode: TypeAlias = Union[
     _NodeWithConfigWriter[StateT_contra],
     _NodeWithConfigStore[StateT_contra],
     _NodeWithConfigWriterStore[StateT_contra],
-    Runnable[StateT_contra, Any],
+    _Invokable[StateT_contra],
 ]
 
 
@@ -536,7 +545,7 @@ class StateGraph(Generic[StateT, InputT, OutputT]):
         if input_schema is not None:
             self._add_schema(input_schema)
         self.nodes[node] = StateNodeSpec(
-            coerce_to_runnable(action, name=node, trace=False),
+            coerce_to_runnable(action, name=node, trace=False),  # type: ignore[arg-type]
             metadata,
             input=input_schema or self.state_schema,
             retry_policy=retry_policy,

libs/langgraph/langgraph/typing.py

@@ -1,7 +1,5 @@
 from __future__ import annotations
 
-from typing import Union
-
 from typing_extensions import TypeVar
 
 from langgraph._typing import StateLike
@@ -19,12 +17,8 @@ InputT = TypeVar("InputT", bound=StateLike, default=StateT)
 Defaults to `StateT`.
 """
 
-ResolvedInputT = TypeVar("ResolvedInputT", bound=StateLike)
-"""Type variable used to represent the resolved input to a state graph.
+OutputT = TypeVar("OutputT", bound=StateLike, default=StateT)
+"""Type variable used to represent the output of a state graph.
 
-No default.
+Defaults to `StateT`.
 """
-
-
-OutputT = TypeVar("OutputT", bound=Union[StateLike, None], default=StateT)
-"""Type variable used to represent the output of a state graph."""

libs/langgraph/tests/test_type_checking.py

@@ -1,6 +1,6 @@
 from dataclasses import dataclass
 from operator import add
-from typing import Annotated, Any
+from typing import Annotated, Any, Union
 
 from langchain_core.runnables import RunnableConfig
 from pydantic import BaseModel
@@ -103,3 +103,21 @@ def test_input_state_specified() -> None:
 
     new_graph.invoke({"something": 1})
     new_graph.invoke({"something": 2, "info": ["hello", "world"]})  # type: ignore[arg-type]
+
+
+def test_invokeable_node_signature() -> None:
+    class State(TypedDict):
+        info: Annotated[list[str], add]
+
+    graph_builder = StateGraph(State)
+
+    class RunnableIsh:
+        def invoke(
+            self,
+            input: State,
+            config: Union[RunnableConfig, None] = None,
+            **kwargs: Any,
+        ) -> dict[str, str]:
+            return {}
+
+    graph_builder.add_node("runnable", RunnableIsh())