diff --git a/src/backend/langflow/graph/edge/base.py b/src/backend/langflow/graph/edge/base.py
index bd585c168..82714e395 100644
--- a/src/backend/langflow/graph/edge/base.py
+++ b/src/backend/langflow/graph/edge/base.py
@@ -72,8 +72,8 @@ class Edge:
         self.source = state["source"]
         self.target = state["target"]
         self.target_param = state["target_param"]
-        self.source_handle = state["source_handle"]
-        self.target_handle = state["target_handle"]
+        self.source_handle = state.get("source_handle")
+        self.target_handle = state.get("target_handle")
 
     def reset(self) -> None:
         self.source._build_params()
diff --git a/src/backend/langflow/graph/vertex/base.py b/src/backend/langflow/graph/vertex/base.py
index c31cca3ec..8ae0bb5f8 100644
--- a/src/backend/langflow/graph/vertex/base.py
+++ b/src/backend/langflow/graph/vertex/base.py
@@ -38,6 +38,8 @@ class Vertex:
         self.task_id: Optional[str] = None
         self.is_task = is_task
         self.params = params or {}
+        self.parent_node_id: Optional[str] = self._data.get("parent_node_id")
+        self.parent_is_top_level = False
 
     def reset_params(self):
         for edge in self.edges:
@@ -88,8 +90,8 @@ class Vertex:
             self._built = False
         self.artifacts: Dict[str, Any] = {}
         self.task_id: Optional[str] = None
-        self.parent_node_id: Optional[str] = self._data.get("parent_node_id")
-        self.parent_is_top_level = False
+        self.parent_node_id = state["parent_node_id"]
+        self.parent_is_top_level = state["parent_is_top_level"]
 
     def set_top_level(self, top_level_nodes: List[str]) -> None:
         self.parent_is_top_level = self.parent_node_id in top_level_nodes
diff --git a/tests/test_graph.py b/tests/test_graph.py
index 0dc9fc372..8d442ce76 100644
--- a/tests/test_graph.py
+++ b/tests/test_graph.py
@@ -497,7 +497,7 @@ def test_update_source_handle():
         "edges": [{"source": "some_node"}],
     }
     updated_edge = update_source_handle(
-        new_edge, {"nodes": flow_data["nodes"], "edges": flow_data["edges"]}
+        new_edge, flow_data["nodes"], flow_data["edges"]
     )
     assert updated_edge["source"] == "last_node"
     assert updated_edge["data"]["sourceHandle"]["id"] == "last_node"