Handle Command returned from node (in addition to GraphCommand)

libs/langgraph/langgraph/graph/state.py

@@ -829,15 +829,16 @@ def _coerce_state(schema: Type[Any], input: dict[str, Any]) -> dict[str, Any]:
 def _control_branch(value: Any) -> Sequence[Union[str, Send]]:
     if isinstance(value, Send):
         return [value]
-    if not isinstance(value, GraphCommand):
+    if not isinstance(value, Command):
         return EMPTY_SEQ
     if value.graph == Command.PARENT:
         raise ParentCommand(value)
     rtn: list[Union[str, Send]] = []
-    if isinstance(value.goto, str):
-        rtn.append(value.goto)
-    else:
-        rtn.extend(value.goto)
+    if isinstance(value, GraphCommand):
+        if isinstance(value.goto, str):
+            rtn.append(value.goto)
+        else:
+            rtn.extend(value.goto)
     if isinstance(value.send, Send):
         rtn.append(value.send)
     else:
@@ -853,10 +854,11 @@ async def _acontrol_branch(value: Any) -> Sequence[Union[str, Send]]:
     if value.graph == Command.PARENT:
         raise ParentCommand(value)
     rtn: list[Union[str, Send]] = []
-    if isinstance(value.goto, str):
-        rtn.append(value.goto)
-    else:
-        rtn.extend(value.goto)
+    if isinstance(value, GraphCommand):
+        if isinstance(value.goto, str):
+            rtn.append(value.goto)
+        else:
+            rtn.extend(value.goto)
     if isinstance(value.send, Send):
         rtn.append(value.send)
     else:

libs/langgraph/tests/test_pregel.py

@@ -1925,7 +1925,7 @@ def test_send_sequences() -> None:
 
     def send_for_fun(state):
         return [
-            Send("2", GraphCommand(send=Send("2", 3))),
+            Send("2", Command(send=Send("2", 3))),
             Send("2", GraphCommand(send=Send("2", 4))),
             "3.1",
         ]

libs/langgraph/tests/test_pregel_async.py

@@ -2573,14 +2573,14 @@ async def test_send_sequences(checkpointer_name: str) -> None:
                 if isinstance(state, list)  # or isinstance(state, Control)
                 else ["|".join((self.name, str(state)))]
             )
-            if isinstance(state, GraphCommand):
+            if isinstance(state, Command):
                 return replace(state, update=update)
             else:
                 return update
 
     async def send_for_fun(state):
         return [
-            Send("2", GraphCommand(send=Send("2", 3))),
+            Send("2", Command(send=Send("2", 3))),
             Send("2", GraphCommand(send=Send("2", 4))),
             "3.1",
         ]