@@ -198,6 +198,7 @@ void TopLevelWindowModel::prependSurface(unityapi::MirSurfaceInterface *surface,
198198 ModelEntry &entry = m_windowModel[i];
199199 if (entry.application == application && entry.window->surface() == nullptr) {
200200 entry.window->setSurface(surface);
201+ m_allSurfaces.insert(surface);
201202 connectSurface(surface);
202203 INFO_MSG << " appId=" << application->appId() << " surface=" << surface
203204 << ", filling out placeholder. after: " << toString();
@@ -224,8 +225,8 @@ void TopLevelWindowModel::prependSurfaceHelper(unityapi::MirSurfaceInterface *su
224225 Window *window = createWindow(surface);
225226
226227 m_windowModel.prepend(ModelEntry(window, application));
227− m_allSurfaces.insert(surface);
228228 if (surface) {
229+ m_allSurfaces.insert(surface);
229230 connectSurface(surface);
230231 }
231232
@@ -436,7 +437,7 @@ void TopLevelWindowModel::onSurfacesAboutToBeRemovedFromWorkspace(const std::sha
436437 m_allSurfaces.remove(surface);
437438 continue;
438439 }
439− while(iter != surfaces.constEnd()) {
440+ while(iter != surfaces.constEnd()) {
440441 int index = indexOf(*iter);
441442 if (index != end+1) {
442443 break;
@@ -459,7 +460,7 @@ void TopLevelWindowModel::onSurfacesAboutToBeRemovedFromWorkspace(const std::sha
459460 window->setFocused(false);
460461
461462 m_windowModel.removeAt(start);
462− m_allSurfaces.remove(window->surface());
463+ m_allSurfaces.remove(surface);
463464 }
464465
465466 if (m_modelState == RemovingState) {
@@ -482,12 +483,13 @@ void TopLevelWindowModel::removeAt(int index)
482483 }
483484
484485 auto window = m_windowModel[index].window;
486+ auto surface = window->surface();
485487
486488 window->setSurface(nullptr);
487489 window->setFocused(false);
488490
489491 m_windowModel.removeAt(index);
490− m_allSurfaces.remove(window->surface());
492+ m_allSurfaces.remove(surface);
491493
492494 if (m_modelState == RemovingState) {
493495 endRemoveRows();
@@ -794,34 +796,32 @@ void TopLevelWindowModel::refreshWindows()
794796 m_windowModel.clear();
795797 m_allSurfaces.clear();
796798
797− if (!m_workspace || !m_applicationManager) return;
798− WMPolicyInterface::instance()->forEachWindowInWorkspace(m_workspace->workspace(), [this](const miral::Window &window) {
799− auto surface = m_surfaceManager->surfaceFor(window);
800− if (surface) {
801− if (surface->parentSurface()) {
802− // Wrap it in a Window so that we keep focusedWindow() up to date.
803− Window *window = createWindow(surface);
804− connect(surface, &QObject::destroyed, window, [=](){
805− window->setSurface(nullptr);
806− window->deleteLater();
807− });
799+ if (!m_workspace || !m_applicationManager || !m_surfaceManager) return;
800+
801+ m_surfaceManager->forEachSurfaceInWorkspace(m_workspace->workspace(), [this](unity::shell::application::MirSurfaceInterface* surface) {
802+ if (surface->parentSurface()) {
803+ // Wrap it in a Window so that we keep focusedWindow() up to date.
804+ Window *window = createWindow(surface);
805+ connect(surface, &QObject::destroyed, window, [=](){
806+ window->setSurface(nullptr);
807+ window->deleteLater();
808+ });
809+ } else {
810+ if (surface->type() == Mir::InputMethodType) {
811+ setInputMethodWindow(createWindow(surface));
808812 } else {
809− if (surface->type() == Mir::InputMethodType) {
810− setInputMethodWindow(createWindow(surface));
813+ auto *application = m_applicationManager->findApplicationWithSurface(surface);
814+ if (application) {
815+ prependSurface(surface, application);
811816 } else {
812− auto *application = m_applicationManager->findApplicationWithSurface(surface);
813− if (application) {
814− prependSurface(surface, application);
815− } else {
816− // Must be a prompt session. No need to do add it as a prompt surface is not top-level.
817− // It will show up in the ApplicationInfoInterface::promptSurfaceList of some application.
818− // Still wrap it in a Window though, so that we keep focusedWindow() up to date.
819− Window *promptWindow = createWindow(surface);
820− connect(surface, &QObject::destroyed, promptWindow, [=](){
821− promptWindow->setSurface(nullptr);
822− promptWindow->deleteLater();
823− });
824− }
817+ // Must be a prompt session. No need to do add it as a prompt surface is not top-level.
818+ // It will show up in the ApplicationInfoInterface::promptSurfaceList of some application.
819+ // Still wrap it in a Window though, so that we keep focusedWindow() up to date.
820+ Window *promptWindow = createWindow(surface);
821+ connect(surface, &QObject::destroyed, promptWindow, [=](){
822+ promptWindow->setSurface(nullptr);
823+ promptWindow->deleteLater();
824+ });
825825 }
826826 }
827827 }