assistant2: Navigate context strip with keyboard (#23128)

Context pills are now focusable and intractable via the keyboard.

- <kbd>←</kbd> and <kbd>→</kbd> move the focus to the previous or next
item (wrapping if necessary)
- <kbd>↓</kbd> and <kbd>↑</kbd> move the focus vertically
- If the cursor is in the first/last row of the assistant/inline editor,
they will move the focus to the strip
- Inside the strip, they will move the focus to the pill horizontally
overlapping the most
- If already in the first/last row of the strip, they will move to the
first/last pill (like in editors)
- If the first/last pill is focused, they will move the focus back to
the editor
- <kbd>⌫</kbd>  removes the focused pill (unless it's the suggested one)
- <kbd>⏎</kbd> accepts the suggested pill if focused
  


https://github.com/user-attachments/assets/040bc71c-a3ae-4961-9886-2d5c3d290a73



Release Notes:

- N/A

assets/keymaps/default-linux.json

@@ -584,11 +584,24 @@
       "enter": "assistant2::Chat"
     }
   },
+  {
+    "context": "ContextStrip",
+    "use_key_equivalents": true,
+    "bindings": {
+      "up": "assistant2::FocusUp",
+      "right": "assistant2::FocusRight",
+      "left": "assistant2::FocusLeft",
+      "down": "assistant2::FocusDown",
+      "backspace": "assistant2::RemoveFocusedContext",
+      "enter": "assistant2::AcceptSuggestedContext"
+    }
+  },
   {
     "context": "PromptEditor",
     "bindings": {
       "ctrl-[": "assistant::CyclePreviousInlineAssist",
-      "ctrl-]": "assistant::CycleNextInlineAssist"
+      "ctrl-]": "assistant::CycleNextInlineAssist",
+      "ctrl-alt-e": "assistant2::RemoveAllContext"
     }
   },
   {

assets/keymaps/default-macos.json

@@ -238,6 +238,18 @@
       "enter": "assistant2::Chat"
     }
   },
