3 files changed, 22 insertions, 21 deletions

diff --git a/azalea-client/src/account.rs b/azalea-client/src/account.rs
index 46d0a728..d045cc72 100755
--- a/azalea-client/src/account.rs
+++ b/azalea-client/src/account.rs
@@ -15,6 +15,8 @@ use uuid::Uuid;
 /// To join a server using this account, use [`Client::join`] or
 /// [`azalea::ClientBuilder`].
 ///
+/// Note that this is also a component that our clients have.
+///
 /// # Examples
 ///
 /// ```rust,no_run
diff --git a/azalea-client/src/client.rs b/azalea-client/src/client.rs
index 561c5a24..0cf01582 100644
--- a/azalea-client/src/client.rs
+++ b/azalea-client/src/client.rs
@@ -251,24 +251,15 @@ impl Client {
         resolved_address: &SocketAddr,
         run_schedule_sender: mpsc::UnboundedSender<()>,
     ) -> Result<(Self, mpsc::UnboundedReceiver<Event>), JoinError> {
+        let entity = ecs_lock.lock().spawn(account.to_owned()).id();
+
         let conn = Connection::new(resolved_address).await?;
         let (conn, game_profile) = Self::handshake(conn, account, address).await?;
         let (read_conn, write_conn) = conn.into_split();
 
-        let (tx, rx) = mpsc::unbounded_channel();
-
-        let mut ecs = ecs_lock.lock();
-
-        // Make the ecs entity for this client
-        let entity = ecs.spawn_empty().id();
+        // we did the handshake, so now we're connected to the server
 
-        // we got the GameConnection, so the server is now connected :)
-        let client = Client::new(
-            game_profile.clone(),
-            entity,
-            ecs_lock.clone(),
-            run_schedule_sender.clone(),
-        );
+        let (tx, rx) = mpsc::unbounded_channel();
 
         let (packet_writer_sender, packet_writer_receiver) = mpsc::unbounded_channel();
 
@@ -295,12 +286,13 @@ impl Client {
             write_packets_task,
         );
 
-        ecs.entity_mut(entity).insert((
-            account.to_owned(),
-            JoinedClientBundle {
+        ecs_lock
+            .lock()
+            .entity_mut(entity)
+            .insert(JoinedClientBundle {
                 local_player,
                 packet_receiver,
-                game_profile: GameProfileComponent(game_profile),
+                game_profile: GameProfileComponent(game_profile.clone()),
                 physics_state: PhysicsState::default(),
                 local_player_events: LocalPlayerEvents(tx),
                 inventory: InventoryComponent::default(),
@@ -318,9 +310,14 @@ impl Client {
                 attack: attack::AttackBundle::default(),
 
                 _local: LocalEntity,
-            },
-        ));
+            });
 
+        let client = Client::new(
+            game_profile,
+            entity,
+            ecs_lock.clone(),
+            run_schedule_sender.clone(),
+        );
         Ok((client, rx))
     }
 
diff --git a/azalea/src/swarm/mod.rs b/azalea/src/swarm/mod.rs
index ba38a17e..66e5ea23 100644
--- a/azalea/src/swarm/mod.rs
+++ b/azalea/src/swarm/mod.rs
@@ -547,9 +547,11 @@ impl Swarm {
                 }
             }
             cloned_bots.lock().remove(&bot.entity);
-            let owned_account = cloned_bot.component::<Account>();
+            let account = cloned_bot
+                .get_component::<Account>()
+                .expect("bot is missing required Account component");
             swarm_tx
-                .send(SwarmEvent::Disconnect(Box::new(owned_account)))
+                .send(SwarmEvent::Disconnect(Box::new(account)))
                 .unwrap();
         });