diff --git a/crates/agent_ui/src/agent_configuration.rs b/crates/agent_ui/src/agent_configuration.rs
index 81de185db7c8586a0db67d8da49d6f2dc59d2e8e..e667a6c1d0507f399b4339c84fec4fc7099eab4e 100644
--- a/crates/agent_ui/src/agent_configuration.rs
+++ b/crates/agent_ui/src/agent_configuration.rs
@@ -983,7 +983,7 @@ impl AgentConfiguration {
                         })
                         .entry("Add Custom Agent", None, {
                             move |window, cx| {
-                                if let Some(workspace) = window.root().flatten() {
+                                if let Some(workspace) = Workspace::for_window(window, cx) {
                                     let workspace = workspace.downgrade();
                                     window
                                         .spawn(cx, async |cx| {
diff --git a/crates/edit_prediction_ui/src/edit_prediction_button.rs b/crates/edit_prediction_ui/src/edit_prediction_button.rs
index ac579d33ab42ff11a58e850639c1850c88262f74..729b901be1556f011c101258d34af9b98b45f272 100644
--- a/crates/edit_prediction_ui/src/edit_prediction_button.rs
+++ b/crates/edit_prediction_ui/src/edit_prediction_button.rs
@@ -993,7 +993,7 @@ impl EditPredictionButton {
                         "Edit Prediction Menu Action",
                         action = "configure_excluded_files",
                     );
-                    if let Some(workspace) = window.root().flatten() {
+                    if let Some(workspace) = Workspace::for_window(window, cx) {
                         let workspace = workspace.downgrade();
                         window
                             .spawn(cx, async |cx| {