+  {
+    "context": "ContextStrip",
+    "use_key_equivalents": true,
+    "bindings": {
+      "up": "assistant2::FocusUp",
+      "right": "assistant2::FocusRight",
+      "left": "assistant2::FocusLeft",
+      "down": "assistant2::FocusDown",
+      "backspace": "assistant2::RemoveFocusedContext",
+      "enter": "assistant2::AcceptSuggestedContext"
+    }
+  },
   {
     "context": "PromptLibrary",
     "use_key_equivalents": true,

crates/assistant2/src/active_thread.rs

@@ -327,13 +327,11 @@ impl ActiveThread {
                     )
                     .when_some(context, |parent, context| {
                         if !context.is_empty() {
-                            parent.child(
-                                h_flex().flex_wrap().gap_1().px_1p5().pb_1p5().children(
+                            parent.child(h_flex().flex_wrap().gap_1().px_1p5().pb_1p5().children(
                                 context.into_iter().map(|context| {
-                                        ContextPill::new_added(context, false, None)
+                                    ContextPill::new_added(context, false, false, None)
                                 }),
-                                ),
-                            )
+                            ))
                         } else {
                             parent
                         }

crates/assistant2/src/assistant.rs

@@ -45,7 +45,13 @@ actions!(
         OpenHistory,
         Chat,
         CycleNextInlineAssist,
-        CyclePreviousInlineAssist
+        CyclePreviousInlineAssist,
+        FocusUp,
+        FocusDown,
+        FocusLeft,
+        FocusRight,
+        RemoveFocusedContext,
+        AcceptSuggestedContext
     ]
 );
 

crates/assistant2/src/context_strip.rs

@@ -4,7 +4,8 @@ use collections::HashSet;
 use editor::Editor;
 use file_icons::FileIcons;
 use gpui::{
-    DismissEvent, EventEmitter, FocusHandle, Model, Subscription, View, WeakModel, WeakView,
+    AppContext, Bounds, DismissEvent, EventEmitter, FocusHandle, FocusableView, Model,
+    Subscription, View, WeakModel, WeakView,
 };
 use itertools::Itertools;
 use language::Buffer;
@@ -17,7 +18,10 @@ use crate::context_store::ContextStore;
 use crate::thread::Thread;
 use crate::thread_store::ThreadStore;
 use crate::ui::ContextPill;
-use crate::{AssistantPanel, RemoveAllContext, ToggleContextPicker};
+use crate::{
+    AcceptSuggestedContext, AssistantPanel, FocusDown, FocusLeft, FocusRight, FocusUp,
+    RemoveAllContext, RemoveFocusedContext, ToggleContextPicker,
+};
 
 pub struct ContextStrip {
     context_store: Model<ContextStore>,
@@ -26,7 +30,9 @@ pub struct ContextStrip {
     focus_handle: FocusHandle,
     suggest_context_kind: SuggestContextKind,
     workspace: WeakView<Workspace>,
-    _context_picker_subscription: Subscription,
+    _subscriptions: Vec<Subscription>,
+    focused_index: Option<usize>,
+    children_bounds: Option<Vec<Bounds<Pixels>>>,
 }
 
 impl ContextStrip {
@@ -34,7 +40,6 @@ impl ContextStrip {
         context_store: Model<ContextStore>,
         workspace: WeakView<Workspace>,
         thread_store: Option<WeakModel<ThreadStore>>,
-        focus_handle: FocusHandle,
         context_picker_menu_handle: PopoverMenuHandle<ContextPicker>,
         suggest_context_kind: SuggestContextKind,
         cx: &mut ViewContext<Self>,
@@ -49,8 +54,13 @@ impl ContextStrip {
             )
         });
 
-        let context_picker_subscription =
-            cx.subscribe(&context_picker, Self::handle_context_picker_event);
+        let focus_handle = cx.focus_handle();
+
+        let subscriptions = vec![
+            cx.subscribe(&context_picker, Self::handle_context_picker_event),
+            cx.on_focus(&focus_handle, Self::handle_focus),
+            cx.on_blur(&focus_handle, Self::handle_blur),
+        ];
 
         Self {
             context_store: context_store.clone(),
@@ -59,7 +69,9 @@ impl ContextStrip {
             focus_handle,
             suggest_context_kind,
             workspace,
-            _context_picker_subscription: context_picker_subscription,
+            _subscriptions: subscriptions,
+            focused_index: None,
+            children_bounds: None,
         }
     }
 
@@ -137,6 +149,199 @@ impl ContextStrip {
     ) {
         cx.emit(ContextStripEvent::PickerDismissed);
     }
+
+    fn handle_focus(&mut self, cx: &mut ViewContext<Self>) {
+        self.focused_index = self.last_pill_index();
+        cx.notify();
+    }
+
+    fn handle_blur(&mut self, cx: &mut ViewContext<Self>) {
+        self.focused_index = None;
+        cx.notify();
+    }
+
+    fn focus_left(&mut self, _: &FocusLeft, cx: &mut ViewContext<Self>) {
+        self.focused_index = match self.focused_index {
+            Some(index) if index > 0 => Some(index - 1),
+            _ => self.last_pill_index(),
+        };
+
+        cx.notify();
+    }
+
+    fn focus_right(&mut self, _: &FocusRight, cx: &mut ViewContext<Self>) {
+        let Some(last_index) = self.last_pill_index() else {
+            return;
+        };
+
+        self.focused_index = match self.focused_index {
+            Some(index) if index < last_index => Some(index + 1),
+            _ => Some(0),
+        };
+
+        cx.notify();
+    }
+
+    fn focus_up(&mut self, _: &FocusUp, cx: &mut ViewContext<Self>) {
+        let Some(focused_index) = self.focused_index else {
+            return;
+        };
+
+        if focused_index == 0 {
+            return cx.emit(ContextStripEvent::BlurredUp);
+        }
+
+        let Some((focused, pills)) = self.focused_bounds(focused_index) else {
+            return;
+        };
+
+        let iter = pills[..focused_index].iter().enumerate().rev();
+        self.focused_index = Self::find_best_horizontal_match(focused, iter).or(Some(0));
+        cx.notify();
+    }
+
+    fn focus_down(&mut self, _: &FocusDown, cx: &mut ViewContext<Self>) {
+        let Some(focused_index) = self.focused_index else {
+            return;
+        };
+
+        let last_index = self.last_pill_index();
+
+        if self.focused_index == last_index {
+            return cx.emit(ContextStripEvent::BlurredDown);
+        }
+
+        let Some((focused, pills)) = self.focused_bounds(focused_index) else {
+            return;
+        };
+
+        let iter = pills.iter().enumerate().skip(focused_index + 1);
+        self.focused_index = Self::find_best_horizontal_match(focused, iter).or(last_index);
+        cx.notify();
+    }
+
+    fn focused_bounds(&self, focused: usize) -> Option<(&Bounds<Pixels>, &[Bounds<Pixels>])> {
+        let pill_bounds = self.pill_bounds()?;
+        let focused = pill_bounds.get(focused)?;
+
+        Some((focused, pill_bounds))
+    }
+
+    fn pill_bounds(&self) -> Option<&[Bounds<Pixels>]> {
+        let bounds = self.children_bounds.as_ref()?;
+        let eraser = if bounds.len() < 3 { 0 } else { 1 };
+        let pills = &bounds[1..bounds.len() - eraser];
+
+        if pills.is_empty() {
+            None
+        } else {
+            Some(pills)
+        }
+    }
+
+    fn last_pill_index(&self) -> Option<usize> {
+        Some(self.pill_bounds()?.len() - 1)
+    }
+
+    fn find_best_horizontal_match<'a>(
+        focused: &'a Bounds<Pixels>,
+        iter: impl Iterator<Item = (usize, &'a Bounds<Pixels>)>,
+    ) -> Option<usize> {
+        let mut best = None;
+
+        let focused_left = focused.left();
+        let focused_right = focused.right();
+
+        for (index, probe) in iter {
+            if probe.origin.y == focused.origin.y {
+                continue;
+            }
+
+            let overlap = probe.right().min(focused_right) - probe.left().max(focused_left);
+
+            best = match best {
+                Some((_, prev_overlap, y)) if probe.origin.y != y || prev_overlap > overlap => {
+                    break;
+                }
+                Some(_) | None => Some((index, overlap, probe.origin.y)),
+            };
+        }
+
+        best.map(|(index, _, _)| index)
+    }
+
+    fn remove_focused_context(&mut self, _: &RemoveFocusedContext, cx: &mut ViewContext<Self>) {
+        if let Some(index) = self.focused_index {
+            let mut is_empty = false;
+
+            self.context_store.update(cx, |this, _cx| {
+                if let Some(item) = this.context().get(index) {
+                    this.remove_context(item.id());
+                }
+
+                is_empty = this.context().is_empty();
+            });
+
+            if is_empty {
+                cx.emit(ContextStripEvent::BlurredEmpty);
+            } else {
+                self.focused_index = Some(index.saturating_sub(1));
+                cx.notify();
+            }
+        }
+    }
+
+    fn is_suggested_focused<T>(&self, context: &Vec<T>) -> bool {
+        // We only suggest one item after the actual context
+        self.focused_index == Some(context.len())
+    }
+
+    fn accept_suggested_context(&mut self, _: &AcceptSuggestedContext, cx: &mut ViewContext<Self>) {
+        if let Some(suggested) = self.suggested_context(cx) {
+            let context_store = self.context_store.read(cx);
+
+            if self.is_suggested_focused(context_store.context()) {
+                self.add_suggested_context(&suggested, cx);
+            }
+        }
+    }
+
+    fn add_suggested_context(&mut self, suggested: &SuggestedContext, cx: &mut ViewContext<Self>) {
+        let task = self.context_store.update(cx, |context_store, cx| {
+            context_store.accept_suggested_context(&suggested, cx)
+        });
+
+        let workspace = self.workspace.clone();
+
+        cx.spawn(|this, mut cx| async move {
+            match task.await {
+                Ok(()) => {
+                    if let Some(this) = this.upgrade() {
+                        this.update(&mut cx, |_, cx| cx.notify())?;
+                    }
+                }
+                Err(err) => {
+                    let Some(workspace) = workspace.upgrade() else {
+                        return anyhow::Ok(());
+                    };
+
+                    workspace.update(&mut cx, |workspace, cx| {
+                        workspace.show_error(&err, cx);
+                    })?;
+                }
+            }
+            anyhow::Ok(())
+        })
+        .detach_and_log_err(cx);
+
+        cx.notify();
+    }
+}
+
+impl FocusableView for ContextStrip {
+    fn focus_handle(&self, _cx: &AppContext) -> FocusHandle {
+        self.focus_handle.clone()
+    }
 }
 
 impl Render for ContextStrip {
@@ -164,6 +369,23 @@ impl Render for ContextStrip {
         h_flex()
             .flex_wrap()
             .gap_1()
+            .track_focus(&focus_handle)
+            .key_context("ContextStrip")
+            .on_action(cx.listener(Self::focus_up))
+            .on_action(cx.listener(Self::focus_right))
+            .on_action(cx.listener(Self::focus_down))
+            .on_action(cx.listener(Self::focus_left))
+            .on_action(cx.listener(Self::remove_focused_context))
+            .on_action(cx.listener(Self::accept_suggested_context))
+            .on_children_prepainted({
+                let view = cx.view().downgrade();
+                move |children_bounds, cx| {
+                    view.update(cx, |this, _| {
+                        this.children_bounds = Some(children_bounds);
+                    })
+                    .ok();
+                }
+            })
             .child(
                 PopoverMenu::new("context-picker")
                     .menu(move |cx| {
@@ -217,10 +439,11 @@ impl Render for ContextStrip {
                     )
                 }
             })
-            .children(context.iter().map(|context| {
+            .children(context.iter().enumerate().map(|(i, context)| {
                 ContextPill::new_added(
                     context.clone(),
                     dupe_names.contains(&context.name),
+                    self.focused_index == Some(i),
                     Some({
                         let id = context.id;
                         let context_store = self.context_store.clone();
@@ -232,43 +455,23 @@ impl Render for ContextStrip {
                         }))
                     }),
                 )
+                .on_click(Rc::new(cx.listener(move |this, _, cx| {
+                    this.focused_index = Some(i);
+                    cx.notify();
+                })))
             }))
             .when_some(suggested_context, |el, suggested| {
-                el.child(ContextPill::new_suggested(
+                el.child(
+                    ContextPill::new_suggested(
                         suggested.name().clone(),
                         suggested.icon_path(),
                         suggested.kind(),
-                    {
-                        let context_store = self.context_store.clone();
-                        Rc::new(cx.listener(move |this, _event, cx| {
-                            let task = context_store.update(cx, |context_store, cx| {
-                                context_store.accept_suggested_context(&suggested, cx)
-                            });
-
-                            let workspace = this.workspace.clone();
-                            cx.spawn(|this, mut cx| async move {
-                                match task.await {
-                                    Ok(()) => {
-                                        if let Some(this) = this.upgrade() {
-                                            this.update(&mut cx, |_, cx| cx.notify())?;
-                                        }
-                                    }
-                                    Err(err) => {
-                                        let Some(workspace) = workspace.upgrade() else {
-                                            return anyhow::Ok(());
-                                        };
-
-                                        workspace.update(&mut cx, |workspace, cx| {
-                                            workspace.show_error(&err, cx);
-                                        })?;
-                                    }
-                                }
-                                anyhow::Ok(())
-                            })
-                            .detach_and_log_err(cx);
-                        }))
-                    },
-                ))
+                        self.is_suggested_focused(&context),
+                    )
+                    .on_click(Rc::new(cx.listener(move |this, _event, cx| {
+                        this.add_suggested_context(&suggested, cx);
+                    }))),
+                )
             })
             .when(!context.is_empty(), {
                 move |parent| {
@@ -300,6 +503,9 @@ impl Render for ContextStrip {
 
 pub enum ContextStripEvent {
     PickerDismissed,
+    BlurredEmpty,
+    BlurredDown,
+    BlurredUp,
 }
 
 impl EventEmitter<ContextStripEvent> for ContextStrip {}

crates/assistant2/src/inline_prompt_editor.rs

@@ -415,6 +415,8 @@ impl<T: 'static> PromptEditor<T> {
                     editor.move_to_end(&Default::default(), cx)
                 });
             }
+        } else {
+            cx.focus_view(&self.context_strip);
         }
     }
 
@@ -738,12 +740,19 @@ impl<T: 'static> PromptEditor<T> {
     fn handle_context_strip_event(
         &mut self,
         _context_strip: View<ContextStrip>,
-        ContextStripEvent::PickerDismissed: &ContextStripEvent,
+        event: &ContextStripEvent,
         cx: &mut ViewContext<Self>,
     ) {
+        match event {
+            ContextStripEvent::PickerDismissed
+            | ContextStripEvent::BlurredEmpty
+            | ContextStripEvent::BlurredUp => {
                 let editor_focus_handle = self.editor.focus_handle(cx);
                 cx.focus(&editor_focus_handle);
             }
+            ContextStripEvent::BlurredDown => {}
+        }
+    }
 }
 
 pub enum PromptEditorMode {
@@ -826,7 +835,6 @@ impl PromptEditor<BufferCodegen> {
                 context_store.clone(),
                 workspace.clone(),
                 thread_store.clone(),
-                prompt_editor.focus_handle(cx),
                 context_picker_menu_handle.clone(),
                 SuggestContextKind::Thread,
                 cx,
@@ -978,7 +986,6 @@ impl PromptEditor<TerminalCodegen> {
                 context_store.clone(),
                 workspace.clone(),
                 thread_store.clone(),
-                prompt_editor.focus_handle(cx),
                 context_picker_menu_handle.clone(),
                 SuggestContextKind::Thread,
                 cx,

crates/assistant2/src/message_editor.rs

@@ -1,5 +1,6 @@
 use std::sync::Arc;
 
+use editor::actions::MoveUp;
 use editor::{Editor, EditorElement, EditorEvent, EditorStyle};
 use fs::Fs;
 use gpui::{
@@ -75,7 +76,6 @@ impl MessageEditor {
                 context_store.clone(),
                 workspace.clone(),
                 Some(thread_store.clone()),
-                editor.focus_handle(cx),
                 context_picker_menu_handle.clone(),
                 SuggestContextKind::File,
                 cx,
@@ -221,12 +221,27 @@ impl MessageEditor {
     fn handle_context_strip_event(
         &mut self,
         _context_strip: View<ContextStrip>,
-        ContextStripEvent::PickerDismissed: &ContextStripEvent,
+        event: &ContextStripEvent,
         cx: &mut ViewContext<Self>,
     ) {
+        match event {
+            ContextStripEvent::PickerDismissed
+            | ContextStripEvent::BlurredEmpty
+            | ContextStripEvent::BlurredDown => {
                 let editor_focus_handle = self.editor.focus_handle(cx);
                 cx.focus(&editor_focus_handle);
             }
+            ContextStripEvent::BlurredUp => {}
+        }
+    }
+
+    fn move_up(&mut self, _: &MoveUp, cx: &mut ViewContext<Self>) {
+        if self.context_picker_menu_handle.is_deployed() {
+            cx.propagate();
+        } else {
+            cx.focus_view(&self.context_strip);
+        }
+    }
 }
 
 impl FocusableView for MessageEditor {
@@ -249,6 +264,7 @@ impl Render for MessageEditor {
             .on_action(cx.listener(Self::toggle_model_selector))
             .on_action(cx.listener(Self::toggle_context_picker))
             .on_action(cx.listener(Self::remove_all_context))
+            .on_action(cx.listener(Self::move_up))
             .size_full()
             .gap_2()
             .p_2()

crates/assistant2/src/ui/context_pill.rs

@@ -10,13 +10,16 @@ pub enum ContextPill {
     Added {
         context: ContextSnapshot,
         dupe_name: bool,
+        focused: bool,
+        on_click: Option<Rc<dyn Fn(&ClickEvent, &mut WindowContext)>>,
         on_remove: Option<Rc<dyn Fn(&ClickEvent, &mut WindowContext)>>,
     },
     Suggested {
         name: SharedString,
         icon_path: Option<SharedString>,
         kind: ContextKind,
-        on_add: Rc<dyn Fn(&ClickEvent, &mut WindowContext)>,
+        focused: bool,
+        on_click: Option<Rc<dyn Fn(&ClickEvent, &mut WindowContext)>>,
     },
 }
 
@@ -24,12 +27,15 @@ impl ContextPill {
     pub fn new_added(
         context: ContextSnapshot,
         dupe_name: bool,
+        focused: bool,
         on_remove: Option<Rc<dyn Fn(&ClickEvent, &mut WindowContext)>>,
     ) -> Self {
         Self::Added {
             context,
             dupe_name,
             on_remove,
+            focused,
+            on_click: None,
         }
     }
 
@@ -37,16 +43,29 @@ impl ContextPill {
         name: SharedString,
         icon_path: Option<SharedString>,
         kind: ContextKind,
-        on_add: Rc<dyn Fn(&ClickEvent, &mut WindowContext)>,
+        focused: bool,
     ) -> Self {
         Self::Suggested {
             name,
             icon_path,
             kind,
-            on_add,
+            focused,
+            on_click: None,
         }
     }
 
+    pub fn on_click(mut self, listener: Rc<dyn Fn(&ClickEvent, &mut WindowContext)>) -> Self {
+        match &mut self {
+            ContextPill::Added { on_click, .. } => {
+                *on_click = Some(listener);
+            }
+            ContextPill::Suggested { on_click, .. } => {
+                *on_click = Some(listener);
+            }
+        }
+        self
+    }
+
     pub fn id(&self) -> ElementId {
         match self {
             Self::Added { context, .. } => {
@@ -93,9 +112,15 @@ impl RenderOnce for ContextPill {
                 context,
                 dupe_name,
                 on_remove,
+                focused,
+                on_click,
             } => base_pill
                 .bg(color.element_background)
-                .border_color(color.border.opacity(0.5))
+                .border_color(if *focused {
+                    color.border_focused
+                } else {
+                    color.border.opacity(0.5)
+                })
                 .pr(if on_remove.is_some() { px(2.) } else { px(4.) })
                 .child(
                     h_flex()
@@ -128,16 +153,25 @@ impl RenderOnce for ContextPill {
                                 move |event, cx| on_remove(event, cx)
                             }),
                     )
+                })
+                .when_some(on_click.as_ref(), |element, on_click| {
+                    let on_click = on_click.clone();
+                    element.on_click(move |event, cx| on_click(event, cx))
                 }),
             ContextPill::Suggested {
                 name,
                 icon_path: _,
                 kind,
-                on_add,
+                focused,
+                on_click,
             } => base_pill
                 .cursor_pointer()
                 .pr_1()
-                .border_color(color.border_variant.opacity(0.5))
+                .border_color(if *focused {
+                    color.border_focused
+                } else {
+                    color.border_variant.opacity(0.5)
+                })
                 .hover(|style| style.bg(color.element_hover.opacity(0.5)))
                 .child(
                     Label::new(name.clone())
@@ -162,9 +196,9 @@ impl RenderOnce for ContextPill {
                         .into_any_element(),
                 )
                 .tooltip(|cx| Tooltip::with_meta("Suggested Context", None, "Click to add it", cx))
-                .on_click({
-                    let on_add = on_add.clone();
-                    move |event, cx| on_add(event, cx)
+                .when_some(on_click.as_ref(), |element, on_click| {
+                    let on_click = on_click.clone();
+                    element.on_click(move |event, cx| on_click(event, cx))
                 }),
         }
     }

crates/gpui/src/elements/div.rs

@@ -1104,6 +1104,7 @@ pub fn div() -> Div {
     Div {
         interactivity,
         children: SmallVec::default(),
+        prepaint_listener: None,
     }
 }
 
@@ -1111,6 +1112,19 @@ pub fn div() -> Div {
 pub struct Div {
     interactivity: Interactivity,
     children: SmallVec<[AnyElement; 2]>,
+    prepaint_listener: Option<Box<dyn Fn(Vec<Bounds<Pixels>>, &mut WindowContext) + 'static>>,
+}
+
+impl Div {
+    /// Add a listener to be called when the children of this `Div` are prepainted.
+    /// This allows you to store the [`Bounds`] of the children for later use.
+    pub fn on_children_prepainted(
+        mut self,
+        listener: impl Fn(Vec<Bounds<Pixels>>, &mut WindowContext) + 'static,
+    ) -> Self {
+        self.prepaint_listener = Some(Box::new(listener));
+        self
+    }
 }
 
 /// A frame state for a `Div` element, which contains layout IDs for its children.
@@ -1177,6 +1191,13 @@ impl Element for Div {
         request_layout: &mut Self::RequestLayoutState,
         cx: &mut WindowContext,
     ) -> Option<Hitbox> {
+        let has_prepaint_listener = self.prepaint_listener.is_some();
+        let mut children_bounds = Vec::with_capacity(if has_prepaint_listener {
+            request_layout.child_layout_ids.len()
+        } else {
+            0
+        });
+
         let mut child_min = point(Pixels::MAX, Pixels::MAX);
         let mut child_max = Point::default();
         if let Some(handle) = self.interactivity.scroll_anchor.as_ref() {
@@ -1189,6 +1210,7 @@ impl Element for Div {
             state.child_bounds = Vec::with_capacity(request_layout.child_layout_ids.len());
             state.bounds = bounds;
             let requested = state.requested_scroll_top.take();
+            // TODO az
 
             for (ix, child_layout_id) in request_layout.child_layout_ids.iter().enumerate() {
                 let child_bounds = cx.layout_bounds(*child_layout_id);
@@ -1209,6 +1231,10 @@ impl Element for Div {
                 let child_bounds = cx.layout_bounds(*child_layout_id);
                 child_min = child_min.min(&child_bounds.origin);
                 child_max = child_max.max(&child_bounds.bottom_right());
+
+                if has_prepaint_listener {
+                    children_bounds.push(child_bounds);
+                }
             }
             (child_max - child_min).into()
         };
@@ -1224,6 +1250,11 @@ impl Element for Div {
                         child.prepaint(cx);
                     }
                 });
+
+                if let Some(listener) = self.prepaint_listener.as_ref() {
+                    listener(children_bounds, cx);
+                }
+
                 hitbox
             },
         )
@@ -2330,6 +2361,18 @@ where
     }
 }
 
+impl Focusable<Div> {
+    /// Add a listener to be called when the children of this `Div` are prepainted.
+    /// This allows you to store the [`Bounds`] of the children for later use.
+    pub fn on_children_prepainted(
+        mut self,
+        listener: impl Fn(Vec<Bounds<Pixels>>, &mut WindowContext) + 'static,
+    ) -> Self {
+        self.element = self.element.on_children_prepainted(listener);
+        self
+    }
+}
+
 impl<E> Element for Focusable<E>
 where
     E: Element,