Query for window instead of capturing (#55059)

This allows us to move entities between windows without breaking all the
callbacks.

Self-Review Checklist:

- [x] I've reviewed my own diff for quality, security, and reliability
- [x] Unsafe blocks (if any) have justifying comments
- [x] The content is consistent with the [UI/UX
checklist](https://github.com/zed-industries/zed/blob/main/CONTRIBUTING.md#uiux-checklist)
- [x] Tests cover the new/changed behavior
- [x] Performance impact has been considered and is acceptable

Release Notes:

- N/A

crates/gpui/Cargo.toml

@@ -180,6 +180,10 @@ naga.workspace = true
 name = "hello_world"
 path = "examples/hello_world.rs"
 
+[[example]]
+name = "move_entity_between_windows"
+path = "examples/move_entity_between_windows.rs"
+
 [[example]]
 name = "image"
 path = "examples/image/image.rs"

crates/gpui/examples/move_entity_between_windows.rs

@@ -0,0 +1,154 @@
+//! An entity registers callbacks via the `_in` API family and then gets
+//! re-hosted in a new window via a click. The point of the example is to
+//! demonstrate that callbacks dispatched after the move correctly target the
+//! entity's *current* window rather than the window it was in at
+//! registration time.
+//!
+//! To run:  cargo run -p gpui --example move_entity_between_windows
+
+#![cfg_attr(target_family = "wasm", no_main)]
+
+use std::time::Duration;
+
+use gpui::{
+    App, AppContext as _, Bounds, Context, EventEmitter, MouseButton, Render, SharedString,
+    Subscription, Task, Window, WindowBounds, WindowOptions, div, prelude::*, px, rgb, size,
+};
+use gpui_platform::application;
+
+struct MoveToNewWindow;
+
+struct HelloWorld {
+    text: SharedString,
+    tick_count: u32,
+    move_count: u32,
+    _tasks: Vec<Task<()>>,
+    _subscriptions: Vec<Subscription>,
+}
+
+impl EventEmitter<MoveToNewWindow> for HelloWorld {}
+
+impl HelloWorld {
+    fn new(window: &mut Window, cx: &mut Context<Self>) -> Self {
+        let self_entity = cx.entity();
+
+        let task = cx.spawn_in(window, async move |this, cx| {
+            loop {
+                cx.background_executor().timer(Duration::from_secs(1)).await;
+                let result = this.update_in(cx, |this, window, _cx| {
+                    this.tick_count += 1;
+                    println!(
+                        "tick #{} fired in entity's current window {}",
+                        this.tick_count,
+                        window.window_handle().window_id().as_u64(),
+                    );
+                });
+                if let Err(err) = result {
+                    println!("tick task giving up: {err}");
+                    return;
+                }
+            }
+        });
+
+        let subscription = cx.subscribe_in::<_, MoveToNewWindow>(
+            &self_entity,
+            window,
+            move |this, _emitter, _event, window, cx| {
+                let entered_window_id = window.window_handle().window_id().as_u64();
+                println!(
+                    "MoveToNewWindow handler fired in entity's current window {entered_window_id}",
+                );
+
+                this.move_count += 1;
+                cx.notify();
+
+                let entity = cx.entity();
+                let old_window = window.window_handle();
+                cx.defer(move |cx| {
+                    let bounds = Bounds::centered(None, size(px(500.0), px(500.0)), cx);
+                    cx.open_window(
+                        WindowOptions {
+                            window_bounds: Some(WindowBounds::Windowed(bounds)),
+                            ..Default::default()
+                        },
+                        move |_, _| entity,
+                    )
+                    .expect("failed to open new window");
+                    old_window
+                        .update(cx, |_, window, _| window.remove_window())
+                        .ok();
+                });
+            },
+        );
+
+        Self {
+            text: "World".into(),
+            tick_count: 0,
+            move_count: 0,
+            _tasks: vec![task],
+            _subscriptions: vec![subscription],
+        }
+    }
+}
+
+impl Render for HelloWorld {
+    fn render(&mut self, window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
+        let window_id = window.window_handle().window_id().as_u64();
+
+        div()
+            .flex()
+            .flex_col()
+            .gap_3()
+            .bg(rgb(0x505050))
+            .size(px(500.0))
+            .justify_center()
+            .items_center()
+            .text_xl()
+            .text_color(rgb(0xffffff))
+            .child(format!("Hello, {}!", &self.text))
+            .child(format!("Rendering in window: {window_id}"))
+            .child(format!("Ticks observed by entity: {}", self.tick_count))
+            .child(format!("Moves observed by entity: {}", self.move_count))
+            .child(
+                div()
+                    .px_4()
+                    .py_2()
+                    .bg(rgb(0x4040ff))
+                    .rounded_md()
+                    .child("Move me to a new window")
+                    .on_mouse_down(
+                        MouseButton::Left,
+                        cx.listener(|_this, _, _window, cx| {
+                            cx.emit(MoveToNewWindow);
+                        }),
+                    ),
+            )
+    }
+}
+
+fn run_example() {
+    application().run(|cx: &mut App| {
+        let bounds = Bounds::centered(None, size(px(500.0), px(500.0)), cx);
+        cx.open_window(
+            WindowOptions {
+                window_bounds: Some(WindowBounds::Windowed(bounds)),
+                ..Default::default()
+            },
+            |window, cx| cx.new(|cx| HelloWorld::new(window, cx)),
+        )
+        .unwrap();
+        cx.activate(true);
+    });
+}
+
+#[cfg(not(target_family = "wasm"))]
+fn main() {
+    run_example();
+}
+
+#[cfg(target_family = "wasm")]
+#[wasm_bindgen::prelude::wasm_bindgen(start)]
+pub fn start() {
+    gpui_platform::web_init();
+    run_example();
+}

crates/gpui/src/app.rs

@@ -638,6 +638,7 @@ pub struct App {
     pub(crate) window_invalidators_by_entity:
         FxHashMap<EntityId, FxHashMap<WindowId, WindowInvalidator>>,
     pub(crate) tracked_entities: FxHashMap<WindowId, FxHashSet<EntityId>>,
+    pub(crate) current_window_by_entity: FxHashMap<EntityId, WindowId>,
     #[cfg(any(feature = "inspector", debug_assertions))]
     pub(crate) inspector_renderer: Option<crate::InspectorRenderer>,
     #[cfg(any(feature = "inspector", debug_assertions))]
@@ -715,6 +716,7 @@ impl App {
                 observers: SubscriberSet::new(),
                 tracked_entities: FxHashMap::default(),
                 window_invalidators_by_entity: FxHashMap::default(),
+                current_window_by_entity: FxHashMap::default(),
                 event_listeners: SubscriberSet::new(),
                 release_listeners: SubscriberSet::new(),
                 keystroke_observers: SubscriberSet::new(),
@@ -952,6 +954,8 @@ impl App {
                 .entry(*entity)
                 .or_default()
                 .insert(window_handle.id, invalidator.clone());
+            self.current_window_by_entity
+                .insert(*entity, window_handle.id);
         }
         tracked_entities.clear();
         tracked_entities.extend(entities.iter().copied());
@@ -1458,6 +1462,8 @@ impl App {
             for (entity_id, mut entity) in dropped {
                 self.observers.remove(&entity_id);
                 self.event_listeners.remove(&entity_id);
+                self.window_invalidators_by_entity.remove(&entity_id);
+                self.current_window_by_entity.remove(&entity_id);
                 for release_callback in self.release_listeners.remove(&entity_id) {
                     release_callback(entity.as_mut(), self);
                 }
@@ -1534,6 +1540,13 @@ impl App {
         tid: TypeId,
         window: Option<WindowId>,
     ) {
+        // Seed the entity's current window from its creation context so
+        // `with_window` resolves correctly before the entity has ever been
+        // rendered.
+        if let Some(id) = window {
+            self.current_window_by_entity.insert(entity.entity_id(), id);
+        }
+
         self.new_entity_observers.clone().retain(&tid, |observer| {
             if let Some(id) = window {
                 self.update_window_id(id, {
@@ -1548,7 +1561,28 @@ impl App {
         });
     }
 
-    fn update_window_id<T, F>(&mut self, id: WindowId, update: F) -> Result<T>
+    /// Run `f` against the entity's *current* window — the most recently
+    /// rendered window that referenced the entity, or its creation window if
+    /// it has yet to be rendered. Returns `None` if the entity has no
+    /// current window, or if that window has been closed, or if it is
+    /// already on the update stack.
+    pub fn with_window<R>(
+        &mut self,
+        entity_id: EntityId,
+        f: impl FnOnce(&mut Window, &mut App) -> R,
+    ) -> Option<R> {
+        let window_id = *self.current_window_by_entity.get(&entity_id)?;
+        self.update_window_id(window_id, |_, window, cx| f(window, cx))
+            .ok()
+    }
+
+    fn ensure_window(&mut self, entity_id: EntityId, window: WindowId) {
+        self.current_window_by_entity
+            .entry(entity_id)
+            .or_insert(window);
+    }
+
+    pub(crate) fn update_window_id<T, F>(&mut self, id: WindowId, update: F) -> Result<T>
     where
         F: FnOnce(AnyView, &mut Window, &mut App) -> T,
     {
@@ -1565,6 +1599,18 @@ impl App {
                 if window.removed {
                     cx.window_handles.remove(&id);
                     cx.windows.remove(id);
+                    if let Some(tracked) = cx.tracked_entities.remove(&id) {
+                        for entity_id in tracked {
+                            if let Some(windows) =
+                                cx.window_invalidators_by_entity.get_mut(&entity_id)
+                            {
+                                windows.remove(&id);
+                            }
+                            if cx.current_window_by_entity.get(&entity_id) == Some(&id) {
+                                cx.current_window_by_entity.remove(&entity_id);
+                            }
+                        }
+                    }
 
                     cx.window_closed_observers.clone().retain(&(), |callback| {
                         callback(cx, id);
@@ -2281,13 +2327,27 @@ impl App {
                 .or_default(),
         );
 
-        if window_invalidators.is_empty() {
+        // `window_invalidators_by_entity` is monotonic, so an entry alone
+        // doesn't mean the window is currently rendering the entity. Filter
+        // through `tracked_entities` to keep invalidation tight to windows
+        // that actually display this entity right now.
+        let live_invalidators: SmallVec<[WindowInvalidator; 2]> = window_invalidators
+            .iter()
+            .filter(|(window_id, _)| {
+                self.tracked_entities
+                    .get(window_id)
+                    .is_some_and(|set| set.contains(&entity_id))
+            })
+            .map(|(_, invalidator)| invalidator.clone())
+            .collect();
+
+        if live_invalidators.is_empty() {
             if self.pending_notifications.insert(entity_id) {
                 self.pending_effects
                     .push_back(Effect::Notify { emitter: entity_id });
             }
         } else {
-            for invalidator in window_invalidators.values() {
+            for invalidator in &live_invalidators {
                 invalidator.invalidate_view(entity_id, self);
             }
         }
@@ -2423,6 +2483,14 @@ impl AppContext for App {
         self.update_window_id(handle.id, update)
     }
 
+    fn with_window<R>(
+        &mut self,
+        entity_id: EntityId,
+        f: impl FnOnce(&mut Window, &mut App) -> R,
+    ) -> Option<R> {
+        App::with_window(self, entity_id, f)
+    }
+
     fn read_window<T, R>(
         &self,
         window: &WindowHandle<T>,

crates/gpui/src/app/async_context.rs

@@ -1,8 +1,8 @@
 use crate::{
     AnyView, AnyWindowHandle, App, AppCell, AppContext, BackgroundExecutor, BorrowAppContext,
-    Entity, EventEmitter, Focusable, ForegroundExecutor, Global, GpuiBorrow, PromptButton,
-    PromptLevel, Render, Reservation, Result, Subscription, Task, VisualContext, Window,
-    WindowHandle,
+    Entity, EntityId, EventEmitter, Focusable, ForegroundExecutor, Global, GpuiBorrow,
+    PromptButton, PromptLevel, Render, Reservation, Result, Subscription, Task, VisualContext,
+    Window, WindowHandle,
 };
 use anyhow::{Context as _, bail};
 use derive_more::{Deref, DerefMut};
@@ -94,6 +94,19 @@ impl AppContext for AsyncApp {
         lock.update_window(window, f)
     }
 
+    fn with_window<R>(
+        &mut self,
+        entity_id: EntityId,
+        f: impl FnOnce(&mut Window, &mut App) -> R,
+    ) -> Option<R> {
+        let app = self.app.upgrade()?;
+        let mut lock = app.try_borrow_mut().ok()?;
+        if lock.quitting {
+            return None;
+        }
+        lock.with_window(entity_id, f)
+    }
+
     fn read_window<T, R>(
         &self,
         window: &WindowHandle<T>,
@@ -365,7 +378,12 @@ impl AppContext for AsyncWindowContext {
     where
         T: 'static,
     {
-        self.app.new(build_entity)
+        // Associate the new entity with our captured window so that
+        // `with_window` can resolve a dispatch target before the entity has
+        // been rendered.
+        self.app
+            .update_window(self.window, |_, _, cx| cx.new(build_entity))
+            .expect("window was unexpectedly closed")
     }
 
     fn reserve_entity<T: 'static>(&mut self) -> Reservation<T> {
@@ -377,7 +395,11 @@ impl AppContext for AsyncWindowContext {
         reservation: Reservation<T>,
         build_entity: impl FnOnce(&mut Context<T>) -> T,
     ) -> Entity<T> {
-        self.app.insert_entity(reservation, build_entity)
+        self.app
+            .update_window(self.window, |_, _, cx| {
+                cx.insert_entity(reservation, build_entity)
+            })
+            .expect("window was unexpectedly closed")
     }
 
     fn update_entity<T: 'static, R>(
@@ -409,6 +431,14 @@ impl AppContext for AsyncWindowContext {
         self.app.update_window(window, update)
     }
 
+    fn with_window<R>(
+        &mut self,
+        entity_id: EntityId,
+        f: impl FnOnce(&mut Window, &mut App) -> R,
+    ) -> Option<R> {
+        self.app.with_window(entity_id, f)
+    }
+
     fn read_window<T, R>(
         &self,
         window: &WindowHandle<T>,
@@ -457,9 +487,12 @@ impl VisualContext for AsyncWindowContext {
         view: &Entity<T>,
         update: impl FnOnce(&mut T, &mut Window, &mut Context<T>) -> R,
     ) -> Result<R> {
-        self.app.update_window(self.window, |_, window, cx| {
-            view.update(cx, |entity, cx| update(entity, window, cx))
-        })
+        let view = view.clone();
+        self.app
+            .with_window(view.entity_id(), |window, app| {
+                view.update(app, |entity, cx| update(entity, window, cx))
+            })
+            .context("entity has no current window")
     }
 
     fn replace_root_view<V>(

crates/gpui/src/app/context.rs

@@ -307,9 +307,13 @@ impl<'a, T: 'static> Context<'a, T> {
         window: &Window,
         f: impl FnOnce(&mut T, &mut Window, &mut Context<T>) + 'static,
     ) {
-        let view = self.entity();
-        window.defer(self, move |window, cx| {
-            view.update(cx, |view, cx| f(view, window, cx))
+        let view = self.weak_entity();
+        let entity_id = self.entity_id();
+        self.ensure_window(entity_id, window.handle.id);
+        self.app.defer(move |cx| {
+            cx.with_window(entity_id, |window, cx| {
+                view.update(cx, |view, cx| f(view, window, cx)).ok();
+            });
         });
     }
 
@@ -326,25 +330,21 @@ impl<'a, T: 'static> Context<'a, T> {
     {
         let observed_id = observed.entity_id();
         let observed = observed.downgrade();
-        let window_handle = window.handle;
         let observer = self.weak_entity();
+        let observer_id = self.entity_id();
+        self.ensure_window(observer_id, window.handle.id);
         self.new_observer(
             observed_id,
             Box::new(move |cx| {
-                window_handle
-                    .update(cx, |_, window, cx| {
-                        if let Some((observer, observed)) =
-                            observer.upgrade().zip(observed.upgrade())
-                        {
-                            observer.update(cx, |observer, cx| {
-                                on_notify(observer, observed, window, cx);
-                            });
-                            true
-                        } else {
-                            false
-                        }
-                    })
-                    .unwrap_or(false)
+                let Some((observer, observed)) = observer.upgrade().zip(observed.upgrade()) else {
+                    return false;
+                };
+                cx.with_window(observer_id, |window, cx| {
+                    observer.update(cx, |observer, cx| {
+                        on_notify(observer, observed, window, cx);
+                    });
+                });
+                true
             }),
         )
     }
@@ -363,28 +363,25 @@ impl<'a, T: 'static> Context<'a, T> {
         Evt: 'static,
     {
         let emitter = emitter.downgrade();
-        let window_handle = window.handle;
         let subscriber = self.weak_entity();
+        let subscriber_id = self.entity_id();
+        self.ensure_window(subscriber_id, window.handle.id);
         self.new_subscription(
             emitter.entity_id(),
             (
                 TypeId::of::<Evt>(),
                 Box::new(move |event, cx| {
-                    window_handle
-                        .update(cx, |_, window, cx| {
-                            if let Some((subscriber, emitter)) =
-                                subscriber.upgrade().zip(emitter.upgrade())
-                            {
-                                let event = event.downcast_ref().expect("invalid event type");
-                                subscriber.update(cx, |subscriber, cx| {
-                                    on_event(subscriber, &emitter, event, window, cx);
-                                });
-                                true
-                            } else {
-                                false
-                            }
-                        })
-                        .unwrap_or(false)
+                    let Some((subscriber, emitter)) = subscriber.upgrade().zip(emitter.upgrade())
+                    else {
+                        return false;
+                    };
+                    let event = event.downcast_ref().expect("invalid event type");
+                    cx.with_window(subscriber_id, |window, cx| {
+                        subscriber.update(cx, |subscriber, cx| {
+                            on_event(subscriber, &emitter, event, window, cx);
+                        });
+                    });
+                    true
                 }),
             ),
         )
@@ -835,6 +832,15 @@ impl<T> AppContext for Context<'_, T> {
         self.app.update_window(window, update)
     }
 
+    #[inline]
+    fn with_window<R>(
+        &mut self,
+        entity_id: EntityId,
+        f: impl FnOnce(&mut Window, &mut App) -> R,
+    ) -> Option<R> {
+        self.app.with_window(entity_id, f)
+    }
+
     #[inline]
     fn read_window<U, R>(
         &self,

crates/gpui/src/app/entity_map.rs

@@ -795,14 +795,13 @@ impl<T: 'static> WeakEntity<T> {
         update: impl FnOnce(&mut T, &mut Window, &mut Context<T>) -> R,
     ) -> Result<R>
     where
-        C: VisualContext,
+        C: AppContext,
     {
-        let window = cx.window_handle();
         let entity = self.upgrade().context("entity released")?;
-
-        window.update(cx, |_, window, cx| {
-            entity.update(cx, |entity, cx| update(entity, window, cx))
+        cx.with_window(entity.entity_id(), |window, app| {
+            entity.update(app, |entity, cx| update(entity, window, cx))
         })
+        .context("entity has no current window")
     }
 
     /// Reads the entity referenced by this handle with the given function if

crates/gpui/src/app/headless_app_context.rs

@@ -10,7 +10,7 @@
 
 use crate::{
     AnyView, AnyWindowHandle, App, AppCell, AppContext, AssetSource, BackgroundExecutor, Bounds,
-    Context, Entity, ForegroundExecutor, Global, Pixels, PlatformHeadlessRenderer,
+    Context, Entity, EntityId, ForegroundExecutor, Global, Pixels, PlatformHeadlessRenderer,
     PlatformTextSystem, Render, Reservation, Size, Task, TestDispatcher, TestPlatform, TextSystem,
     Window, WindowBounds, WindowHandle, WindowOptions,
     app::{GpuiBorrow, GpuiMode},
@@ -246,6 +246,15 @@ impl AppContext for HeadlessAppContext {
         lock.update_window(window, f)
     }
 
+    fn with_window<R>(
+        &mut self,
+        entity_id: EntityId,
+        f: impl FnOnce(&mut Window, &mut App) -> R,
+    ) -> Option<R> {
+        let mut lock = self.app.borrow_mut();
+        lock.with_window(entity_id, f)
+    }
+
     fn read_window<T, R>(
         &self,
         window: &WindowHandle<T>,

crates/gpui/src/app/test_context.rs

@@ -1,9 +1,9 @@
 use crate::{
     Action, AnyView, AnyWindowHandle, App, AppCell, AppContext, AsyncApp, AvailableSpace,
     BackgroundExecutor, BorrowAppContext, Bounds, Capslock, ClipboardItem, DrawPhase, Drawable,
-    Element, Empty, EventEmitter, ForegroundExecutor, Global, InputEvent, Keystroke, Modifiers,
-    ModifiersChangedEvent, MouseButton, MouseDownEvent, MouseMoveEvent, MouseUpEvent, Pixels,
-    Platform, Point, Render, Result, Size, Task, TestDispatcher, TestPlatform,
+    Element, Empty, EntityId, EventEmitter, ForegroundExecutor, Global, InputEvent, Keystroke,
+    Modifiers, ModifiersChangedEvent, MouseButton, MouseDownEvent, MouseMoveEvent, MouseUpEvent,
+    Pixels, Platform, Point, Render, Result, Size, Task, TestDispatcher, TestPlatform,
     TestScreenCaptureSource, TestWindow, TextSystem, VisualContext, Window, WindowBounds,
     WindowHandle, WindowOptions, app::GpuiMode, window::ElementArenaScope,
 };
@@ -84,6 +84,15 @@ impl AppContext for TestAppContext {
         lock.update_window(window, f)
     }
 
+    fn with_window<R>(
+        &mut self,
+        entity_id: EntityId,
+        f: impl FnOnce(&mut Window, &mut App) -> R,
+    ) -> Option<R> {
+        let mut lock = self.app.borrow_mut();
+        lock.with_window(entity_id, f)
+    }
+
     fn read_window<T, R>(
         &self,
         window: &WindowHandle<T>,
@@ -193,12 +202,6 @@ impl TestAppContext {
         &self.foreground_executor
     }
 
-    #[expect(clippy::wrong_self_convention)]
-    fn new<T: 'static>(&mut self, build_entity: impl FnOnce(&mut Context<T>) -> T) -> Entity<T> {
-        let mut cx = self.app.borrow_mut();
-        cx.new(build_entity)
-    }
-
     /// Gives you an `&mut App` for the duration of the closure
     pub fn update<R>(&self, f: impl FnOnce(&mut App) -> R) -> R {
         let mut cx = self.app.borrow_mut();
@@ -940,7 +943,9 @@ impl VisualTestContext {
 
 impl AppContext for VisualTestContext {
     fn new<T: 'static>(&mut self, build_entity: impl FnOnce(&mut Context<T>) -> T) -> Entity<T> {
-        self.cx.new(build_entity)
+        self.window
+            .update(&mut self.cx, |_, _, cx| cx.new(build_entity))
+            .expect("window was unexpectedly closed")
     }
 
     fn reserve_entity<T: 'static>(&mut self) -> crate::Reservation<T> {
@@ -952,7 +957,11 @@ impl AppContext for VisualTestContext {
         reservation: crate::Reservation<T>,
         build_entity: impl FnOnce(&mut Context<T>) -> T,
     ) -> Entity<T> {
-        self.cx.insert_entity(reservation, build_entity)
+        self.window
+            .update(&mut self.cx, |_, _, cx| {
+                cx.insert_entity(reservation, build_entity)
+            })
+            .expect("window was unexpectedly closed")
     }
 
     fn update_entity<T, R>(
@@ -987,6 +996,14 @@ impl AppContext for VisualTestContext {
         self.cx.update_window(window, f)
     }
 
+    fn with_window<R>(
+        &mut self,
+        entity_id: EntityId,
+        f: impl FnOnce(&mut Window, &mut App) -> R,
+    ) -> Option<R> {
+        self.cx.with_window(entity_id, f)
+    }
+
     fn read_window<T, R>(
         &self,
         window: &WindowHandle<T>,
@@ -1037,11 +1054,14 @@ impl VisualContext for VisualTestContext {
         view: &Entity<V>,
         update: impl FnOnce(&mut V, &mut Window, &mut Context<V>) -> R,
     ) -> R {
-        self.window
-            .update(&mut self.cx, |_, window, cx| {
-                view.update(cx, |v, cx| update(v, window, cx))
+        let view = view.clone();
+        self.cx
+            .app
+            .borrow_mut()
+            .with_window(view.entity_id(), |window, app| {
+                view.update(app, |v, cx| update(v, window, cx))
             })
-            .expect("window was unexpectedly closed")
+            .expect("entity has no current window; use `update` instead of `update_in`")
     }
 
     fn replace_root_view<V>(

crates/gpui/src/app/visual_test_context.rs

@@ -1,9 +1,9 @@
 use crate::{
     Action, AnyView, AnyWindowHandle, App, AppCell, AppContext, AssetSource, BackgroundExecutor,
-    Bounds, ClipboardItem, Context, Entity, ForegroundExecutor, Global, InputEvent, Keystroke,
-    Modifiers, MouseButton, MouseDownEvent, MouseMoveEvent, MouseUpEvent, Pixels, Platform, Point,
-    Render, Result, Size, Task, TestDispatcher, TextSystem, VisualTestPlatform, Window,
-    WindowBounds, WindowHandle, WindowOptions, app::GpuiMode,
+    Bounds, ClipboardItem, Context, Entity, EntityId, ForegroundExecutor, Global, InputEvent,
+    Keystroke, Modifiers, MouseButton, MouseDownEvent, MouseMoveEvent, MouseUpEvent, Pixels,
+    Platform, Point, Render, Result, Size, Task, TestDispatcher, TextSystem, VisualTestPlatform,
+    Window, WindowBounds, WindowHandle, WindowOptions, app::GpuiMode,
 };
 use anyhow::anyhow;
 use image::RgbaImage;
@@ -446,6 +446,15 @@ impl AppContext for VisualTestAppContext {
         lock.update_window(window, f)
     }
 
+    fn with_window<R>(
+        &mut self,
+        entity_id: EntityId,
+        f: impl FnOnce(&mut Window, &mut App) -> R,
+    ) -> Option<R> {
+        let mut lock = self.app.borrow_mut();
+        lock.with_window(entity_id, f)
+    }
+
     fn read_window<T, R>(
         &self,
         window: &WindowHandle<T>,

crates/gpui/src/gpui.rs

@@ -171,6 +171,16 @@ pub trait AppContext {
     where
         F: FnOnce(AnyView, &mut Window, &mut App) -> T;
 
+    /// Run `f` against the entity's *current* window — the most recently
+    /// rendered window that referenced the entity. Returns `None` if the
+    /// entity has no current window or that window is unavailable. See
+    /// [`App::with_window`] for the underlying lookup.
+    fn with_window<R>(
+        &mut self,
+        entity_id: EntityId,
+        f: impl FnOnce(&mut Window, &mut App) -> R,
+    ) -> Option<R>;
+
     /// Read a window off of the application context.
     fn read_window<T, R>(
         &self,

crates/gpui_macros/src/derive_app_context.rs

@@ -79,6 +79,15 @@ pub fn derive_app_context(input: TokenStream) -> TokenStream {
                 self.#app_variable.update_window(window, f)
             }
 
+            fn with_window<R>(
+                &mut self,
+                entity_id: gpui::EntityId,
+                f: impl FnOnce(&mut gpui::Window, &mut gpui::App) -> R,
+            ) -> Option<R>
+            {
+                self.#app_variable.with_window(entity_id, f)
+            }
+
             fn read_window<T, R>(
                 &self,
                 window: &gpui::WindowHandle<T>,

crates/sidebar/src/sidebar_tests.rs

@@ -10981,7 +10981,7 @@ async fn test_remote_archive_thread_with_disconnected_remote(
     // Disconnect the remote connection before archiving. We don't
     // `run_until_parked` here because the disconnect itself triggers
     // reconnection work that can't complete in the test environment.
-    remote_client.update_in(cx, |client, _window, cx| {
+    remote_client.update(cx, |client, cx| {
         client.simulate_disconnect(cx).detach();
     });