diff --git a/sway/desktop/xwayland.c b/sway/desktop/xwayland.c
index db21dc78..89b1a481 100644
--- a/sway/desktop/xwayland.c
+++ b/sway/desktop/xwayland.c
@@ -92,20 +92,15 @@ static void unmanaged_handle_unmap(struct wl_listener *listener, void *data) {
 	wl_list_remove(&surface->commit.link);
 
 	struct sway_seat *seat = input_manager_current_seat();
-	if (seat->wlr_seat->keyboard_state.focused_surface ==
-			xsurface->surface) {
-
-		// Try to find another unmanaged surface from the same process to pass
-		// focus to. This is necessary because some applications (e.g. Jetbrains
-		// IDEs) represent their multi-level menus as unmanaged surfaces, and
-		// when closing a submenu, the main menu should get input focus.
-		struct sway_xwayland_unmanaged *current;
-		wl_list_for_each(current, &root->xwayland_unmanaged, link) {
-			struct wlr_xwayland_surface *prev_xsurface =
-				current->wlr_xwayland_surface;
-			if (prev_xsurface->pid == xsurface->pid &&
-					wlr_xwayland_or_surface_wants_focus(prev_xsurface)) {
-				seat_set_focus_surface(seat, prev_xsurface->surface, false);
+	if (seat->wlr_seat->keyboard_state.focused_surface == xsurface->surface) {
+		// This simply returns focus to the parent surface if there's one available.
+		// This seems to handle JetBrains issues.
+		if (xsurface->parent &&
+				wlr_surface_is_xwayland_surface(xsurface->parent->surface)) {
+			struct wlr_xwayland_surface *next_surface =
+				wlr_xwayland_surface_from_wlr_surface(xsurface->parent->surface);
+			if (wlr_xwayland_or_surface_wants_focus(next_surface)) {
+				seat_set_focus_surface(seat, xsurface->parent->surface, false);
 				return;
 			}
 		}