diff --git a/app/views/units/_form_close.html.erb b/app/views/units/_form_close.html.erb
index bea914b..5226d99 100644
--- a/app/views/units/_form_close.html.erb
+++ b/app/views/units/_form_close.html.erb
@@ -1,4 +1,3 @@
-<%= turbo_stream.remove @unit %>
-<%= turbo_stream.enable_all 'td.actions .button' %>
-<%= turbo_stream.enable :add_unit %>
-<%= turbo_stream.focus :add_unit %>
+<%= @unit.persisted? ? turbo_stream.replace(@unit) : turbo_stream.remove(@unit) %>
+<%= turbo_stream.enable id %>
+<%= turbo_stream.focus id %>
diff --git a/app/views/units/_unit.html.erb b/app/views/units/_unit.html.erb
index cd2332b..b787939 100644
--- a/app/views/units/_unit.html.erb
+++ b/app/views/units/_unit.html.erb
@@ -1,6 +1,7 @@
 <tr id="<%= dom_id(unit) %>">
   <td class="<%= class_names('link', {subunit: unit.base}) %>">
-    <%= link_to unit.symbol, edit_unit_path(unit), data: {turbo_stream: true} %>
+    <%= link_to unit.symbol, edit_unit_path(unit), id: dom_id(unit, :edit),
+      data: {turbo_stream: true} %>
   </td>
   <td><%= unit.name %></td>
   <td class="number"><%= scientifize(unit.multiplier) if unit.multiplier %></td>
@@ -9,7 +10,8 @@
     <td class="actions">
       <% if unit.base.nil? %>
         <%= image_link_to t(".add_subunit"), "plus-outline", new_unit_path(unit),
-          onclick: 'this.blur();', data: {turbo_stream: true} %>
+          id: dom_id(unit, :add), onclick: 'this.blur();',
+          data: {turbo_stream: true} %>
       <% end %>
 
       <%= image_button_to t(".delete_unit"), "delete-outline", unit_path(unit),
diff --git a/app/views/units/edit.turbo_stream.erb b/app/views/units/edit.turbo_stream.erb
index ebb2aed..e5df3f7 100644
--- a/app/views/units/edit.turbo_stream.erb
+++ b/app/views/units/edit.turbo_stream.erb
@@ -1,10 +1,7 @@
 <%# TODO: make sure turbo_stream layout is used in new/edit %>
-<%= turbo_stream.disable :add_unit %>
-<%= turbo_stream.disable_all 'td.actions .button' %>
-
-<%= turbo_stream.replace @unit, partial: 'form' %>
+<%= turbo_stream.replace @unit, partial: 'form', locals: {id: dom_id(@unit, :edit)} %>
 
 <%= turbo_stream.update :unit_form_frame do %>
   <%= form_with model: @unit, html: {id: :unit_form} do %>
   <% end %>
-<% end %>
+<% end unless @unit.errors.present? %>
diff --git a/app/views/units/new.turbo_stream.erb b/app/views/units/new.turbo_stream.erb
index dd4d73f..ff2d753 100644
--- a/app/views/units/new.turbo_stream.erb
+++ b/app/views/units/new.turbo_stream.erb
@@ -1,14 +1,16 @@
-<%= turbo_stream.disable :add_unit %>
-<%# TODO: disable only form-opening actions %>
-<%= turbo_stream.disable_all 'td.actions .button' %>
+<% options = {partial: 'form', locals: {id: dom_id(@unit.base || @unit, :add)}} %>
 
 <% if @unit.errors.present? %>
-  <%= turbo_stream.replace @unit, partial: 'form' %>
+  <%= turbo_stream.replace @unit, **options -%>
 <% else %>
+  <%= turbo_stream.disable options[:locals][:id] -%>
+  <%= turbo_stream.click_all 'tbody a[name=cancel]' -%>
+  <%#= turbo_stream.blur_all %>
+
   <% if @unit.base.nil? %>
-    <%= turbo_stream.prepend :units, partial: 'form' %>
+    <%= turbo_stream.prepend :units, **options -%>
   <% else %>
-    <%= turbo_stream.after @unit.base, partial: 'form' %>
+    <%= turbo_stream.after @unit.base, **options -%>
   <% end %>
 
   <%= turbo_stream.update :unit_form_frame do %>