diff --git a/tests/__snapshots__/test_pregel.ambr b/tests/__snapshots__/test_pregel.ambr
index b00d086ab..b68acaf76 100644
--- a/tests/__snapshots__/test_pregel.ambr
+++ b/tests/__snapshots__/test_pregel.ambr
@@ -97,18 +97,6 @@
         "source": "left",
         "target": "condition"
       },
-      {
-        "source": "condition",
-        "target": "left",
-        "data": "left",
-        "conditional": true
-      },
-      {
-        "source": "condition",
-        "target": "right",
-        "data": "right",
-        "conditional": true
-      },
       {
         "source": "condition",
         "target": "__end__",
@@ -121,39 +109,33 @@
 # ---
 # name: test_conditional_entrypoint_graph.3
   '''
-            +-----------+          
-            | __start__ |          
-            +-----------+          
-                   *               
-                   *               
-                   *               
-           +--------------+        
-           | should_start |        
-           +--------------+        
-             ..         ..         
-           ..             ..       
-          .                 ..     
-    +------+                  .    
-    | left |                  .    
-    +------+                  .    
-        *                     .    
-        *                     .    
-        *                     .    
-  +-----------+               .    
-  | condition |               .    
-  +-----------+...            .    
-        .         ...         .    
-        .            ....     .    
-        .                ..   .    
-        .                +-------+ 
-         ..              | right | 
-           ..            +-------+ 
-             ..         **         
-               ..     **           
-                 .   *             
-              +---------+          
-              | __end__ |          
-              +---------+          
+            +-----------+        
+            | __start__ |        
+            +-----------+        
+                  *              
+                  *              
+                  *              
+          +--------------+       
+          | should_start |       
+          +--------------+       
+             .         ..        
+           ..            .       
+          .               ..     
+    +------+                .    
+    | left |                .    
+    +------+                .    
+        *                   .    
+        *                   .    
+        *                   .    
+  +-----------+        +-------+ 
+  | condition |        | right | 
+  +-----------+        +-------+ 
+             .         *         
+              ..     **          
+                .   *            
+            +---------+          
+            | __end__ |          
+            +---------+          
   '''
 # ---
 # name: test_conditional_entrypoint_graph_state
@@ -254,18 +236,6 @@
         "source": "left",
         "target": "condition"
       },
-      {
-        "source": "condition",
-        "target": "left",
-        "data": "left",
-        "conditional": true
-      },
-      {
-        "source": "condition",
-        "target": "right",
-        "data": "right",
-        "conditional": true
-      },
       {
         "source": "condition",
         "target": "__end__",
@@ -278,39 +248,33 @@
 # ---
 # name: test_conditional_entrypoint_graph_state.3
   '''
-            +-----------+          
-            | __start__ |          
-            +-----------+          
-                   *               
-                   *               
-                   *               
-           +--------------+        
-           | should_start |        
-           +--------------+        
-             ..         ..         
-           ..             ..       
-          .                 ..     
-    +------+                  .    
-    | left |                  .    
-    +------+                  .    
-        *                     .    
-        *                     .    
-        *                     .    
-  +-----------+               .    
-  | condition |               .    
-  +-----------+...            .    
-        .         ...         .    
-        .            ....     .    
-        .                ..   .    
-        .                +-------+ 
-         ..              | right | 
-           ..            +-------+ 
-             ..         **         
-               ..     **           
-                 .   *             
-              +---------+          
-              | __end__ |          
-              +---------+          
+            +-----------+        
+            | __start__ |        
+            +-----------+        
+                  *              
+                  *              
+                  *              
+          +--------------+       
+          | should_start |       
+          +--------------+       
+             .         ..        
+           ..            .       
+          .               ..     
+    +------+                .    
+    | left |                .    
+    +------+                .    
+        *                   .    
+        *                   .    
+        *                   .    
+  +-----------+        +-------+ 
+  | condition |        | right | 
+  +-----------+        +-------+ 
+             .         *         
+              ..     **          
+                .   *            
+            +---------+          
+            | __end__ |          
+            +---------+          
   '''
 # ---
 # name: test_conditional_graph[end_of_run]
@@ -807,9 +771,9 @@
             +-------+           
             | agent |           
             +-------+           
-           *         .          
-         **           ..        
-        *               .       
+           .         .          
+         ..           ..        
+        .               .       
   +-------+         +---------+ 
   | tools |         | __end__ | 
   +-------+         +---------+ 
diff --git a/tests/test_pregel.py b/tests/test_pregel.py
index 1d84f3199..3b0c40c4f 100644
--- a/tests/test_pregel.py
+++ b/tests/test_pregel.py
@@ -1367,7 +1367,7 @@ def test_conditional_entrypoint_graph(snapshot: SnapshotAssertion) -> None:
         should_start, {"go-left": "left", "go-right": "right"}
     )
 
-    workflow.add_conditional_edges("left", lambda data: END)
+    workflow.add_conditional_edges("left", lambda data: END, {END: END})
     workflow.add_edge("right", END)
 
     app = workflow.compile()
@@ -1819,7 +1819,7 @@ def test_conditional_entrypoint_graph_state(snapshot: SnapshotAssertion) -> None
         should_start, {"go-left": "left", "go-right": "right"}
     )
 
-    workflow.add_conditional_edges("left", lambda data: END)
+    workflow.add_conditional_edges("left", lambda data: END, {END: END})
     workflow.add_edge("right", END)
 
     app = workflow.compile()