Moves EventDispatcher from Display to WindowManagerState

EventDispatcher maintains state that is per WindowManager
(accelerators, root...), so it makes sense that the EventDispatcher is
owned by the WindowManagerState.

This patch also makes Display update WindowManagerState and
EventDispatcher appropriately when the active user changes. When the
active user changes capture and any other event related state is reset.

BUG=590955
TEST=covered by tests
R=ben@chromium.org

Review URL: https://codereview.chromium.org/1775133003

Cr-Commit-Position: refs/heads/master@{#380005}

23 files changed, 742 insertions, 408 deletions

diff --git a/components/mus/ws/connection_manager.cc b/components/mus/ws/connection_manager.cc
index a710d9f..4ef8461 100644
--- a/components/mus/ws/connection_manager.cc
+++ b/components/mus/ws/connection_manager.cc
@@ -482,7 +482,6 @@ void ConnectionManager::OnWindowBoundsChanged(ServerWindow* window,
   if (!window->parent())
     return;
 
-  // TODO(sky): optimize this.
   SchedulePaint(window->parent(), old_bounds);
   SchedulePaint(window->parent(), new_bounds);
 
@@ -572,11 +571,6 @@ void ConnectionManager::OnTransientWindowRemoved(
   }
 }
 
-void ConnectionManager::OnWillDestroyDisplay(Display* display) {
-  for (auto& pair : tree_map_)
-    pair.second->OnWillDestroyDisplay(display);
-}
-
 void ConnectionManager::OnFirstDisplayReady() {
   delegate_->OnFirstDisplayReady();
 }
diff --git a/components/mus/ws/connection_manager.h b/components/mus/ws/connection_manager.h
index b6d81f7..fe3bc95 100644
--- a/components/mus/ws/connection_manager.h
+++ b/components/mus/ws/connection_manager.h
@@ -40,9 +40,6 @@ class WindowTreeBinding;
 
 // ConnectionManager manages the set of connections to the window server (all
 // the WindowTrees) as well as providing the root of the hierarchy.
-//
-// TODO(sky): this class is doing too much. Refactor to make responsibilities
-// clearer.
 class ConnectionManager : public ServerWindowDelegate,
                           public ServerWindowObserver,
                           public DisplayManagerDelegate {
@@ -54,6 +51,7 @@ class ConnectionManager : public ServerWindowDelegate,
   ConnectionManagerDelegate* delegate() { return delegate_; }
 
   UserIdTracker* user_id_tracker() { return &user_id_tracker_; }
+  const UserIdTracker* user_id_tracker() const { return &user_id_tracker_; }
 
   DisplayManager* display_manager() { return display_manager_.get(); }
   const DisplayManager* display_manager() const {
@@ -256,7 +254,6 @@ class ConnectionManager : public ServerWindowDelegate,
                                 ServerWindow* transient_child) override;
 
   // DisplayManagerDelegate:
-  void OnWillDestroyDisplay(Display* display) override;
   void OnFirstDisplayReady() override;
   void OnNoMoreDisplays() override;
 
diff --git a/components/mus/ws/display.cc b/components/mus/ws/display.cc
index ef878ea..b2a97fe 100644
--- a/components/mus/ws/display.cc
+++ b/components/mus/ws/display.cc
@@ -19,72 +19,10 @@
 #include "components/mus/ws/window_tree_binding.h"
 #include "mojo/common/common_type_converters.h"
 #include "mojo/converters/geometry/geometry_type_converters.h"
-#include "mojo/converters/input_events/input_events_type_converters.h"
 #include "mojo/shell/public/interfaces/connector.mojom.h"
 
 namespace mus {
 namespace ws {
-namespace {
-
-base::TimeDelta GetDefaultAckTimerDelay() {
-#if defined(NDEBUG)
-  return base::TimeDelta::FromMilliseconds(100);
-#else
-  return base::TimeDelta::FromMilliseconds(1000);
-#endif
-}
-
-bool EventsCanBeCoalesced(const ui::Event& one, const ui::Event& two) {
-  if (one.type() != two.type() || one.flags() != two.flags())
-    return false;
-
-  // TODO(sad): wheel events can also be merged.
-  if (one.type() != ui::ET_POINTER_MOVED)
-    return false;
-
-  return one.AsPointerEvent()->pointer_id() ==
-         two.AsPointerEvent()->pointer_id();
-}
-
-scoped_ptr<ui::Event> CoalesceEvents(scoped_ptr<ui::Event> first,
-                                     scoped_ptr<ui::Event> second) {
-  DCHECK(first->type() == ui::ET_POINTER_MOVED)
-      << " Non-move events cannot be merged yet.";
-  // For mouse moves, the new event just replaces the old event.
-  return second;
-}
-
-}  // namespace
-
-class Display::ProcessedEventTarget {
- public:
-  ProcessedEventTarget(ServerWindow* window, bool in_nonclient_area)
-      : in_nonclient_area_(in_nonclient_area) {
-    tracker_.Add(window);
-  }
-
-  ~ProcessedEventTarget() {}
-
-  // Return true if the event is still valid. The event becomes invalid if
-  // the window is destroyed while waiting to dispatch.
-  bool IsValid() const { return !tracker_.windows().empty(); }
-
-  ServerWindow* window() {
-    DCHECK(IsValid());
-    return tracker_.windows().front();
-  }
-
-  bool in_nonclient_area() const { return in_nonclient_area_; }
-
- private:
-  ServerWindowTracker tracker_;
-  const bool in_nonclient_area_;
-
-  DISALLOW_COPY_AND_ASSIGN(ProcessedEventTarget);
-};
-
-Display::QueuedEvent::QueuedEvent() {}
-Display::QueuedEvent::~QueuedEvent() {}
 
 Display::Display(ConnectionManager* connection_manager,
                  mojo::Connector* connector,
@@ -92,17 +30,18 @@ Display::Display(ConnectionManager* connection_manager,
                  const scoped_refptr<SurfacesState>& surfaces_state)
     : id_(connection_manager->display_manager()->GetAndAdvanceNextDisplayId()),
       connection_manager_(connection_manager),
-      event_dispatcher_(this),
       platform_display_(
           PlatformDisplay::Create(connector, gpu_state, surfaces_state)),
-      tree_awaiting_input_ack_(nullptr),
       last_cursor_(0) {
   platform_display_->Init(this);
 
   connection_manager_->window_manager_factory_registry()->AddObserver(this);
+  connection_manager_->user_id_tracker()->AddObserver(this);
 }
 
 Display::~Display() {
+  connection_manager_->user_id_tracker()->RemoveObserver(this);
+
   connection_manager_->window_manager_factory_registry()->RemoveObserver(this);
 
   DestroyFocusController();
@@ -213,18 +152,17 @@ const WindowManagerState* Display::GetWindowManagerStateForUser(
   return iter == window_manager_state_map_.end() ? nullptr : iter->second.get();
 }
 
-void Display::SetCapture(ServerWindow* window, bool in_nonclient_area) {
-  ServerWindow* capture_window = event_dispatcher_.capture_window();
-  if (capture_window == window)
-    return;
-  event_dispatcher_.SetCaptureWindow(window, in_nonclient_area);
-}
-
 mojom::Rotation Display::GetRotation() const {
   return platform_display_->GetRotation();
 }
 
+const WindowManagerState* Display::GetActiveWindowManagerState() const {
+  return GetWindowManagerStateForUser(
+      connection_manager_->user_id_tracker()->active_id());
+}
+
 void Display::SetFocusedWindow(ServerWindow* new_focused_window) {
+  // TODO(sky): this is wrong. Focus is global, not per Display.
   ServerWindow* old_focused_window = focus_controller_->GetFocusedWindow();
   if (old_focused_window == new_focused_window)
     return;
@@ -276,21 +214,30 @@ void Display::OnWillDestroyTree(WindowTree* tree) {
     }
   }
 
-  if (tree_awaiting_input_ack_ != tree)
-    return;
-  // The WindowTree is dying. So it's not going to ack the event.
-  OnEventAck(tree_awaiting_input_ack_);
+  for (const auto& pair : window_manager_state_map_)
+    pair.second->OnWillDestroyTree(tree);
+}
+
+void Display::UpdateNativeCursor(int32_t cursor_id) {
+  if (cursor_id != last_cursor_) {
+    platform_display_->SetCursorById(cursor_id);
+    last_cursor_ = cursor_id;
+  }
 }
 
 void Display::OnCursorUpdated(ServerWindow* window) {
-  if (window == event_dispatcher_.mouse_cursor_source_window())
+  WindowManagerState* wms = GetActiveWindowManagerState();
+  if (wms && window == wms->event_dispatcher()->mouse_cursor_source_window())
     UpdateNativeCursor(window->cursor());
 }
 
 void Display::MaybeChangeCursorOnWindowTreeChange() {
-  event_dispatcher_.UpdateCursorProviderByLastKnownLocation();
+  WindowManagerState* wms = GetActiveWindowManagerState();
+  if (!wms)
+    return;
+  wms->event_dispatcher()->UpdateCursorProviderByLastKnownLocation();
   ServerWindow* cursor_source_window =
-      event_dispatcher_.mouse_cursor_source_window();
+      wms->event_dispatcher()->mouse_cursor_source_window();
   if (cursor_source_window)
     UpdateNativeCursor(cursor_source_window->cursor());
 }
@@ -303,24 +250,14 @@ void Display::SetTitle(const mojo::String& title) {
   platform_display_->SetTitle(title.To<base::string16>());
 }
 
-void Display::OnEventAck(mojom::WindowTree* tree) {
-  if (tree_awaiting_input_ack_ != tree) {
-    // TODO(sad): The ack must have arrived after the timeout. We should do
-    // something here, and in OnEventAckTimeout().
-    return;
-  }
-  tree_awaiting_input_ack_ = nullptr;
-  event_ack_timer_.Stop();
-  ProcessNextEventFromQueue();
-}
-
 void Display::InitWindowManagersIfNecessary() {
   if (!init_called_ || !root_)
     return;
 
   display_manager()->OnDisplayAcceleratedWidgetAvailable(this);
   if (binding_) {
-    scoped_ptr<WindowManagerState> wms_ptr(new WindowManagerState(this));
+    scoped_ptr<WindowManagerState> wms_ptr(new WindowManagerState(
+        this, platform_display_.get(), top_level_surface_id_));
     WindowManagerState* wms = wms_ptr.get();
     // For this case we never create additional WindowManagerStates, so any
     // id works.
@@ -332,77 +269,6 @@ void Display::InitWindowManagersIfNecessary() {
   }
 }
 
-void Display::OnEventAckTimeout() {
-  // TODO(sad): Figure out what we should do.
-  NOTIMPLEMENTED() << "Event ACK timed out.";
-  OnEventAck(tree_awaiting_input_ack_);
-}
-
-void Display::QueueEvent(
-    const ui::Event& event,
-    scoped_ptr<ProcessedEventTarget> processed_event_target) {
-  scoped_ptr<QueuedEvent> queued_event(new QueuedEvent);
-  queued_event->event = ui::Event::Clone(event);
-  queued_event->processed_target = std::move(processed_event_target);
-  event_queue_.push(std::move(queued_event));
-}
-
-void Display::ProcessNextEventFromQueue() {
-  // Loop through |event_queue_| stopping after dispatching the first valid
-  // event.
-  while (!event_queue_.empty()) {
-    scoped_ptr<QueuedEvent> queued_event = std::move(event_queue_.front());
-    event_queue_.pop();
-    if (!queued_event->processed_target) {
-      event_dispatcher_.ProcessEvent(*queued_event->event);
-      return;
-    }
-    if (queued_event->processed_target->IsValid()) {
-      DispatchInputEventToWindowImpl(
-          queued_event->processed_target->window(),
-          queued_event->processed_target->in_nonclient_area(),
-          *queued_event->event);
-      return;
-    }
-  }
-}
-
-void Display::DispatchInputEventToWindowImpl(ServerWindow* target,
-                                             bool in_nonclient_area,
-                                             const ui::Event& event) {
-  if (target == root_.get()) {
-    // TODO(sky): use active windowmanager here.
-    target = GetFirstWindowManagerState()->root();
-  }
-  if (event.IsMousePointerEvent()) {
-    DCHECK(event_dispatcher_.mouse_cursor_source_window());
-    UpdateNativeCursor(
-        event_dispatcher_.mouse_cursor_source_window()->cursor());
-  }
-
-  // If the event is in the non-client area the event goes to the owner of
-  // the window. Otherwise if the window is an embed root, forward to the
-  // embedded window.
-  WindowTree* tree =
-      in_nonclient_area
-          ? connection_manager_->GetTreeWithId(target->id().connection_id)
-          : connection_manager_->GetTreeWithRoot(target);
-  if (!tree) {
-    DCHECK(!in_nonclient_area);
-    tree = connection_manager_->GetTreeWithId(target->id().connection_id);
-  }
-
-  // TOOD(sad): Adjust this delay, possibly make this dynamic.
-  const base::TimeDelta max_delay = base::debug::BeingDebugged()
-                                        ? base::TimeDelta::FromDays(1)
-                                        : GetDefaultAckTimerDelay();
-  event_ack_timer_.Start(FROM_HERE, max_delay, this,
-                         &Display::OnEventAckTimeout);
-
-  tree_awaiting_input_ack_ = tree;
-  tree->DispatchInputEvent(target, mojom::Event::From(event));
-}
-
 void Display::CreateWindowManagerStatesFromRegistry() {
   std::vector<WindowManagerFactoryService*> services =
       connection_manager_->window_manager_factory_registry()->GetServices();
@@ -415,43 +281,28 @@ void Display::CreateWindowManagerStatesFromRegistry() {
 void Display::CreateWindowManagerStateFromService(
     WindowManagerFactoryService* service) {
   scoped_ptr<WindowManagerState> wms_ptr(
-      new WindowManagerState(this, service->user_id()));
+      new WindowManagerState(this, platform_display_.get(),
+                             top_level_surface_id_, service->user_id()));
   WindowManagerState* wms = wms_ptr.get();
   window_manager_state_map_[service->user_id()] = std::move(wms_ptr);
   wms->tree_ = connection_manager_->CreateTreeForWindowManager(
       this, service->window_manager_factory(), wms->root());
 }
 
-void Display::UpdateNativeCursor(int32_t cursor_id) {
-  if (cursor_id != last_cursor_) {
-    platform_display_->SetCursorById(cursor_id);
-    last_cursor_ = cursor_id;
-  }
-}
-
 ServerWindow* Display::GetRootWindow() {
   return root_.get();
 }
 
 void Display::OnEvent(const ui::Event& event) {
-  mojom::EventPtr mojo_event(mojom::Event::From(event));
-  // If this is still waiting for an ack from a previously sent event, then
-  // queue up the event to be dispatched once the ack is received.
-  if (event_ack_timer_.IsRunning()) {
-    if (!event_queue_.empty() && !event_queue_.back()->processed_target &&
-        EventsCanBeCoalesced(*event_queue_.back()->event, event)) {
-      event_queue_.back()->event = CoalesceEvents(
-          std::move(event_queue_.back()->event), ui::Event::Clone(event));
-      return;
-    }
-    QueueEvent(event, nullptr);
-    return;
-  }
-  event_dispatcher_.ProcessEvent(event);
+  WindowManagerState* wms = GetActiveWindowManagerState();
+  if (wms)
+    wms->ProcessEvent(event);
 }
 
 void Display::OnNativeCaptureLost() {
-  SetCapture(nullptr, false);
+  WindowManagerState* state = GetActiveWindowManagerState();
+  if (state)
+    state->SetCapture(nullptr, false);
 }
 
 void Display::OnDisplayClosed() {
@@ -470,7 +321,6 @@ void Display::OnViewportMetricsChanged(
     focus_controller_.reset(new FocusController(this, root_.get()));
     focus_controller_->AddObserver(this);
     InitWindowManagersIfNecessary();
-    event_dispatcher_.set_root(root_.get());
   } else {
     root_->SetBounds(gfx::Rect(new_metrics.size_in_pixels.To<gfx::Size>()));
     const gfx::Rect wm_bounds(root_->bounds().size());
@@ -484,7 +334,9 @@ void Display::OnViewportMetricsChanged(
 }
 
 void Display::OnTopLevelSurfaceChanged(cc::SurfaceId surface_id) {
-  event_dispatcher_.set_surface_id(surface_id);
+  DCHECK(!root_);
+  // This should only be called once, and before we've created root_.
+  top_level_surface_id_ = surface_id;
 }
 
 void Display::OnCompositorFrameDrawn() {
@@ -594,55 +446,28 @@ void Display::OnFocusChanged(FocusControllerChangeSource change_source,
                        new_focused_window->text_input_state());
 }
 
-void Display::OnAccelerator(uint32_t accelerator_id, const ui::Event& event) {
-  // TODO(sky): accelerators need to be maintained per windowmanager and pushed
-  // to the eventdispatcher when the active userid changes.
-  GetFirstWindowManagerState()->tree()->OnAccelerator(
-      accelerator_id, mojom::Event::From(event));
-}
-
-void Display::SetFocusedWindowFromEventDispatcher(
-    ServerWindow* new_focused_window) {
-  SetFocusedWindow(new_focused_window);
-}
-
-ServerWindow* Display::GetFocusedWindowForEventDispatcher() {
-  return GetFocusedWindow();
-}
-
-void Display::SetNativeCapture() {
-  platform_display_->SetCapture();
-}
-
-void Display::ReleaseNativeCapture() {
-  platform_display_->ReleaseCapture();
-}
-
-void Display::OnServerWindowCaptureLost(ServerWindow* window) {
-  DCHECK(window);
-  connection_manager_->ProcessLostCapture(window);
-}
-
-void Display::DispatchInputEventToWindow(ServerWindow* target,
-                                         bool in_nonclient_area,
-                                         const ui::Event& event) {
-  if (event_ack_timer_.IsRunning()) {
-    scoped_ptr<ProcessedEventTarget> processed_event_target(
-        new ProcessedEventTarget(target, in_nonclient_area));
-    QueueEvent(event, std::move(processed_event_target));
-    return;
-  }
-
-  DispatchInputEventToWindowImpl(target, in_nonclient_area, event);
-}
-
 void Display::OnWindowDestroyed(ServerWindow* window) {
   windows_needing_frame_destruction_.erase(window);
   window->RemoveObserver(this);
 }
 
-void Display::OnActiveUserIdChanged(const UserId& id) {
-  // TODO(sky): this likely needs to cancel any pending events and all that.
+void Display::OnActiveUserIdChanged(const UserId& previously_active_id,
+                                    const UserId& active_id) {
+  WindowManagerState* previous_wms =
+      GetWindowManagerStateForUser(previously_active_id);
+  const gfx::Point mouse_location_on_screen =
+      previous_wms
+          ? previous_wms->event_dispatcher()->mouse_pointer_last_location()
+          : gfx::Point();
+  if (previous_wms)
+    previous_wms->event_dispatcher()->Reset();
+
+  WindowManagerState* active_wms = GetWindowManagerStateForUser(active_id);
+  if (active_wms) {
+    active_wms->event_dispatcher()->Reset();
+    active_wms->event_dispatcher()->SetMousePointerScreenLocation(
+        mouse_location_on_screen);
+  }
 }
 
 void Display::OnUserIdAdded(const UserId& id) {}
diff --git a/components/mus/ws/display.h b/components/mus/ws/display.h
index 69cfb9e..1ec923f 100644
--- a/components/mus/ws/display.h
+++ b/components/mus/ws/display.h
@@ -16,8 +16,6 @@
 #include "components/mus/common/types.h"
 #include "components/mus/public/interfaces/window_manager_constants.mojom.h"
 #include "components/mus/public/interfaces/window_tree_host.mojom.h"
-#include "components/mus/ws/event_dispatcher.h"
-#include "components/mus/ws/event_dispatcher_delegate.h"
 #include "components/mus/ws/focus_controller_delegate.h"
 #include "components/mus/ws/focus_controller_observer.h"
 #include "components/mus/ws/platform_display.h"
@@ -54,7 +52,6 @@ class Display : public PlatformDisplayDelegate,
                 public mojom::WindowTreeHost,
                 public FocusControllerObserver,
                 public FocusControllerDelegate,
-                public EventDispatcherDelegate,
                 public ServerWindowObserver,
                 public UserIdTrackerObserver,
                 public WindowManagerFactoryRegistryObserver {
@@ -97,8 +94,6 @@ class Display : public PlatformDisplayDelegate,
 
   ConnectionManager* connection_manager() { return connection_manager_; }
 
-  EventDispatcher* event_dispatcher() { return &event_dispatcher_; }
-
   // Returns the root of the Display. The root's children are the roots
   // of the corresponding WindowManagers.
   ServerWindow* root_window() { return root_.get(); }
@@ -116,12 +111,15 @@ class Display : public PlatformDisplayDelegate,
   }
   const WindowManagerState* GetWindowManagerStateForUser(
       const UserId& user_id) const;
+  WindowManagerState* GetActiveWindowManagerState() {
+    return const_cast<WindowManagerState*>(
+        const_cast<const Display*>(this)->GetActiveWindowManagerState());
+  }
+  const WindowManagerState* GetActiveWindowManagerState() const;
   size_t num_window_manger_states() const {
     return window_manager_state_map_.size();
   }
 
-  void SetCapture(ServerWindow* window, bool in_nonclient_area);
-
   void SetFocusedWindow(ServerWindow* window);
   ServerWindow* GetFocusedWindow();
   void DestroyFocusController();
@@ -134,17 +132,11 @@ class Display : public PlatformDisplayDelegate,
                             const ui::TextInputState& state);
   void SetImeVisibility(ServerWindow* window, bool visible);
 
-  // Returns the window that has captured input.
-  ServerWindow* GetCaptureWindow() {
-    return event_dispatcher_.capture_window();
-  }
-  const ServerWindow* GetCaptureWindow() const {
-    return event_dispatcher_.capture_window();
-  }
-
   // Called just before |tree| is destroyed.
   void OnWillDestroyTree(WindowTree* tree);
 
+  void UpdateNativeCursor(int32_t cursor_id);
+
   // Called when a client updates a cursor. This will update the cursor on the
   // native display if the cursor is currently under |window|.
   void OnCursorUpdated(ServerWindow* window);
@@ -158,50 +150,15 @@ class Display : public PlatformDisplayDelegate,
   void SetSize(mojo::SizePtr size) override;
   void SetTitle(const mojo::String& title) override;
 
-  void OnEventAck(mojom::WindowTree* tree);
-
  private:
-  class ProcessedEventTarget;
   friend class test::DisplayTestApi;
 
   using WindowManagerStateMap =
       std::map<UserId, scoped_ptr<WindowManagerState>>;
 
-  // There are two types of events that may be queued, both occur only when
-  // waiting for an ack from a client.
-  // . We get an event from the PlatformDisplay. This results in |event| being
-  //   set, but |processed_target| is null.
-  // . We get an event from the EventDispatcher. In this case both |event| and
-  //   |processed_target| are valid.
-  // The second case happens if EventDispatcher generates more than one event
-  // at a time.
-  struct QueuedEvent {
-    QueuedEvent();
-    ~QueuedEvent();
-
-    scoped_ptr<ui::Event> event;
-    scoped_ptr<ProcessedEventTarget> processed_target;
-  };
-
   // Inits the necessary state once the display is ready.
   void InitWindowManagersIfNecessary();
 
-  void OnEventAckTimeout();
-
-  // Schedules an event to be processed later.
-  void QueueEvent(const ui::Event& event,
-                  scoped_ptr<ProcessedEventTarget> processed_event_target);
-
-  // Processes the next valid event in |event_queue_|. If the event has already
-  // been processed it is dispatched, otherwise the event is passed to the
-  // EventDispatcher for processing.
-  void ProcessNextEventFromQueue();
-
-  // Dispatches the event to the appropriate client and starts the ack timer.
-  void DispatchInputEventToWindowImpl(ServerWindow* target,
-                                      bool in_nonclient_area,
-                                      const ui::Event& event);
-
   // Creates the set of WindowManagerStates from the
   // WindowManagerFactoryRegistry.
   void CreateWindowManagerStatesFromRegistry();
@@ -209,8 +166,6 @@ class Display : public PlatformDisplayDelegate,
   void CreateWindowManagerStateFromService(
       WindowManagerFactoryService* service);
 
-  void UpdateNativeCursor(int32_t cursor_id);
-
   // PlatformDisplayDelegate:
   ServerWindow* GetRootWindow() override;
   void OnEvent(const ui::Event& event) override;
@@ -232,22 +187,12 @@ class Display : public PlatformDisplayDelegate,
                       ServerWindow* old_focused_window,
                       ServerWindow* new_focused_window) override;
 
-  // EventDispatcherDelegate:
-  void OnAccelerator(uint32_t accelerator_id, const ui::Event& event) override;
-  void SetFocusedWindowFromEventDispatcher(ServerWindow* window) override;
-  ServerWindow* GetFocusedWindowForEventDispatcher() override;
-  void SetNativeCapture() override;
-  void ReleaseNativeCapture() override;
-  void OnServerWindowCaptureLost(ServerWindow* window) override;
-  void DispatchInputEventToWindow(ServerWindow* target,
-                                  bool in_nonclient_area,
-                                  const ui::Event& event) override;
-
   // ServerWindowObserver:
   void OnWindowDestroyed(ServerWindow* window) override;
 
   // UserIdTrackerObserver:
-  void OnActiveUserIdChanged(const UserId& id) override;
+  void OnActiveUserIdChanged(const UserId& previously_active_id,
+                             const UserId& active_id) override;
   void OnUserIdAdded(const UserId& id) override;
   void OnUserIdRemoved(const UserId& id) override;
 
@@ -259,11 +204,9 @@ class Display : public PlatformDisplayDelegate,
   // Set once Init() has been called.
   bool init_called_ = false;
   ConnectionManager* const connection_manager_;
-  EventDispatcher event_dispatcher_;
   scoped_ptr<ServerWindow> root_;
   scoped_ptr<PlatformDisplay> platform_display_;
   scoped_ptr<FocusController> focus_controller_;
-  mojom::WindowTree* tree_awaiting_input_ack_;
 
   // The last cursor set. Used to track whether we need to change the cursor.
   int32_t last_cursor_;
@@ -274,11 +217,10 @@ class Display : public PlatformDisplayDelegate,
   // draws.
   std::set<ServerWindow*> windows_needing_frame_destruction_;
 
-  std::queue<scoped_ptr<QueuedEvent>> event_queue_;
-  base::OneShotTimer event_ack_timer_;
-
   WindowManagerStateMap window_manager_state_map_;
 
+  cc::SurfaceId top_level_surface_id_;
+
   DISALLOW_COPY_AND_ASSIGN(Display);
 };
 
diff --git a/components/mus/ws/display_manager.cc b/components/mus/ws/display_manager.cc
index 7b321e6..b5dbb09 100644
--- a/components/mus/ws/display_manager.cc
+++ b/components/mus/ws/display_manager.cc
@@ -39,8 +39,6 @@ void DisplayManager::AddDisplay(Display* display) {
 }
 
 void DisplayManager::DestroyDisplay(Display* display) {
-  delegate_->OnWillDestroyDisplay(display);
-
   if (pending_displays_.count(display)) {
     pending_displays_.erase(display);
   } else {
@@ -89,11 +87,6 @@ const Display* DisplayManager::GetDisplayContaining(
   return nullptr;
 }
 
-Display* DisplayManager::GetActiveDisplay() {
-  // TODO(sky): this isn't active, but first. Make it active.
-  return displays_.size() ? *displays_.begin() : nullptr;
-}
-
 WindowManagerAndDisplayConst DisplayManager::GetWindowManagerAndDisplay(
     const ServerWindow* window) const {
   const ServerWindow* last = window;
diff --git a/components/mus/ws/display_manager.h b/components/mus/ws/display_manager.h
index 686b15a0..0a9ccd5 100644
--- a/components/mus/ws/display_manager.h
+++ b/components/mus/ws/display_manager.h
@@ -58,8 +58,6 @@ class DisplayManager {
   Display* GetDisplayContaining(ServerWindow* window);
   const Display* GetDisplayContaining(const ServerWindow* window) const;
 
-  Display* GetActiveDisplay();
-
   WindowManagerAndDisplayConst GetWindowManagerAndDisplay(
       const ServerWindow* window) const;
   WindowManagerAndDisplay GetWindowManagerAndDisplay(
diff --git a/components/mus/ws/display_manager_delegate.h b/components/mus/ws/display_manager_delegate.h
index 0af7760..0c6ee32 100644
--- a/components/mus/ws/display_manager_delegate.h
+++ b/components/mus/ws/display_manager_delegate.h
@@ -12,7 +12,6 @@ class Display;
 
 class DisplayManagerDelegate {
  public:
-  virtual void OnWillDestroyDisplay(Display* display) = 0;
   virtual void OnFirstDisplayReady() = 0;
   virtual void OnNoMoreDisplays() = 0;
 
diff --git a/components/mus/ws/display_unittest.cc b/components/mus/ws/display_unittest.cc
index 128d3a97..c9f06a1 100644
--- a/components/mus/ws/display_unittest.cc
+++ b/components/mus/ws/display_unittest.cc
@@ -24,6 +24,7 @@
 #include "components/mus/ws/window_tree.h"
 #include "components/mus/ws/window_tree_binding.h"
 #include "testing/gtest/include/gtest/gtest.h"
+#include "ui/events/event.h"
 #include "ui/gfx/geometry/rect.h"
 
 namespace mus {
@@ -153,6 +154,93 @@ TEST_F(DisplayTest, Destruction) {
   EXPECT_TRUE(connection_manager_delegate_.got_on_no_more_displays());
 }
 
+TEST_F(DisplayTest, EventStateResetOnUserSwitch) {
+  connection_manager_delegate_.set_num_displays_to_create(1);
+
+  const UserId kTestId1 = "20";
+  const UserId kTestId2 = "201";
+  WindowManagerFactoryRegistryTestApi(
+      connection_manager_->window_manager_factory_registry())
+      .AddService(kTestId1, &test_window_manager_factory_);
+  WindowManagerFactoryRegistryTestApi(
+      connection_manager_->window_manager_factory_registry())
+      .AddService(kTestId2, &test_window_manager_factory_);
+
+  connection_manager_->user_id_tracker()->SetActiveUserId(kTestId1);
+
+  DisplayManager* display_manager = connection_manager_->display_manager();
+  ASSERT_EQ(1u, display_manager->displays().size());
+  Display* display = *display_manager->displays().begin();
+  WindowManagerState* active_wms = display->GetActiveWindowManagerState();
+  ASSERT_TRUE(active_wms);
+  EXPECT_EQ(kTestId1, active_wms->user_id());
+
+  static_cast<PlatformDisplayDelegate*>(display)->OnEvent(ui::PointerEvent(
+      ui::MouseEvent(ui::ET_MOUSE_PRESSED, gfx::Point(20, 25),
+                     gfx::Point(20, 25), base::TimeDelta(),
+                     ui::EF_LEFT_MOUSE_BUTTON, ui::EF_LEFT_MOUSE_BUTTON)));
+
+  EXPECT_TRUE(EventDispatcherTestApi(active_wms->event_dispatcher())
+                  .AreAnyPointersDown());
+  EXPECT_EQ(gfx::Point(20, 25),
+            active_wms->event_dispatcher()->mouse_pointer_last_location());
+
+  // Switch the user. Should trigger resetting state in old event dispatcher
+  // and update state in new event dispatcher.
+  connection_manager_->user_id_tracker()->SetActiveUserId(kTestId2);
+  EXPECT_NE(active_wms, display->GetActiveWindowManagerState());
+  EXPECT_FALSE(EventDispatcherTestApi(active_wms->event_dispatcher())
+                   .AreAnyPointersDown());
+  active_wms = display->GetActiveWindowManagerState();
+  EXPECT_EQ(kTestId2, active_wms->user_id());
+  EXPECT_EQ(gfx::Point(20, 25),
+            active_wms->event_dispatcher()->mouse_pointer_last_location());
+  EXPECT_FALSE(EventDispatcherTestApi(active_wms->event_dispatcher())
+                   .AreAnyPointersDown());
+}
+
+// Verifies capture fails when wm is inactive and succeeds when wm is active.
+TEST_F(DisplayTest, SetCaptureFromWindowManager) {
+  connection_manager_delegate_.set_num_displays_to_create(1);
+  const UserId kTestId1 = "20";
+  const UserId kTestId2 = "201";
+  WindowManagerFactoryRegistryTestApi(
+      connection_manager_->window_manager_factory_registry())
+      .AddService(kTestId1, &test_window_manager_factory_);
+  WindowManagerFactoryRegistryTestApi(
+      connection_manager_->window_manager_factory_registry())
+      .AddService(kTestId2, &test_window_manager_factory_);
+  connection_manager_->user_id_tracker()->SetActiveUserId(kTestId1);
+  DisplayManager* display_manager = connection_manager_->display_manager();
+  ASSERT_EQ(1u, display_manager->displays().size());
+  Display* display = *display_manager->displays().begin();
+  WindowManagerState* wms_for_id2 =
+      display->GetWindowManagerStateForUser(kTestId2);
+  ASSERT_TRUE(wms_for_id2);
+  EXPECT_FALSE(wms_for_id2->IsActive());
+  ClientWindowId root_client_id;
+
+  // Create a child of the root that we can set capture on.
+  WindowTree* tree = wms_for_id2->tree();
+  const ServerWindow* root = *tree->roots().begin();
+  ASSERT_TRUE(tree->IsWindowKnown(root, &root_client_id));
+  ClientWindowId child_window_id(
+      WindowIdToTransportId(WindowId(tree->id(), 101)));
+  ASSERT_TRUE(tree->NewWindow(child_window_id, ServerWindow::Properties()));
+  ASSERT_TRUE(tree->SetWindowVisibility(child_window_id, true));
+  ASSERT_TRUE(tree->AddWindow(root_client_id, child_window_id));
+
+  WindowTreeTestApi(tree).EnableCapture();
+
+  // SetCapture() should fail for user id2 as it is inactive.
+  EXPECT_FALSE(tree->SetCapture(child_window_id));
+
+  // Make the second user active and verify capture works.
+  connection_manager_->user_id_tracker()->SetActiveUserId(kTestId2);
+  EXPECT_TRUE(wms_for_id2->IsActive());
+  EXPECT_TRUE(tree->SetCapture(child_window_id));
+}
+
 }  // namespace test
 }  // namespace ws
 }  // namespace mus
diff --git a/components/mus/ws/event_dispatcher.cc b/components/mus/ws/event_dispatcher.cc
index 71dea60..c7eca9e 100644
--- a/components/mus/ws/event_dispatcher.cc
+++ b/components/mus/ws/event_dispatcher.cc
@@ -216,6 +216,25 @@ EventDispatcher::~EventDispatcher() {
   pointer_targets_.clear();
 }
 
+void EventDispatcher::Reset() {
+  if (capture_window_) {
+    CancelPointerEventsToTarget(capture_window_);
+    DCHECK(capture_window_ == nullptr);
+  }
+
+  while (!pointer_targets_.empty())
+    StopTrackingPointer(pointer_targets_.begin()->first);
+
+  mouse_button_down_ = false;
+}
+
+void EventDispatcher::SetMousePointerScreenLocation(
+    const gfx::Point& screen_location) {
+  DCHECK(pointer_targets_.empty());
+  mouse_pointer_last_location_ = screen_location;
+  UpdateCursorProviderByLastKnownLocation();
+}
+
 void EventDispatcher::SetCaptureWindow(ServerWindow* window,
                                        bool in_nonclient_area) {
   if (window == capture_window_)
diff --git a/components/mus/ws/event_dispatcher.h b/components/mus/ws/event_dispatcher.h
index 432213e..800ba8a 100644
--- a/components/mus/ws/event_dispatcher.h
+++ b/components/mus/ws/event_dispatcher.h
@@ -30,6 +30,10 @@ class EventDispatcherDelegate;
 class EventMatcher;
 class ServerWindow;
 
+namespace test {
+class EventDispatcherTestApi;
+}
+
 // Handles dispatching events to the right location as well as updating focus.
 class EventDispatcher : public ServerWindowObserver {
  public:
@@ -40,6 +44,15 @@ class EventDispatcher : public ServerWindowObserver {
 
   void set_surface_id(cc::SurfaceId surface_id) { surface_id_ = surface_id; }
 
+  // Cancels capture and stops tracking any pointer events. This does not send
+  // any events to the delegate.
+  void Reset();
+
+  void SetMousePointerScreenLocation(const gfx::Point& screen_location);
+  const gfx::Point& mouse_pointer_last_location() const {
+    return mouse_pointer_last_location_;
+  }
+
   // |capture_window_| will receive all input. See window_tree.mojom for
   // details.
   ServerWindow* capture_window() { return capture_window_; }
@@ -67,7 +80,7 @@ class EventDispatcher : public ServerWindowObserver {
   void ProcessEvent(const ui::Event& event);
 
  private:
-  friend class EventDispatcherTest;
+  friend class test::EventDispatcherTestApi;
 
   // Keeps track of state associated with an active pointer.
   struct PointerTarget {
diff --git a/components/mus/ws/event_dispatcher_unittest.cc b/components/mus/ws/event_dispatcher_unittest.cc
index 5dac279..ea34f1c 100644
--- a/components/mus/ws/event_dispatcher_unittest.cc
+++ b/components/mus/ws/event_dispatcher_unittest.cc
@@ -15,11 +15,13 @@
 #include "components/mus/ws/server_window.h"
 #include "components/mus/ws/server_window_surface_manager_test_api.h"
 #include "components/mus/ws/test_server_window_delegate.h"
+#include "components/mus/ws/test_utils.h"
 #include "testing/gtest/include/gtest/gtest.h"
 #include "ui/events/event.h"
 
 namespace mus {
 namespace ws {
+namespace test {
 namespace {
 
 // Identifies a generated event.
@@ -202,7 +204,7 @@ class EventDispatcherTest : public testing::Test {
 };
 
 bool EventDispatcherTest::AreAnyPointersDown() const {
-  return event_dispatcher_->AreAnyPointersDown();
+  return EventDispatcherTestApi(event_dispatcher_.get()).AreAnyPointersDown();
 }
 
 void EventDispatcherTest::ClearSetup() {
@@ -221,20 +223,18 @@ ServerWindow* EventDispatcherTest::CreateChildWindow(const WindowId& id) {
 }
 
 bool EventDispatcherTest::IsMouseButtonDown() const {
-  return event_dispatcher_->mouse_button_down_;
+  return EventDispatcherTestApi(event_dispatcher_.get()).is_mouse_button_down();
 }
 
 bool EventDispatcherTest::IsWindowPointerTarget(ServerWindow* window) const {
-  return event_dispatcher_->IsObservingWindow(window);
+  return EventDispatcherTestApi(event_dispatcher_.get())
+      .IsObservingWindow(window);
 }
 
 int EventDispatcherTest::NumberPointerTargetsForWindow(
     ServerWindow* window) const {
-  int count = 0;
-  for (const auto& pair : event_dispatcher_->pointer_targets_)
-    if (pair.second.window == window)
-      count++;
-  return count;
+  return EventDispatcherTestApi(event_dispatcher_.get())
+      .NumberPointerTargetsForWindow(window);
 }
 
 void EventDispatcherTest::SetUp() {
@@ -1088,5 +1088,43 @@ TEST_F(EventDispatcherTest, ProcessPointerEvents) {
   }
 }
 
+TEST_F(EventDispatcherTest, ResetClearsPointerDown) {
+  scoped_ptr<ServerWindow> child(CreateChildWindow(WindowId(1, 3)));
+
+  root_window()->SetBounds(gfx::Rect(0, 0, 100, 100));
+  child->SetBounds(gfx::Rect(10, 10, 20, 20));
+
+  // Send event that is over child.
+  const ui::PointerEvent ui_event(ui::MouseEvent(
+      ui::ET_MOUSE_PRESSED, gfx::Point(20, 25), gfx::Point(20, 25),
+      base::TimeDelta(), ui::EF_LEFT_MOUSE_BUTTON, ui::EF_LEFT_MOUSE_BUTTON));
+  event_dispatcher()->ProcessEvent(ui_event);
+
+  scoped_ptr<DispatchedEventDetails> details =
+      test_event_dispatcher_delegate()->GetAndAdvanceDispatchedEventDetails();
+  ASSERT_TRUE(details);
+  ASSERT_EQ(child.get(), details->window);
+
+  EXPECT_TRUE(AreAnyPointersDown());
+
+  event_dispatcher()->Reset();
+  EXPECT_FALSE(test_event_dispatcher_delegate()->has_queued_events());
+  EXPECT_FALSE(AreAnyPointersDown());
+}
+
+TEST_F(EventDispatcherTest, ResetClearsCapture) {
+  ServerWindow* root = root_window();
+  root->SetBounds(gfx::Rect(0, 0, 100, 100));
+
+  root->SetClientArea(gfx::Insets(5, 5, 5, 5), std::vector<gfx::Rect>());
+  EventDispatcher* dispatcher = event_dispatcher();
+  dispatcher->SetCaptureWindow(root, true);
+
+  event_dispatcher()->Reset();
+  EXPECT_FALSE(test_event_dispatcher_delegate()->has_queued_events());
+  EXPECT_EQ(nullptr, event_dispatcher()->capture_window());
+}
+
+}  // namespace test
 }  // namespace ws
 }  // namespace mus
diff --git a/components/mus/ws/test_utils.cc b/components/mus/ws/test_utils.cc
index a02b675..1f3b0dc 100644
--- a/components/mus/ws/test_utils.cc
+++ b/components/mus/ws/test_utils.cc
@@ -104,6 +104,17 @@ WindowTreeTestApi::~WindowTreeTestApi() {}
 DisplayTestApi::DisplayTestApi(Display* display) : display_(display) {}
 DisplayTestApi::~DisplayTestApi() {}
 
+// EventDispatcherTestApi  ----------------------------------------------------
+
+int EventDispatcherTestApi::NumberPointerTargetsForWindow(
+    ServerWindow* window) {
+  int count = 0;
+  for (const auto& pair : ed_->pointer_targets_)
+    if (pair.second.window == window)
+      count++;
+  return count;
+}
+
 // TestWindowTreeClient -------------------------------------------------------
 
 TestWindowTreeClient::TestWindowTreeClient()
diff --git a/components/mus/ws/test_utils.h b/components/mus/ws/test_utils.h
index 642fd1b..d431e24 100644
--- a/components/mus/ws/test_utils.h
+++ b/components/mus/ws/test_utils.h
@@ -11,11 +11,13 @@
 #include "components/mus/public/interfaces/window_tree.mojom.h"
 #include "components/mus/ws/connection_manager_delegate.h"
 #include "components/mus/ws/display.h"
+#include "components/mus/ws/event_dispatcher.h"
 #include "components/mus/ws/platform_display.h"
 #include "components/mus/ws/platform_display_factory.h"
 #include "components/mus/ws/test_change_tracker.h"
 #include "components/mus/ws/user_display_manager.h"
 #include "components/mus/ws/window_manager_factory_registry.h"
+#include "components/mus/ws/window_manager_state.h"
 #include "components/mus/ws/window_tree.h"
 #include "components/mus/ws/window_tree_binding.h"
 
@@ -68,6 +70,8 @@ class WindowTreeTestApi {
     tree_->window_manager_internal_ = wm_internal;
   }
 
+  void EnableCapture() { tree_->event_ack_id_ = 1u; }
+
  private:
   WindowTree* tree_;
 
@@ -83,14 +87,47 @@ class DisplayTestApi {
 
   void OnEvent(const ui::Event& event) { display_->OnEvent(event); }
 
+ private:
+  Display* display_;
+
+  DISALLOW_COPY_AND_ASSIGN(DisplayTestApi);
+};
+
+// -----------------------------------------------------------------------------
+
+class EventDispatcherTestApi {
+ public:
+  explicit EventDispatcherTestApi(EventDispatcher* ed) : ed_(ed) {}
+  ~EventDispatcherTestApi() {}
+
+  bool AreAnyPointersDown() const { return ed_->AreAnyPointersDown(); }
+  bool is_mouse_button_down() const { return ed_->mouse_button_down_; }
+  bool IsObservingWindow(ServerWindow* window) {
+    return ed_->IsObservingWindow(window);
+  }
+  int NumberPointerTargetsForWindow(ServerWindow* window);
+
+ private:
+  EventDispatcher* ed_;
+
+  DISALLOW_COPY_AND_ASSIGN(EventDispatcherTestApi);
+};
+
+// -----------------------------------------------------------------------------
+
+class WindowManagerStateTestApi {
+ public:
+  explicit WindowManagerStateTestApi(WindowManagerState* wms) : wms_(wms) {}
+  ~WindowManagerStateTestApi() {}
+
   mojom::WindowTree* tree_awaiting_input_ack() {
-    return display_->tree_awaiting_input_ack_;
+    return wms_->tree_awaiting_input_ack_;
   }
 
  private:
-  Display* display_;
+  WindowManagerState* wms_;
 
-  DISALLOW_COPY_AND_ASSIGN(DisplayTestApi);
+  DISALLOW_COPY_AND_ASSIGN(WindowManagerStateTestApi);
 };
 
 // -----------------------------------------------------------------------------
diff --git a/components/mus/ws/user_id_tracker.cc b/components/mus/ws/user_id_tracker.cc
index 453fde6..3a54f64 100644
--- a/components/mus/ws/user_id_tracker.cc
+++ b/components/mus/ws/user_id_tracker.cc
@@ -24,9 +24,10 @@ void UserIdTracker::SetActiveUserId(const UserId& id) {
   if (id == active_id_)
     return;
 
+  const UserId previously_active_id = active_id_;
   active_id_ = id;
   FOR_EACH_OBSERVER(UserIdTrackerObserver, observers_,
-                    OnActiveUserIdChanged(id));
+                    OnActiveUserIdChanged(previously_active_id, id));
 }
 
 void UserIdTracker::AddUserId(const UserId& id) {
diff --git a/components/mus/ws/user_id_tracker.h b/components/mus/ws/user_id_tracker.h
index 45d577d..e1f5cf3 100644
--- a/components/mus/ws/user_id_tracker.h
+++ b/components/mus/ws/user_id_tracker.h
@@ -26,6 +26,8 @@ class UserIdTracker {
 
   bool IsValidUserId(const UserId& id) const;
 
+  const UserId& active_id() const { return active_id_; }
+
   void SetActiveUserId(const UserId& id);
   void AddUserId(const UserId& id);
   void RemoveUserId(const UserId& id);
diff --git a/components/mus/ws/user_id_tracker_observer.h b/components/mus/ws/user_id_tracker_observer.h
index 840c410..4632dea 100644
--- a/components/mus/ws/user_id_tracker_observer.h
+++ b/components/mus/ws/user_id_tracker_observer.h
@@ -14,7 +14,8 @@ namespace ws {
 
 class UserIdTrackerObserver {
  public:
-  virtual void OnActiveUserIdChanged(const UserId& id) = 0;
+  virtual void OnActiveUserIdChanged(const UserId& previously_active_id,
+                                     const UserId& active_id) = 0;
   virtual void OnUserIdAdded(const UserId& id) = 0;
   virtual void OnUserIdRemoved(const UserId& id) = 0;
 
diff --git a/components/mus/ws/window_manager_factory_registry.cc b/components/mus/ws/window_manager_factory_registry.cc
index d7bb576..8c76131 100644
--- a/components/mus/ws/window_manager_factory_registry.cc
+++ b/components/mus/ws/window_manager_factory_registry.cc
@@ -93,7 +93,9 @@ void WindowManagerFactoryRegistry::OnWindowManagerFactorySet(
     connection_manager_->OnFirstWindowManagerFactorySet();
 }
 
-void WindowManagerFactoryRegistry::OnActiveUserIdChanged(const UserId& id) {}
+void WindowManagerFactoryRegistry::OnActiveUserIdChanged(
+    const UserId& previously_active_id,
+    const UserId& active_id) {}
 
 void WindowManagerFactoryRegistry::OnUserIdAdded(const UserId& id) {}
 
diff --git a/components/mus/ws/window_manager_factory_registry.h b/components/mus/ws/window_manager_factory_registry.h
index efd29a4..ece6d39 100644
--- a/components/mus/ws/window_manager_factory_registry.h
+++ b/components/mus/ws/window_manager_factory_registry.h
@@ -56,7 +56,8 @@ class WindowManagerFactoryRegistry : public UserIdTrackerObserver {
   void OnWindowManagerFactorySet(WindowManagerFactoryService* service);
 
   // UserIdTrackerObserver:
-  void OnActiveUserIdChanged(const UserId& id) override;
+  void OnActiveUserIdChanged(const UserId& previously_active_id,
+                             const UserId& active_id) override;
   void OnUserIdAdded(const UserId& id) override;
   void OnUserIdRemoved(const UserId& id) override;
 
diff --git a/components/mus/ws/window_manager_state.cc b/components/mus/ws/window_manager_state.cc
index f7ccc63..2a6d983 100644
--- a/components/mus/ws/window_manager_state.cc
+++ b/components/mus/ws/window_manager_state.cc
@@ -6,27 +6,142 @@
 
 #include "components/mus/ws/connection_manager.h"
 #include "components/mus/ws/display_manager.h"
+#include "components/mus/ws/platform_display.h"
 #include "components/mus/ws/server_window.h"
 #include "components/mus/ws/user_display_manager.h"
+#include "components/mus/ws/user_id_tracker.h"
+#include "components/mus/ws/window_tree.h"
+#include "mojo/converters/input_events/input_events_type_converters.h"
 #include "mojo/shell/public/interfaces/connector.mojom.h"
+#include "ui/events/event.h"
 
 namespace mus {
 namespace ws {
+namespace {
 
-WindowManagerState::WindowManagerState(Display* display)
-    : WindowManagerState(display, false, mojo::shell::mojom::kRootUserID) {}
+base::TimeDelta GetDefaultAckTimerDelay() {
+#if defined(NDEBUG)
+  return base::TimeDelta::FromMilliseconds(100);
+#else
+  return base::TimeDelta::FromMilliseconds(1000);
+#endif
+}
+
+bool EventsCanBeCoalesced(const ui::Event& one, const ui::Event& two) {
+  if (one.type() != two.type() || one.flags() != two.flags())
+    return false;
+
+  // TODO(sad): wheel events can also be merged.
+  if (one.type() != ui::ET_POINTER_MOVED)
+    return false;
+
+  return one.AsPointerEvent()->pointer_id() ==
+         two.AsPointerEvent()->pointer_id();
+}
+
+scoped_ptr<ui::Event> CoalesceEvents(scoped_ptr<ui::Event> first,
+                                     scoped_ptr<ui::Event> second) {
+  DCHECK(first->type() == ui::ET_POINTER_MOVED)
+      << " Non-move events cannot be merged yet.";
+  // For mouse moves, the new event just replaces the old event.
+  return second;
+}
+
+}  // namespace
+
+class WindowManagerState::ProcessedEventTarget {
+ public:
+  ProcessedEventTarget(ServerWindow* window, bool in_nonclient_area)
+      : in_nonclient_area_(in_nonclient_area) {
+    tracker_.Add(window);
+  }
+
+  ~ProcessedEventTarget() {}
+
+  // Return true if the event is still valid. The event becomes invalid if
+  // the window is destroyed while waiting to dispatch.
+  bool IsValid() const { return !tracker_.windows().empty(); }
 
-WindowManagerState::WindowManagerState(Display* display, const UserId& user_id)
-    : WindowManagerState(display, true, user_id) {}
+  ServerWindow* window() {
+    DCHECK(IsValid());
+    return tracker_.windows().front();
+  }
+
+  bool in_nonclient_area() const { return in_nonclient_area_; }
+
+ private:
+  ServerWindowTracker tracker_;
+  const bool in_nonclient_area_;
+
+  DISALLOW_COPY_AND_ASSIGN(ProcessedEventTarget);
+};
+
+WindowManagerState::QueuedEvent::QueuedEvent() {}
+WindowManagerState::QueuedEvent::~QueuedEvent() {}
+
+WindowManagerState::WindowManagerState(Display* display,
+                                       PlatformDisplay* platform_display,
+                                       cc::SurfaceId surface_id)
+    : WindowManagerState(display,
+                         platform_display,
+                         surface_id,
+                         false,
+                         mojo::shell::mojom::kRootUserID) {}
+
+WindowManagerState::WindowManagerState(Display* display,
+                                       PlatformDisplay* platform_display,
+                                       cc::SurfaceId surface_id,
+                                       const UserId& user_id)
+    : WindowManagerState(display, platform_display, surface_id, true, user_id) {
+}
 
 WindowManagerState::~WindowManagerState() {}
 
+void WindowManagerState::SetFrameDecorationValues(
+    mojom::FrameDecorationValuesPtr values) {
+  got_frame_decoration_values_ = true;
+  frame_decoration_values_ = values.Clone();
+  display_->display_manager()
+      ->GetUserDisplayManager(user_id_)
+      ->OnFrameDecorationValuesChanged(this);
+}
+
+void WindowManagerState::SetCapture(ServerWindow* window,
+                                    bool in_nonclient_area) {
+  // TODO(sky): capture should be a singleton. Need to route to
+  // ConnectionManager so that all other EventDispatchers are updated.
+  DCHECK(IsActive());
+  if (capture_window() == window)
+    return;
+  DCHECK(!window || root_->Contains(window));
+  event_dispatcher_.SetCaptureWindow(window, in_nonclient_area);
+}
+
+mojom::DisplayPtr WindowManagerState::ToMojomDisplay() const {
+  mojom::DisplayPtr display_ptr = display_->ToMojomDisplay();
+  // TODO(sky): set work area.
+  display_ptr->work_area = display_ptr->bounds.Clone();
+  display_ptr->frame_decoration_values = frame_decoration_values_.Clone();
+  return display_ptr;
+}
+
+void WindowManagerState::OnWillDestroyTree(WindowTree* tree) {
+  if (tree_awaiting_input_ack_ != tree)
+    return;
+  // The WindowTree is dying. So it's not going to ack the event.
+  OnEventAck(tree_awaiting_input_ack_);
+}
+
 WindowManagerState::WindowManagerState(Display* display,
+                                       PlatformDisplay* platform_display,
+                                       cc::SurfaceId surface_id,
                                        bool is_user_id_valid,
                                        const UserId& user_id)
     : display_(display),
+      platform_display_(platform_display),
       is_user_id_valid_(is_user_id_valid),
-      user_id_(user_id) {
+      user_id_(user_id),
+      event_dispatcher_(this) {
   frame_decoration_values_ = mojom::FrameDecorationValues::New();
   frame_decoration_values_->normal_client_area_insets = mojo::Insets::New();
   frame_decoration_values_->maximized_client_area_insets = mojo::Insets::New();
@@ -42,23 +157,162 @@ WindowManagerState::WindowManagerState(Display* display,
   root_->SetBounds(gfx::Rect(display->root_window()->bounds().size()));
   root_->SetVisible(true);
   display->root_window()->Add(root_.get());
+
+  event_dispatcher_.set_root(root_.get());
+  event_dispatcher_.set_surface_id(surface_id);
 }
 
-void WindowManagerState::SetFrameDecorationValues(
-    mojom::FrameDecorationValuesPtr values) {
-  got_frame_decoration_values_ = true;
-  frame_decoration_values_ = values.Clone();
-  display_->display_manager()
-      ->GetUserDisplayManager(user_id_)
-      ->OnFrameDecorationValuesChanged(this);
+bool WindowManagerState::IsActive() const {
+  return display()->GetActiveWindowManagerState() == this;
 }
 
-mojom::DisplayPtr WindowManagerState::ToMojomDisplay() const {
-  mojom::DisplayPtr display_ptr = display_->ToMojomDisplay();
-  // TODO(sky): set work area.
-  display_ptr->work_area = display_ptr->bounds.Clone();
-  display_ptr->frame_decoration_values = frame_decoration_values_.Clone();
-  return display_ptr;
+void WindowManagerState::ProcessEvent(const ui::Event& event) {
+  mojom::EventPtr mojo_event(mojom::Event::From(event));
+  // If this is still waiting for an ack from a previously sent event, then
+  // queue up the event to be dispatched once the ack is received.
+  if (event_ack_timer_.IsRunning()) {
+    if (!event_queue_.empty() && !event_queue_.back()->processed_target &&
+        EventsCanBeCoalesced(*event_queue_.back()->event, event)) {
+      event_queue_.back()->event = CoalesceEvents(
+          std::move(event_queue_.back()->event), ui::Event::Clone(event));
+      return;
+    }
+    QueueEvent(event, nullptr);
+    return;
+  }
+  event_dispatcher_.ProcessEvent(event);
+}
+
+void WindowManagerState::OnEventAck(mojom::WindowTree* tree) {
+  if (tree_awaiting_input_ack_ != tree) {
+    // TODO(sad): The ack must have arrived after the timeout. We should do
+    // something here, and in OnEventAckTimeout().
+    return;
+  }
+  tree_awaiting_input_ack_ = nullptr;
+  event_ack_timer_.Stop();
+  ProcessNextEventFromQueue();
+}
+
+ConnectionManager* WindowManagerState::connection_manager() {
+  return display_->connection_manager();
+}
+
+void WindowManagerState::OnEventAckTimeout() {
+  // TODO(sad): Figure out what we should do.
+  NOTIMPLEMENTED() << "Event ACK timed out.";
+  OnEventAck(tree_awaiting_input_ack_);
+}
+
+void WindowManagerState::QueueEvent(
+    const ui::Event& event,
+    scoped_ptr<ProcessedEventTarget> processed_event_target) {
+  scoped_ptr<QueuedEvent> queued_event(new QueuedEvent);
+  queued_event->event = ui::Event::Clone(event);
+  queued_event->processed_target = std::move(processed_event_target);
+  event_queue_.push(std::move(queued_event));
+}
+
+void WindowManagerState::ProcessNextEventFromQueue() {
+  // Loop through |event_queue_| stopping after dispatching the first valid
+  // event.
+  while (!event_queue_.empty()) {
+    scoped_ptr<QueuedEvent> queued_event = std::move(event_queue_.front());
+    event_queue_.pop();
+    if (!queued_event->processed_target) {
+      event_dispatcher_.ProcessEvent(*queued_event->event);
+      return;
+    }
+    if (queued_event->processed_target->IsValid()) {
+      DispatchInputEventToWindowImpl(
+          queued_event->processed_target->window(),
+          queued_event->processed_target->in_nonclient_area(),
+          *queued_event->event);
+      return;
+    }
+  }
+}
+
+void WindowManagerState::DispatchInputEventToWindowImpl(
+    ServerWindow* target,
+    bool in_nonclient_area,
+    const ui::Event& event) {
+  if (target == root_->parent())
+    target = root_.get();
+
+  if (event.IsMousePointerEvent()) {
+    DCHECK(event_dispatcher_.mouse_cursor_source_window());
+    display_->UpdateNativeCursor(
+        event_dispatcher_.mouse_cursor_source_window()->cursor());
+  }
+
+  // If the event is in the non-client area the event goes to the owner of
+  // the window. Otherwise if the window is an embed root, forward to the
+  // embedded window.
+  WindowTree* tree =
+      in_nonclient_area
+          ? connection_manager()->GetTreeWithId(target->id().connection_id)
+          : connection_manager()->GetTreeWithRoot(target);
+  if (!tree) {
+    DCHECK(!in_nonclient_area);
+    tree = connection_manager()->GetTreeWithId(target->id().connection_id);
+  }
+
+  // TOOD(sad): Adjust this delay, possibly make this dynamic.
+  const base::TimeDelta max_delay = base::debug::BeingDebugged()
+                                        ? base::TimeDelta::FromDays(1)
+                                        : GetDefaultAckTimerDelay();
+  event_ack_timer_.Start(FROM_HERE, max_delay, this,
+                         &WindowManagerState::OnEventAckTimeout);
+
+  tree_awaiting_input_ack_ = tree;
+  tree->DispatchInputEvent(target, mojom::Event::From(event));
+}
+
+void WindowManagerState::OnAccelerator(uint32_t accelerator_id,
+                                       const ui::Event& event) {
+  DCHECK(IsActive());
+  tree_->OnAccelerator(accelerator_id, mojom::Event::From(event));
+}
+
+void WindowManagerState::SetFocusedWindowFromEventDispatcher(
+    ServerWindow* new_focused_window) {
+  DCHECK(IsActive());
+  display_->SetFocusedWindow(new_focused_window);
+}
+
+ServerWindow* WindowManagerState::GetFocusedWindowForEventDispatcher() {
+  return display()->GetFocusedWindow();
+}
+
+void WindowManagerState::SetNativeCapture() {
+  DCHECK(IsActive());
+  platform_display_->SetCapture();
+}
+
+void WindowManagerState::ReleaseNativeCapture() {
+  platform_display_->ReleaseCapture();
+}
+
+void WindowManagerState::OnServerWindowCaptureLost(ServerWindow* window) {
+  DCHECK(window);
+  display_->connection_manager()->ProcessLostCapture(window);
+}
+
+void WindowManagerState::DispatchInputEventToWindow(ServerWindow* target,
+                                                    bool in_nonclient_area,
+                                                    const ui::Event& event) {
+  DCHECK(IsActive());
+  // TODO(sky): this needs to see if another wms has capture and if so forward
+  // to it.
+  if (event_ack_timer_.IsRunning()) {
+    scoped_ptr<ProcessedEventTarget> processed_event_target(
+        new ProcessedEventTarget(target, in_nonclient_area));
+    QueueEvent(event, std::move(processed_event_target));
+    return;
+  }
+
+  DispatchInputEventToWindowImpl(target, in_nonclient_area, event);
 }
 
 }  // namespace ws
diff --git a/components/mus/ws/window_manager_state.h b/components/mus/ws/window_manager_state.h
index c0fe99a..af6f389 100644
--- a/components/mus/ws/window_manager_state.h
+++ b/components/mus/ws/window_manager_state.h
@@ -8,27 +8,47 @@
 #include <stdint.h>
 
 #include "base/memory/scoped_ptr.h"
+#include "base/timer/timer.h"
 #include "components/mus/public/interfaces/display.mojom.h"
+#include "components/mus/ws/connection_manager.h"
+#include "components/mus/ws/event_dispatcher.h"
+#include "components/mus/ws/event_dispatcher_delegate.h"
 #include "components/mus/ws/user_id.h"
 
+namespace cc {
+struct SurfaceId;
+}
+
 namespace mus {
 namespace ws {
 
 class Display;
+class PlatformDisplay;
 class ServerWindow;
 class WindowTree;
 
+namespace test {
+class WindowManagerStateTestApi;
+}
+
 // Manages the state associated with a connection to a WindowManager for
 // a specific user.
-class WindowManagerState {
+class WindowManagerState : public EventDispatcherDelegate {
  public:
   // Creates a WindowManagerState that can host content from any user.
-  explicit WindowManagerState(Display* display);
-  WindowManagerState(Display* display, const UserId& user_id);
-  ~WindowManagerState();
+  WindowManagerState(Display* display,
+                     PlatformDisplay* platform_display,
+                     cc::SurfaceId surface_id);
+  WindowManagerState(Display* display,
+                     PlatformDisplay* platform_display,
+                     cc::SurfaceId surface_id,
+                     const UserId& user_id);
+  ~WindowManagerState() override;
 
   bool is_user_id_valid() const { return is_user_id_valid_; }
 
+  const UserId& user_id() const { return user_id_; }
+
   ServerWindow* root() { return root_.get(); }
   const ServerWindow* root() const { return root_.get(); }
 
@@ -46,22 +66,97 @@ class WindowManagerState {
     return got_frame_decoration_values_;
   }
 
+  void SetCapture(ServerWindow* window, bool in_nonclient_area);
+  ServerWindow* capture_window() { return event_dispatcher_.capture_window(); }
+  const ServerWindow* capture_window() const {
+    return event_dispatcher_.capture_window();
+  }
+
+  // Returns true if this is the WindowManager of the active user.
+  bool IsActive() const;
+
+  // TODO(sky): EventDispatcher is really an implementation detail and should
+  // not be exposed.
+  EventDispatcher* event_dispatcher() { return &event_dispatcher_; }
+
+  // Processes an event from PlatformDisplay.
+  void ProcessEvent(const ui::Event& event);
+
+  // Called when the ack from an event dispatched to a WindowTree is received.
+  // TODO(sky): make this private and use a callback.
+  void OnEventAck(mojom::WindowTree* tree);
+
   // Returns a mojom::Display for the specified display. WindowManager specific
   // values are not set.
   mojom::DisplayPtr ToMojomDisplay() const;
 
+  void OnWillDestroyTree(WindowTree* tree);
+
  private:
+  class ProcessedEventTarget;
   friend class Display;
-
-  WindowManagerState(Display* display, bool is_user_id_valid,
+  friend class test::WindowManagerStateTestApi;
+
+  // There are two types of events that may be queued, both occur only when
+  // waiting for an ack from a client.
+  // . We get an event from the PlatformDisplay. This results in |event| being
+  //   set, but |processed_target| is null.
+  // . We get an event from the EventDispatcher. In this case both |event| and
+  //   |processed_target| are valid.
+  // The second case happens if EventDispatcher generates more than one event
+  // at a time.
+  struct QueuedEvent {
+    QueuedEvent();
+    ~QueuedEvent();
+
+    scoped_ptr<ui::Event> event;
+    scoped_ptr<ProcessedEventTarget> processed_target;
+  };
+
+  WindowManagerState(Display* display,
+                     PlatformDisplay* platform_display,
+                     cc::SurfaceId surface_id,
+                     bool is_user_id_valid,
                      const UserId& user_id);
 
+  ConnectionManager* connection_manager();
+
+  void OnEventAckTimeout();
+
+  // Schedules an event to be processed later.
+  void QueueEvent(const ui::Event& event,
+                  scoped_ptr<ProcessedEventTarget> processed_event_target);
+
+  // Processes the next valid event in |event_queue_|. If the event has already
+  // been processed it is dispatched, otherwise the event is passed to the
+  // EventDispatcher for processing.
+  void ProcessNextEventFromQueue();
+
+  // Dispatches the event to the appropriate client and starts the ack timer.
+  void DispatchInputEventToWindowImpl(ServerWindow* target,
+                                      bool in_nonclient_area,
+                                      const ui::Event& event);
+
+  // EventDispatcherDelegate:
+  void OnAccelerator(uint32_t accelerator_id, const ui::Event& event) override;
+  void SetFocusedWindowFromEventDispatcher(ServerWindow* window) override;
+  ServerWindow* GetFocusedWindowForEventDispatcher() override;
+  void SetNativeCapture() override;
+  void ReleaseNativeCapture() override;
+  void OnServerWindowCaptureLost(ServerWindow* window) override;
+  void DispatchInputEventToWindow(ServerWindow* target,
+                                  bool in_nonclient_area,
+                                  const ui::Event& event) override;
+
   Display* display_;
+  PlatformDisplay* platform_display_;
   // If this was created implicitly by a call
   // WindowTreeHostFactory::CreateWindowTreeHost(), then |is_user_id_valid_|
   // is false.
   const bool is_user_id_valid_;
   const UserId user_id_;
+  // Root ServerWindow of this WindowManagerState. |root_| has a parent, the
+  // root ServerWindow of the Display.
   scoped_ptr<ServerWindow> root_;
   WindowTree* tree_ = nullptr;
 
@@ -69,6 +164,12 @@ class WindowManagerState {
   bool got_frame_decoration_values_ = false;
   mojom::FrameDecorationValuesPtr frame_decoration_values_;
 
+  mojom::WindowTree* tree_awaiting_input_ack_ = nullptr;
+  std::queue<scoped_ptr<QueuedEvent>> event_queue_;
+  base::OneShotTimer event_ack_timer_;
+
+  EventDispatcher event_dispatcher_;
+
   DISALLOW_COPY_AND_ASSIGN(WindowManagerState);
 };
 
diff --git a/components/mus/ws/window_tree.cc b/components/mus/ws/window_tree.cc
index 7782e8b..90e5c57 100644
--- a/components/mus/ws/window_tree.cc
+++ b/components/mus/ws/window_tree.cc
@@ -72,7 +72,6 @@ WindowTree::WindowTree(ConnectionManager* connection_manager,
       id_(connection_manager_->GetAndAdvanceNextConnectionId()),
       next_window_id_(1),
       event_ack_id_(0),
-      event_source_display_(nullptr),
       is_embed_root_(false),
       window_manager_internal_(nullptr) {
   if (root)
@@ -158,7 +157,19 @@ const Display* WindowTree::GetDisplay(const ServerWindow* window) const {
   return window ? display_manager()->GetDisplayContaining(window) : nullptr;
 }
 
+const WindowManagerState* WindowTree::GetWindowManagerState(
+    const ServerWindow* window) const {
+  return window
+             ? display_manager()
+                   ->GetWindowManagerAndDisplay(window)
+                   .window_manager_state
+             : nullptr;
+}
+
 void WindowTree::OnWindowDestroyingTreeImpl(WindowTree* tree) {
+  if (event_source_wms_ && event_source_wms_->tree() == tree)
+    event_source_wms_ = nullptr;
+
   // Notify our client if |tree| was embedded in any of our views.
   for (const auto* tree_root : tree->roots_) {
     const bool owns_tree_root = tree_root->id().connection_id == id_;
@@ -169,11 +180,6 @@ void WindowTree::OnWindowDestroyingTreeImpl(WindowTree* tree) {
   }
 }
 
-void WindowTree::OnWillDestroyDisplay(Display* display) {
-  if (event_source_display_ == display)
-    event_source_display_ = nullptr;
-}
-
 void WindowTree::NotifyChangeCompleted(
     uint32_t change_id,
     mojom::WindowManagerErrorCode error_code) {
@@ -181,6 +187,21 @@ void WindowTree::NotifyChangeCompleted(
       change_id, error_code == mojom::WindowManagerErrorCode::SUCCESS);
 }
 
+bool WindowTree::SetCapture(const ClientWindowId& client_window_id) {
+  ServerWindow* window = GetWindowByClientId(client_window_id);
+  WindowManagerState* wms = GetWindowManagerState(window);
+  ServerWindow* current_capture_window = wms ? wms->capture_window() : nullptr;
+  if (window && wms && wms->IsActive() &&
+      access_policy_->CanSetCapture(window) &&
+      (!current_capture_window ||
+       access_policy_->CanSetCapture(current_capture_window)) &&
+      event_ack_id_) {
+    wms->SetCapture(window, !HasRoot(window));
+    return true;
+  }
+  return false;
+}
+
 bool WindowTree::NewWindow(
     const ClientWindowId& client_window_id,
     const std::map<std::string, std::vector<uint8_t>>& properties) {
@@ -877,9 +898,9 @@ void WindowTree::DispatchInputEventImpl(ServerWindow* target,
   // the event pointer.
   event_ack_id_ =
       0x1000000 | (reinterpret_cast<uintptr_t>(event.get()) & 0xffffff);
-  event_source_display_ = GetDisplay(target);
+  event_source_wms_ = GetWindowManagerState(target);
   // Should only get events from windows attached to a host.
-  DCHECK(event_source_display_);
+  DCHECK(event_source_wms_);
   client()->OnWindowInputEvent(
       event_ack_id_, ClientWindowIdForWindow(target).id, std::move(event));
 }
@@ -903,8 +924,11 @@ void WindowTree::NewTopLevelWindow(
     mojo::Map<mojo::String, mojo::Array<uint8_t>> transport_properties) {
   DCHECK(!waiting_for_top_level_window_info_);
   const ClientWindowId client_window_id(transport_window_id);
-  Display* display = display_manager()->GetActiveDisplay();
-  // TODO(sky): need a way for client to provide context.
+  // TODO(sky): need a way for client to provide context to figure out display.
+  // TODO(sky): get WMS for user_id for this connection.
+  Display* display = display_manager()->displays().empty()
+                         ? nullptr
+                         : *(display_manager()->displays().begin());
   WindowManagerState* wms =
       display ? display->GetFirstWindowManagerState() : nullptr;
   if (!wms || wms->tree() == this || !IsValidIdForNewWindow(client_window_id)) {
@@ -1011,33 +1035,20 @@ void WindowTree::GetWindowTree(
 }
 
 void WindowTree::SetCapture(uint32_t change_id, Id window_id) {
-  ServerWindow* window = GetWindowByClientId(ClientWindowId(window_id));
-  Display* display = GetDisplay(window);
-  ServerWindow* current_capture_window =
-      display ? display->GetCaptureWindow() : nullptr;
-  bool success = window && access_policy_->CanSetCapture(window) && display &&
-                 (!current_capture_window ||
-                  access_policy_->CanSetCapture(current_capture_window)) &&
-                 event_ack_id_;
-  if (success) {
-    Operation op(this, connection_manager_, OperationType::SET_CAPTURE);
-    display->SetCapture(window, !HasRoot(window));
-  }
-  client()->OnChangeCompleted(change_id, success);
+  client()->OnChangeCompleted(change_id, SetCapture(ClientWindowId(window_id)));
 }
 
 void WindowTree::ReleaseCapture(uint32_t change_id, Id window_id) {
   ServerWindow* window = GetWindowByClientId(ClientWindowId(window_id));
-  Display* display = GetDisplay(window);
-  ServerWindow* current_capture_window =
-      display ? display->GetCaptureWindow() : nullptr;
-  bool success = window && display &&
+  WindowManagerState* wms = GetWindowManagerState(window);
+  ServerWindow* current_capture_window = wms ? wms->capture_window() : nullptr;
+  bool success = window && wms && wms->IsActive() &&
                  (!current_capture_window ||
                   access_policy_->CanSetCapture(current_capture_window)) &&
                  window == current_capture_window;
   if (success) {
     Operation op(this, connection_manager_, OperationType::RELEASE_CAPTURE);
-    display->SetCapture(nullptr, false);
+    wms->SetCapture(nullptr, false);
   }
   client()->OnChangeCompleted(change_id, success);
 }
@@ -1152,10 +1163,10 @@ void WindowTree::OnWindowInputEventAck(uint32_t event_id) {
   }
   event_ack_id_ = 0;
 
-  Display* event_source_display = event_source_display_;
-  event_source_display_ = nullptr;
-  if (event_source_display)
-    event_source_display->OnEventAck(this);
+  WindowManagerState* event_source_wms = event_source_wms_;
+  event_source_wms_ = nullptr;
+  if (event_source_wms)
+    event_source_wms->OnEventAck(this);
 
   if (!event_queue_.empty()) {
     DCHECK(!event_ack_id_);
@@ -1253,18 +1264,15 @@ void WindowTree::GetWindowManagerClient(
 void WindowTree::AddAccelerator(uint32_t id,
                                 mojom::EventMatcherPtr event_matcher,
                                 const AddAcceleratorCallback& callback) {
-  Display* host = GetDisplayForWindowManager();
+  WindowManagerState* wms = GetWindowManagerStateForWindowManager();
   const bool success =
-      host &&
-      host->event_dispatcher()->AddAccelerator(id, std::move(event_matcher));
+      wms->event_dispatcher()->AddAccelerator(id, std::move(event_matcher));
   callback.Run(success);
 }
 
 void WindowTree::RemoveAccelerator(uint32_t id) {
-  Display* host = GetDisplayForWindowManager();
-  if (!host)
-    return;
-  host->event_dispatcher()->RemoveAccelerator(id);
+  WindowManagerState* wms = GetWindowManagerStateForWindowManager();
+  wms->event_dispatcher()->RemoveAccelerator(id);
 }
 
 void WindowTree::AddActivationParent(Id transport_window_id) {
diff --git a/components/mus/ws/window_tree.h b/components/mus/ws/window_tree.h
index 378186d..4dcb880 100644
--- a/components/mus/ws/window_tree.h
+++ b/components/mus/ws/window_tree.h
@@ -109,13 +109,19 @@ class WindowTree : public mojom::WindowTree,
         const_cast<const WindowTree*>(this)->GetDisplay(window));
   }
 
+  const WindowManagerState* GetWindowManagerState(
+      const ServerWindow* window) const;
+  WindowManagerState* GetWindowManagerState(const ServerWindow* window) {
+    return const_cast<WindowManagerState*>(
+        const_cast<const WindowTree*>(this)->GetWindowManagerState(window));
+  }
+
   // Invoked when a tree is about to be destroyed.
   void OnWindowDestroyingTreeImpl(WindowTree* tree);
 
-  void OnWillDestroyDisplay(Display* display);
-
   // These functions are synchronous variants of those defined in the mojom. The
   // WindowTree implementations all call into these. See the mojom for details.
+  bool SetCapture(const ClientWindowId& client_window_id);
   bool NewWindow(const ClientWindowId& client_window_id,
                  const std::map<std::string, std::vector<uint8_t>>& properties);
   bool AddWindow(const ClientWindowId& parent_id,
@@ -406,8 +412,8 @@ class WindowTree : public mojom::WindowTree,
 
   uint32_t event_ack_id_;
 
-  // Display the current event came from.
-  Display* event_source_display_;
+  // WindowManager the current event came from.
+  WindowManagerState* event_source_wms_ = nullptr;
 
   bool is_embed_root_;
 
diff --git a/components/mus/ws/window_tree_unittest.cc b/components/mus/ws/window_tree_unittest.cc
index e0b5eb1..61b1d66 100644
--- a/components/mus/ws/window_tree_unittest.cc
+++ b/components/mus/ws/window_tree_unittest.cc
@@ -154,6 +154,10 @@ ClientWindowId FirstRootId(WindowTree* tree) {
              : ClientWindowId();
 }
 
+ServerWindow* GetCaptureWindow(Display* display) {
+  return display->GetActiveWindowManagerState()->capture_window();
+}
+
 }  // namespace
 
 // -----------------------------------------------------------------------------
@@ -197,7 +201,7 @@ class WindowTreeTest : public testing::Test {
   }
 
   void AckPreviousEvent() {
-    DisplayTestApi test_api(display_);
+    WindowManagerStateTestApi test_api(display_->GetActiveWindowManagerState());
     while (test_api.tree_awaiting_input_ack())
       test_api.tree_awaiting_input_ack()->OnWindowInputEventAck(0);
   }
@@ -637,28 +641,28 @@ TEST_F(WindowTreeTest, ExplicitSetCapture) {
   uint32_t change_id = 42;
   mojom_window_tree->SetCapture(change_id, WindowIdToTransportId(window->id()));
   Display* display = tree->GetDisplay(window);
-  EXPECT_NE(window, display->GetCaptureWindow());
+  EXPECT_NE(window, GetCaptureWindow(display));
 
   // Setting capture after the event is acknowledged should fail
   DispatchEventAndAckImmediately(CreatePointerDownEvent(10, 10));
   mojom_window_tree->SetCapture(++change_id,
                                 WindowIdToTransportId(window->id()));
-  EXPECT_NE(window, display->GetCaptureWindow());
+  EXPECT_NE(window, GetCaptureWindow(display));
 
   // Settings while the event is being process should pass
   DispatchEventWithoutAck(CreatePointerDownEvent(10, 10));
   mojom_window_tree->SetCapture(++change_id,
                                 WindowIdToTransportId(window->id()));
-  EXPECT_EQ(window, display->GetCaptureWindow());
+  EXPECT_EQ(window, GetCaptureWindow(display));
   AckPreviousEvent();
 
   // Only the capture window should be able to release capture
   mojom_window_tree->ReleaseCapture(++change_id,
                                     WindowIdToTransportId(root_window->id()));
-  EXPECT_EQ(window, display->GetCaptureWindow());
+  EXPECT_EQ(window, GetCaptureWindow(display));
   mojom_window_tree->ReleaseCapture(++change_id,
                                     WindowIdToTransportId(window->id()));
-  EXPECT_EQ(nullptr, display->GetCaptureWindow());
+  EXPECT_EQ(nullptr, GetCaptureWindow(display));
 }
 
 }  // namespace test