diff --git a/crates/agent_ui/src/agent_configuration.rs b/crates/agent_ui/src/agent_configuration.rs
index 81de185db7c..e667a6c1d05 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 ac579d33ab4..729b901be15 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| {