disconnect fixes

1 files changed, 15 insertions, 6 deletions

diff --git a/azalea-client/src/disconnect.rs b/azalea-client/src/disconnect.rs
index ac663a66..35334199 100644
--- a/azalea-client/src/disconnect.rs
+++ b/azalea-client/src/disconnect.rs
@@ -1,18 +1,20 @@
 //! Disconnect a client from the server.
 
+use azalea_chat::FormattedText;
+use azalea_entity::LocalEntity;
 use bevy_app::{App, Plugin, PostUpdate};
 use bevy_ecs::{
     component::Component,
     entity::Entity,
     event::{EventReader, EventWriter},
     prelude::Event,
-    query::Changed,
+    query::{Changed, With},
     schedule::IntoSystemConfigs,
     system::{Commands, Query},
 };
 use derive_more::Deref;
 
-use crate::{client::JoinedClientBundle, raw_connection::RawConnection};
+use crate::{client::JoinedClientBundle, events::LocalPlayerEvents, raw_connection::RawConnection};
 
 pub struct DisconnectPlugin;
 impl Plugin for DisconnectPlugin {
@@ -33,7 +35,7 @@ impl Plugin for DisconnectPlugin {
 #[derive(Event)]
 pub struct DisconnectEvent {
     pub entity: Entity,
-    pub reason: Option<String>,
+    pub reason: Option<FormattedText>,
 }
 
 /// System that removes the [`JoinedClientBundle`] from the entity when it
@@ -43,7 +45,11 @@ pub fn remove_components_from_disconnected_players(
     mut events: EventReader<DisconnectEvent>,
 ) {
     for DisconnectEvent { entity, .. } in events.read() {
-        commands.entity(*entity).remove::<JoinedClientBundle>();
+        commands
+            .entity(*entity)
+            .remove::<JoinedClientBundle>()
+            // swarm detects when this tx gets dropped to fire SwarmEvent::Disconnect
+            .remove::<LocalPlayerEvents>();
     }
 }
 
@@ -60,12 +66,15 @@ fn update_read_packets_task_running_component(
     }
 }
 fn disconnect_on_connection_dead(
-    query: Query<(Entity, &IsConnectionAlive), Changed<IsConnectionAlive>>,
+    query: Query<(Entity, &IsConnectionAlive), (Changed<IsConnectionAlive>, With<LocalEntity>)>,
     mut disconnect_events: EventWriter<DisconnectEvent>,
 ) {
     for (entity, &is_connection_alive) in &query {
         if !*is_connection_alive {
-            disconnect_events.send(DisconnectEvent { entity, reason: None });
+            disconnect_events.send(DisconnectEvent {
+                entity,
+                reason: None,
+            });
         }
     }
 }