langgraph: fix edge case with string enums as node names

libs/langgraph/langgraph/pregel/algo.py

@@ -1,4 +1,4 @@
 from collections import defaultdict, deque
 from functools import partial
 from hashlib import sha1
 from typing import (
@@ -13,6 +13,7 @@
     Protocol,
     Sequence,
     Union,
+    cast,
     overload,
 )
 from uuid import UUID
@@ -471,7 +472,7 @@
         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]

libs/langgraph/tests/test_pregel.py

@@ -1,3 +1,4 @@
+import enum
 import json
 import operator
 import re
@@ -11517,3 +11518,23 @@ def node(input: State, config: RunnableConfig, store: BaseStore):
         "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!"}