From 6eace78c53d5a34998ae247d470ec8e71fb14af9 Mon Sep 17 00:00:00 2001
From: Sydney Runkle <54324534+sydney-runkle@users.noreply.github.com>
Date: Wed, 9 Jul 2025 13:23:28 -0400
Subject: [PATCH] patch[langgraph]: Fix hint for `invoke`/`stream` to allow for
 `Command` and `None` (#5414)

use Command and None as well
---
 libs/langgraph/langgraph/pregel/__init__.py | 9 +++++----
 libs/langgraph/langgraph/pregel/protocol.py | 9 +++++----
 2 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py
index 2397c1978..21d4d593b 100644
--- a/libs/langgraph/langgraph/pregel/__init__.py
+++ b/libs/langgraph/langgraph/pregel/__init__.py
@@ -99,6 +99,7 @@ from langgraph.types import (
     All,
     CachePolicy,
     Checkpointer,
+    Command,
     Interrupt,
     Send,
     StateSnapshot,
@@ -2342,7 +2343,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
 
     def stream(
         self,
-        input: InputT,
+        input: InputT | Command | None,
         config: RunnableConfig | None = None,
         *,
         stream_mode: StreamMode | Sequence[StreamMode] | None = None,
@@ -2564,7 +2565,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
 
     async def astream(
         self,
-        input: InputT,
+        input: InputT | Command | None,
         config: RunnableConfig | None = None,
         *,
         stream_mode: StreamMode | Sequence[StreamMode] | None = None,
@@ -2808,7 +2809,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
 
     def invoke(
         self,
-        input: InputT,
+        input: InputT | Command | None,
         config: RunnableConfig | None = None,
         *,
         stream_mode: StreamMode = "values",
@@ -2883,7 +2884,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
 
     async def ainvoke(
         self,
-        input: InputT,
+        input: InputT | Command | None,
         config: RunnableConfig | None = None,
         *,
         stream_mode: StreamMode = "values",
diff --git a/libs/langgraph/langgraph/pregel/protocol.py b/libs/langgraph/langgraph/pregel/protocol.py
index e55be3cbc..dce0ef870 100644
--- a/libs/langgraph/langgraph/pregel/protocol.py
+++ b/libs/langgraph/langgraph/pregel/protocol.py
@@ -9,6 +9,7 @@ from langchain_core.runnables.graph import Graph as DrawableGraph
 from typing_extensions import Self
 
 from langgraph.pregel.types import All, StateSnapshot, StateUpdate, StreamMode
+from langgraph.types import Command
 from langgraph.typing import InputT, OutputT, StateT
 
 
@@ -98,7 +99,7 @@ class PregelProtocol(Runnable[InputT, Any], Generic[StateT, InputT, OutputT], AB
     @abstractmethod
     def stream(
         self,
-        input: InputT,
+        input: InputT | Command | None,
         config: RunnableConfig | None = None,
         *,
         stream_mode: StreamMode | list[StreamMode] | None = None,
@@ -110,7 +111,7 @@ class PregelProtocol(Runnable[InputT, Any], Generic[StateT, InputT, OutputT], AB
     @abstractmethod
     def astream(
         self,
-        input: InputT,
+        input: InputT | Command | None,
         config: RunnableConfig | None = None,
         *,
         stream_mode: StreamMode | list[StreamMode] | None = None,
@@ -122,7 +123,7 @@ class PregelProtocol(Runnable[InputT, Any], Generic[StateT, InputT, OutputT], AB
     @abstractmethod
     def invoke(
         self,
-        input: InputT,
+        input: InputT | Command | None,
         config: RunnableConfig | None = None,
         *,
         interrupt_before: All | Sequence[str] | None = None,
@@ -132,7 +133,7 @@ class PregelProtocol(Runnable[InputT, Any], Generic[StateT, InputT, OutputT], AB
     @abstractmethod
     async def ainvoke(
         self,
-        input: InputT,
+        input: InputT | Command | None,
         config: RunnableConfig | None = None,
         *,
         interrupt_before: All | Sequence[str] | None = None,