diff --git a/crates/graph/src/lib.rs b/crates/graph/src/lib.rs
index c069870..b198f5f 100644
--- a/crates/graph/src/lib.rs
+++ b/crates/graph/src/lib.rs
@@ -58,7 +58,7 @@ impl<O: 'static, S: Synchronicity> Graph<O, S> {
         self.node_graph.borrow().node_count()
     }
 
-    pub fn modify<F>(&mut self, mut f: F) -> Result<(), BuildGraphError>
+    pub fn modify<F>(mut self, mut f: F) -> Result<Self, BuildGraphError>
     where
         F: FnMut(&mut GraphBuilder<O, S>) -> (),
     {
@@ -81,14 +81,13 @@ impl<O: 'static, S: Synchronicity> Graph<O, S> {
         let old_output = self.output.node_idx;
 
         let mut graph = GraphBuilder {
-            // TODO: is using the same graph as self correct? if the modify fails, is it left in a bad state?
             node_graph: Rc::clone(&self.node_graph),
             output: Some(self.output.clone()),
             output_type: std::marker::PhantomData,
             is_valid: Rc::clone(&self.is_valid),
         };
         f(&mut graph);
-        *self = graph.build()?;
+        self = graph.build()?;
 
         // Any new inboud edges invalidate their target nodes.
         let mut graph = self.node_graph.borrow_mut();
@@ -113,7 +112,8 @@ impl<O: 'static, S: Synchronicity> Graph<O, S> {
             }
         }
 
-        Ok(())
+        drop(graph);
+        Ok(self)
     }
 }
 
@@ -427,7 +427,7 @@ mod tests {
         builder.set_output(ConstantRule(1));
         let mut graph = builder.build().unwrap();
         assert_eq!(*graph.evaluate(), 1);
-        graph
+        graph = graph
             .modify(|g| {
                 g.set_output(ConstantRule(2));
             })
@@ -444,7 +444,7 @@ mod tests {
         *input.borrow_mut() = Some(builder.add_value(1));
         let mut graph = builder.build().unwrap();
         assert_eq!(*graph.evaluate(), 1);
-        graph
+        graph = graph
             .modify(|g| {
                 *input.borrow_mut() = Some(g.add_value(2));
             })