From 2fe38f394021c2e3f4d7169c791a1ee8d84054b9 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Thu, 21 Nov 2024 09:07:19 -0800
Subject: [PATCH] Fix get_state

---
 libs/langgraph/langgraph/func/__init__.py | 1 +
 libs/langgraph/tests/test_pregel.py       | 2 ++
 2 files changed, 3 insertions(+)

diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py
index 4db234b81..d990ae4a5 100644
--- a/libs/langgraph/langgraph/func/__init__.py
+++ b/libs/langgraph/langgraph/func/__init__.py
@@ -71,6 +71,7 @@ def entrypoint(
             channels={START: EphemeralValue(Any), END: LastValue(Any, END)},
             input_channels=START,
             output_channels=END,
+            stream_channels=END,
             stream_mode="updates",
             checkpointer=checkpointer,
             store=store,
diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py
index 8de87ef0e..cb7998430 100644
--- a/libs/langgraph/tests/test_pregel.py
+++ b/libs/langgraph/tests/test_pregel.py
@@ -2046,6 +2046,8 @@ def test_imp_stream_order(
         {"graph": {"a": "0foobarbaz", "c": "something else"}},
     ]
 
+    assert graph.get_state(thread1).values == {"a": "0foobarbaz", "c": "something else"}
+
 
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
 def test_send_dedupe_on_resume(