3 files changed, 29 insertions, 7 deletions
diff --git a/azalea-client/src/client.rs b/azalea-client/src/client.rs
index 131aef16..9de97cc1 100644
--- a/azalea-client/src/client.rs
+++ b/azalea-client/src/client.rs
@@ -76,7 +76,7 @@ use crate::{
     mining::{self, MinePlugin},
     movement::{LastSentLookDirection, PhysicsState, PlayerMovePlugin},
     packet_handling::{
-        login::{self, LoginSendPacketQueue},
+        login::{self, InLoginState, LoginSendPacketQueue},
         PacketHandlerPlugin,
     },
     player::retroactively_add_game_profile_component,
@@ -371,6 +371,7 @@ impl Client {
         ecs_lock.lock().entity_mut(entity).insert((
             LoginSendPacketQueue { tx: ecs_packets_tx },
             login::IgnoreQueryIds::default(),
+            InLoginState,
         ));
 
         // login
@@ -457,6 +458,7 @@ impl Client {
                         p.game_profile
                     );
                     conn.write(ServerboundLoginAcknowledged {}).await?;
+
                     break (conn.config(), p.game_profile);
                 }
                 ClientboundLoginPacket::LoginDisconnect(p) => {
@@ -485,7 +487,8 @@ impl Client {
             .lock()
             .entity_mut(entity)
             .remove::<login::IgnoreQueryIds>()
-            .remove::<LoginSendPacketQueue>();
+            .remove::<LoginSendPacketQueue>()
+            .remove::<InLoginState>();
 
         Ok((conn, profile))
     }
diff --git a/azalea-client/src/configuration.rs b/azalea-client/src/configuration.rs
index bf07710b..0f198dbe 100644
--- a/azalea-client/src/configuration.rs
+++ b/azalea-client/src/configuration.rs
@@ -9,25 +9,30 @@ use azalea_protocol::{
 };
 use bevy_app::prelude::*;
 use bevy_ecs::prelude::*;
+use tracing::{debug, warn};
 
-use crate::{client::InConfigState, packet_handling::configuration::SendConfigurationEvent};
+use crate::{
+    client::InConfigState,
+    packet_handling::{configuration::SendConfigurationEvent, login::InLoginState},
+};
 
 pub struct ConfigurationPlugin;
 impl Plugin for ConfigurationPlugin {
     fn build(&self, app: &mut App) {
         app.add_systems(
             Update,
-            handle_in_configuration_state
+            handle_end_login_state
                 .before(crate::packet_handling::configuration::handle_send_packet_event),
         );
     }
 }
 
-fn handle_in_configuration_state(
-    query: Query<(Entity, &ClientInformation), Added<InConfigState>>,
+fn handle_end_login_state(
+    mut removed: RemovedComponents<InLoginState>,
+    query: Query<&ClientInformation>,
     mut send_packet_events: EventWriter<SendConfigurationEvent>,
 ) {
-    for (entity, client_information) in query.iter() {
+    for entity in removed.read() {
         let mut brand_data = Vec::new();
         // they don't have to know :)
         "vanilla".azalea_write(&mut brand_data).unwrap();
@@ -39,6 +44,15 @@ fn handle_in_configuration_state(
             },
         ));
 
+        let client_information = match query.get(entity).ok() {
+            Some(i) => i,
+            None => {
+                warn!("ClientInformation component was not set before leaving login state, using a default");
+                &ClientInformation::default()
+            }
+        };
+
+        debug!("Writing ClientInformation while in config state: {client_information:?}");
         send_packet_events.send(SendConfigurationEvent::new(
             entity,
             ServerboundClientInformation {
diff --git a/azalea-client/src/packet_handling/login.rs b/azalea-client/src/packet_handling/login.rs
index 11c0b8e9..dec4aa06 100644
--- a/azalea-client/src/packet_handling/login.rs
+++ b/azalea-client/src/packet_handling/login.rs
@@ -48,6 +48,11 @@ pub struct LoginSendPacketQueue {
     pub tx: mpsc::UnboundedSender<ServerboundLoginPacket>,
 }
 
+/// A marker component for local players that are currently in the
+/// `login` state.
+#[derive(Component, Clone, Debug)]
+pub struct InLoginState;
+
 pub fn handle_send_packet_event(
     mut send_packet_events: EventReader<SendLoginPacketEvent>,
     mut query: Query<&mut LoginSendPacketQueue>,