diff --git a/crates/border-wars/src/lib.rs b/crates/border-wars/src/lib.rs
index 62b07b8..700aa3b 100644
--- a/crates/border-wars/src/lib.rs
+++ b/crates/border-wars/src/lib.rs
@@ -1,9 +1,10 @@
 //! The file that contains utility functions, enums, structs of the game.
 
 use std::default::Default;
-
 use bevy::prelude::*;
 
+pub mod menu;
+
 #[derive(Debug, Clone, Copy, Default, Eq, PartialEq, Hash, States)]
 pub enum GameState {
     #[default]
diff --git a/crates/border-wars/src/main.rs b/crates/border-wars/src/main.rs
index c3561cf..816ec63 100644
--- a/crates/border-wars/src/main.rs
+++ b/crates/border-wars/src/main.rs
@@ -1,11 +1,8 @@
 //! The main entry point of the game.
 
 use bevy::prelude::*;
-use lib::GameState;
-use menu::MenuPlugin;
-
-mod lib;
-mod menu;
+use border_wars::menu::MenuPlugin;
+use border_wars::GameState;
 
 fn main() {
     App::new()
diff --git a/crates/border-wars/src/menu.rs b/crates/border-wars/src/menu.rs
index 90b1255..30f5371 100644
--- a/crates/border-wars/src/menu.rs
+++ b/crates/border-wars/src/menu.rs
@@ -12,6 +12,10 @@ pub struct MenuPlugin;
 impl Plugin for MenuPlugin {
     /// A function that is called when the plugin is added to the application.
     fn build(&self, app: &mut App) {
+        app.add_plugins(EguiPlugin).add_systems(
+            Update,
+            (ui_connect_window, ui_host_window).run_if(in_state(GameState::Menu)),
+        );
         println!("{:?}", app.world.resource::<State<GameState>>());
         app.add_plugins(EguiPlugin).add_systems(
             Update,