feat(langgraph): push_messages should directly write to the state (#4791)

libs/langgraph/langgraph/graph/message.py

@@ -23,6 +23,7 @@ from langchain_core.messages import (
 )
 from typing_extensions import TypedDict
 
+from langgraph.constants import CONF, CONFIG_KEY_SEND
 from langgraph.graph.state import StateGraph
 
 Messages = Union[list[MessageLikeRepresentation], MessageLikeRepresentation]
@@ -298,8 +299,13 @@ def _format_messages(messages: Sequence[BaseMessage]) -> list[BaseMessage]:
 
 def push_message(
     message: Union[MessageLikeRepresentation, BaseMessageChunk],
+    *,
+    state_key: Optional[str] = "messages",
 ) -> AnyMessage:
-    """Write a message manually to the `messages` / `messages-tuple` stream mode."""
+    """Write a message manually to the `messages` / `messages-tuple` stream mode.
+
+    Will automatically write to the channel specified in the `state_key` unless `state_key` is `None`.
+    """
 
     from langchain_core.callbacks.base import (
         BaseCallbackHandler,
@@ -334,4 +340,7 @@ def push_message(
         )
         stream_handler._emit(message_meta, message, dedupe=False)
 
+    if state_key:
+        config[CONF][CONFIG_KEY_SEND]([(state_key, message)])
+
     return message

libs/langgraph/langgraph/graph/ui.py

@@ -54,7 +54,7 @@ def push_ui_message(
     id: Optional[str] = None,
     metadata: Optional[dict[str, Any]] = None,
     message: Optional[AnyMessage] = None,
-    state_key: str = "ui",
+    state_key: Optional[str] = "ui",
     merge: bool = False,
 ) -> UIMessage:
     """Push a new UI message to update the UI state.
@@ -111,7 +111,8 @@ def push_ui_message(
     }
 
     writer(evt)
-    config[CONF][CONFIG_KEY_SEND]([(state_key, evt)])
+    if state_key:
+        config[CONF][CONFIG_KEY_SEND]([(state_key, evt)])
 
     return evt
 

libs/langgraph/tests/test_messages_state.py

@@ -342,13 +342,9 @@ def test_push_messages_in_graph():
         with pytest.raises(ValueError, match="Message ID is required"):
             push_message(AIMessage(content="No ID"))
 
-        return {
-            "messages": [
-                push_message(AIMessage(content="First", id="1")),
-                push_message(HumanMessage(content="Second", id="2")),
-                push_message(AIMessage(content="Third", id="3")),
-            ]
-        }
+        push_message(AIMessage(content="First", id="1"))
+        push_message(HumanMessage(content="Second", id="2"))
+        push_message(AIMessage(content="Third", id="3"))
 
     builder = StateGraph(MessagesState)
     builder.add_node(chat)