diff --git a/src/backend/base/langflow/components/prototypes/Listen.py b/src/backend/base/langflow/components/prototypes/Listen.py
index 23b9c3ee1..6e5de723a 100644
--- a/src/backend/base/langflow/components/prototypes/Listen.py
+++ b/src/backend/base/langflow/components/prototypes/Listen.py
@@ -18,5 +18,11 @@ class ListenComponent(CustomComponent):
 
     def build(self, name: str) -> Data:
         state = self.get_state(name)
+        self._set_successors_ids()
         self.status = state
         return state
+
+    def _set_successors_ids(self):
+        self.vertex.is_state = True
+        successors = self.vertex.graph.successor_map.get(self.vertex.id, [])
+        return successors + self.vertex.graph.activated_vertices
diff --git a/src/backend/base/langflow/components/prototypes/Notify.py b/src/backend/base/langflow/components/prototypes/Notify.py
index c07baae34..b83331e3a 100644
--- a/src/backend/base/langflow/components/prototypes/Notify.py
+++ b/src/backend/base/langflow/components/prototypes/Notify.py
@@ -39,4 +39,10 @@ class NotifyComponent(CustomComponent):
         else:
             self.status = "No record provided."
         self.status = data
+        self._set_successors_ids()
         return data
+
+    def _set_successors_ids(self):
+        self.vertex.is_state = True
+        successors = self.vertex.graph.successor_map.get(self.vertex.id, [])
+        return successors + self.vertex.graph.activated_vertices
diff --git a/src/backend/base/langflow/graph/vertex/base.py b/src/backend/base/langflow/graph/vertex/base.py
index 6d1d27f4e..1fc6a08cc 100644
--- a/src/backend/base/langflow/graph/vertex/base.py
+++ b/src/backend/base/langflow/graph/vertex/base.py
@@ -71,6 +71,7 @@ class Vertex:
         self._built_object = UnbuiltObject()
         self._built_result = None
         self._built = False
+        self._successors_ids: Optional[List[str]] = None
         self.artifacts: Dict[str, Any] = {}
         self.artifacts_raw: Dict[str, Any] = {}
         self.artifacts_type: Dict[str, str] = {}
diff --git a/src/backend/base/langflow/graph/vertex/types.py b/src/backend/base/langflow/graph/vertex/types.py
index cd330aa0a..78c541f65 100644
--- a/src/backend/base/langflow/graph/vertex/types.py
+++ b/src/backend/base/langflow/graph/vertex/types.py
@@ -404,16 +404,18 @@ class InterfaceVertex(ComponentVertex):
         return self.vertex_type == InterfaceComponentTypes.ChatInput and self.is_input
 
 
-class StateVertex(Vertex):
+class StateVertex(ComponentVertex):
     def __init__(self, data: Dict, graph):
-        super().__init__(data, graph=graph, base_type="custom_components")
+        super().__init__(data, graph=graph)
         self.steps = [self._build]
-        self.is_state = True
+        self.is_state = False
 
     @property
     def successors_ids(self) -> List[str]:
-        successors = self.graph.successor_map.get(self.id, [])
-        return successors + self.graph.activated_vertices
+        if self._successors_ids is None:
+            self.is_state = False
+            return super().successors_ids
+        return self._successors_ids
 
     def _built_object_repr(self):
         if self.artifacts and "repr" in self.artifacts: