diff --git a/crates/storybook2/src/ui.rs b/crates/storybook2/src/ui.rs
index 8b09f6f0ff..b74ac0641a 100644
--- a/crates/storybook2/src/ui.rs
+++ b/crates/storybook2/src/ui.rs
@@ -1,10 +1,12 @@
 mod children;
 mod components;
+mod element_ext;
 mod elements;
 pub mod prelude;
 mod tokens;
 
 pub use children::*;
 pub use components::*;
+pub use element_ext::*;
 pub use elements::*;
 pub use tokens::*;
diff --git a/crates/storybook2/src/ui/components/list.rs b/crates/storybook2/src/ui/components/list.rs
index a491045757..02f480c37f 100644
--- a/crates/storybook2/src/ui/components/list.rs
+++ b/crates/storybook2/src/ui/components/list.rs
@@ -111,16 +111,16 @@ impl<S: 'static + Send + Sync + Clone> ListHeader<S> {
             .flex_1()
             .w_full()
             .fill(background_color)
-            // .when(self.state == InteractionState::Focused, |this| {
-            //     this.border()
-            //         .border_color(theme.lowest.accent.default.border)
-            // })
+            .when(self.state == InteractionState::Focused, |this| {
+                this.border()
+                    .border_color(theme.lowest.accent.default.border)
+            })
             .relative()
             .py_1()
             .child(
                 div()
                     .h_6()
-                    // .when(self.variant == ListItemVariant::Inset, |this| this.px_2())
+                    .when(self.variant == ListItemVariant::Inset, |this| this.px_2())
                     .flex()
                     .flex_1()
                     .w_full()
@@ -178,7 +178,7 @@ impl<S: 'static + Send + Sync + Clone> ListSubHeader<S> {
         h_stack().flex_1().w_full().relative().py_1().child(
             div()
                 .h_6()
-                // .when(self.variant == ListItemVariant::Inset, |this| this.px_2())
+                .when(self.variant == ListItemVariant::Inset, |this| this.px_2())
                 .flex()
                 .flex_1()
                 .w_full()
@@ -412,15 +412,15 @@ impl<S: 'static + Send + Sync + Clone> ListEntry<S> {
 
         div()
             .fill(background_color)
-            // .when(self.state == InteractionState::Focused, |this| {
-            //     this.border()
-            //         .border_color(theme.lowest.accent.default.border)
-            // })
+            .when(self.state == InteractionState::Focused, |this| {
+                this.border()
+                    .border_color(theme.lowest.accent.default.border)
+            })
             .relative()
             .py_1()
             .child(
                 sized_item
-                    // .when(self.variant == ListItemVariant::Inset, |this| this.px_2())
+                    .when(self.variant == ListItemVariant::Inset, |this| this.px_2())
                     // .ml(rems(0.75 * self.indent_level as f32))
                     .children((0..self.indent_level).map(|_| {
                         div()
diff --git a/crates/storybook2/src/ui/element_ext.rs b/crates/storybook2/src/ui/element_ext.rs
new file mode 100644
index 0000000000..41c91c0450
--- /dev/null
+++ b/crates/storybook2/src/ui/element_ext.rs
@@ -0,0 +1,18 @@
+use gpui3::Element;
+
+pub trait ElementExt<S: 'static + Send + Sync>: Element<State = S> {
+    /// Applies a given function `then` to the current element if `condition` is true.
+    /// This function is used to conditionally modify the element based on a given condition.
+    /// If `condition` is false, it just returns the current element as it is.
+    fn when(mut self, condition: bool, then: impl FnOnce(Self) -> Self) -> Self
+    where
+        Self: Sized,
+    {
+        if condition {
+            self = then(self);
+        }
+        self
+    }
+}
+
+impl<S: 'static + Send + Sync, E: Element<State = S>> ElementExt<S> for E {}
diff --git a/crates/storybook2/src/ui/prelude.rs b/crates/storybook2/src/ui/prelude.rs
index a2ab86d372..252ef42853 100644
--- a/crates/storybook2/src/ui/prelude.rs
+++ b/crates/storybook2/src/ui/prelude.rs
@@ -3,7 +3,7 @@ pub use gpui3::{
     WindowContext,
 };
 
-pub use crate::ui::{HackyChildren, HackyChildrenPayload};
+pub use crate::ui::{HackyChildren, HackyChildrenPayload, ElementExt};
 
 use gpui3::{hsla, rgb, Hsla};
 use strum::EnumIter;