diff --git a/src/hugr/rewrite/replace.rs b/src/hugr/rewrite/replace.rs
index 1e9e9f3a4..863965376 100644
--- a/src/hugr/rewrite/replace.rs
+++ b/src/hugr/rewrite/replace.rs
@@ -521,7 +521,9 @@ mod test {
                 inputs: vec![listy.clone()].into(),
                 tuple_sum_rows: vec![type_row![]],
                 other_outputs: vec![listy.clone()].into(),
-                extension_delta: ExtensionSet::singleton(&collections::EXTENSION_NAME),
+                // This should be ExtensionSet::singleton(&collections::EXTENSION_NAME),
+                // at least when https://github.com/CQCL/issues/388 is fixed
+                extension_delta: ExtensionSet::new(),
             },
         )?;
         let r_df1 = replacement.add_node_with_parent(
diff --git a/src/ops/controlflow.rs b/src/ops/controlflow.rs
index c28b3fab0..0214ec15f 100644
--- a/src/ops/controlflow.rs
+++ b/src/ops/controlflow.rs
@@ -4,7 +4,7 @@ use smol_str::SmolStr;
 
 use crate::extension::ExtensionSet;
 use crate::types::{EdgeKind, FunctionType, Type, TypeRow};
-use crate::{type_row, Direction};
+use crate::Direction;
 
 use super::dataflow::{DataflowOpTrait, DataflowParent};
 use super::OpTag;
@@ -179,10 +179,8 @@ impl OpTrait for DataflowBlock {
         Some(EdgeKind::ControlFlow)
     }
 
-    fn dataflow_signature(&self) -> Option<FunctionType> {
-        Some(
-            FunctionType::new(type_row![], type_row![]).with_extension_delta(&self.extension_delta),
-        )
+    fn extension_delta(&self) -> ExtensionSet {
+        self.extension_delta.clone()
     }
 
     fn non_df_port_count(&self, dir: Direction) -> usize {
@@ -210,10 +208,6 @@ impl OpTrait for ExitBlock {
         Some(EdgeKind::ControlFlow)
     }
 
-    fn dataflow_signature(&self) -> Option<FunctionType> {
-        Some(FunctionType::new(type_row![], type_row![]))
-    }
-
     fn non_df_port_count(&self, dir: Direction) -> usize {
         match dir {
             Direction::Incoming => 1,