From 1e6d9584344722455cae76e002c389bdd00bc330 Mon Sep 17 00:00:00 2001
From: Vadym Barda <vadym@langchain.dev>
Date: Tue, 18 Feb 2025 18:53:11 -0500
Subject: [PATCH] langgraph: fallback on no-op writer in get_stream_writer
 (#3488)

---
 libs/langgraph/langgraph/config.py  |  7 ++++++-
 libs/langgraph/tests/test_pregel.py | 32 +++++++++++++++++++++++++++++
 2 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/libs/langgraph/langgraph/config.py b/libs/langgraph/langgraph/config.py
index df2b637b7..b2ef57cfb 100644
--- a/libs/langgraph/langgraph/config.py
+++ b/libs/langgraph/langgraph/config.py
@@ -1,5 +1,6 @@
 import asyncio
 import sys
+from typing import Any
 
 from langchain_core.runnables import RunnableConfig
 from langchain_core.runnables.config import var_child_runnable_config
@@ -9,6 +10,10 @@ from langgraph.store.base import BaseStore
 from langgraph.types import StreamWriter
 
 
+def _no_op_stream_writer(c: Any) -> None:
+    pass
+
+
 def get_config() -> RunnableConfig:
     if sys.version_info < (3, 11):
         try:
@@ -177,4 +182,4 @@ def get_stream_writer() -> StreamWriter:
         ```
     """
     config = get_config()
-    return config[CONF][CONFIG_KEY_STREAM_WRITER]
+    return config[CONF].get(CONFIG_KEY_STREAM_WRITER, _no_op_stream_writer)
diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py
index 7e15db62f..992b4664f 100644
--- a/libs/langgraph/tests/test_pregel.py
+++ b/libs/langgraph/tests/test_pregel.py
@@ -54,6 +54,7 @@ from langgraph.checkpoint.base import (
     CheckpointTuple,
 )
 from langgraph.checkpoint.memory import InMemorySaver, MemorySaver
+from langgraph.config import get_stream_writer
 from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, START
 from langgraph.errors import InvalidUpdateError
 from langgraph.func import entrypoint, task
@@ -6495,3 +6496,34 @@ def test_pydantic_none_state_update() -> None:
 
     graph = StateGraph(State).add_node(node_a).add_edge(START, "node_a").compile()
     assert graph.invoke({"foo": ""}) == {"foo": None}
+
+
+def test_get_stream_writer() -> None:
+    class State(TypedDict):
+        foo: str
+
+    def my_node(state):
+        writer = get_stream_writer()
+        writer("custom!")
+        return state
+
+    graph = StateGraph(State).add_node(my_node).add_edge(START, "my_node").compile()
+    assert list(graph.stream({"foo": "bar"}, stream_mode="custom")) == ["custom!"]
+    assert list(graph.stream({"foo": "bar"}, stream_mode="values")) == [
+        {"foo": "bar"},
+        {"foo": "bar"},
+    ]
+    assert list(graph.stream({"foo": "bar"}, stream_mode=["custom", "updates"])) == [
+        (
+            "custom",
+            "custom!",
+        ),
+        (
+            "updates",
+            {
+                "my_node": {
+                    "foo": "bar",
+                },
+            },
+        ),
+    ]