diff --git a/crates/border-wars/src/lib.rs b/crates/border-wars/src/lib.rs
index 700aa3b..d967357 100644
--- a/crates/border-wars/src/lib.rs
+++ b/crates/border-wars/src/lib.rs
@@ -1,14 +1,19 @@
 //! The file that contains utility functions, enums, structs of the game.
 
 use std::default::Default;
+
 use bevy::prelude::*;
 
 pub mod menu;
 
+/// The state of the game.
 #[derive(Debug, Clone, Copy, Default, Eq, PartialEq, Hash, States)]
 pub enum GameState {
+    /// The main menu.
     #[default]
     Menu,
+    /// The lobby where the host will wait for other players to join.
     Lobby,
+    /// The game.
     Game,
 }
diff --git a/crates/border-wars/src/menu.rs b/crates/border-wars/src/menu.rs
index 30f5371..11a400e 100644
--- a/crates/border-wars/src/menu.rs
+++ b/crates/border-wars/src/menu.rs
@@ -12,10 +12,8 @@ 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)),
-        );
+        app.add_plugins(EguiPlugin)
+            .add_systems(Update, menu_ui.run_if(in_state(GameState::Menu)));
         println!("{:?}", app.world.resource::<State<GameState>>());
         app.add_plugins(EguiPlugin).add_systems(
             Update,
@@ -24,7 +22,7 @@ impl Plugin for MenuPlugin {
     }
 }
 /// Display a window, with a button to join a game and an input to enter
-/// the address or create a new game.
+/// TODO
 fn menu_ui(mut ctx: EguiContexts, mut connection_string: Local<String>) {
     egui::CentralPanel::default().show(ctx.ctx_mut(), |ui| {
         ui.heading("Border Wars");