From da935e78054a735b18d24c796b676655ca3452be Mon Sep 17 00:00:00 2001
From: Vadym Barda <vadym@langchain.dev>
Date: Mon, 30 Sep 2024 18:00:33 -0400
Subject: [PATCH] langgraph: fix edge case with string enums as node names
 (#1926)

---
 libs/langgraph/langgraph/pregel/algo.py |  3 ++-
 libs/langgraph/tests/test_pregel.py     | 21 +++++++++++++++++++++
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py
index daa942a39..98ac8576f 100644
--- a/libs/langgraph/langgraph/pregel/algo.py
+++ b/libs/langgraph/langgraph/pregel/algo.py
@@ -13,6 +13,7 @@ from typing import (
     Protocol,
     Sequence,
     Union,
+    cast,
     overload,
 )
 from uuid import UUID
@@ -471,7 +472,7 @@ def prepare_single_task(
         else:
             return PregelTask(task_id, packet.node, task_path)
     elif task_path[0] == PULL:
-        name = str(task_path[1])
+        name = cast(str, task_path[1])
         if name not in processes:
             return
         proc = processes[name]
diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py
index 0bd5772d4..29e287da2 100644
--- a/libs/langgraph/tests/test_pregel.py
+++ b/libs/langgraph/tests/test_pregel.py
@@ -1,3 +1,4 @@
+import enum
 import json
 import operator
 import re
@@ -11517,3 +11518,23 @@ def test_store_injected(
         "some_val": 0,
     }  # Overwrites the whole doc
     assert len(the_store.search(("foo", "bar"))) == 1  # still overwriting the same one
+
+
+def test_enum_node_names():
+    class NodeName(str, enum.Enum):
+        BAZ = "baz"
+
+    class State(TypedDict):
+        foo: str
+        bar: str
+
+    def baz(state: State):
+        return {"bar": state["foo"] + "!"}
+
+    graph = StateGraph(State)
+    graph.add_node(NodeName.BAZ, baz)
+    graph.add_edge(START, NodeName.BAZ)
+    graph.add_edge(NodeName.BAZ, END)
+    graph = graph.compile()
+
+    assert graph.invoke({"foo": "hello"}) == {"foo": "hello", "bar": "hello!"}