From 173b03396340261b42f89a088b9c517b8f8e7e55 Mon Sep 17 00:00:00 2001
From: Nova <technobaboo@proton.me>
Date: Mon, 7 Apr 2025 09:16:59 -0700
Subject: [PATCH] fix(input): unresponsive clients get uncaptured

---
 src/core/client.rs        | 19 +++++++++++++++++--
 src/main.rs               |  3 ++-
 src/nodes/input/method.rs | 10 +++++++++-
 3 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/src/core/client.rs b/src/core/client.rs
index adf7e30..fbbcda7 100644
--- a/src/core/client.rs
+++ b/src/core/client.rs
@@ -23,12 +23,14 @@ use std::{
 	iter::FromIterator,
 	path::PathBuf,
 	sync::{Arc, OnceLock},
+	time::Instant,
 };
-use tokio::{net::UnixStream, task::JoinHandle};
+use tokio::{net::UnixStream, sync::watch, task::JoinHandle};
 use tracing::info;
 
 lazy_static! {
 	pub static ref CLIENTS: OwnedRegistry<Client> = OwnedRegistry::new();
+	static ref INTERNAL_CLIENT_MESSAGE_TIMES: (watch::Sender<Instant>, watch::Receiver<Instant>) = watch::channel(Instant::now());
 	pub static ref INTERNAL_CLIENT: Arc<Client> = CLIENTS.add(Client {
 		pid: None,
 		// env: None,
@@ -38,14 +40,18 @@ lazy_static! {
 		flush_join_handle: OnceLock::new(),
 		disconnect_status: OnceLock::new(),
 
-		message_sender_handle: None,
 		id_counter: CounterU32::new(0),
+		message_last_received: INTERNAL_CLIENT_MESSAGE_TIMES.1.clone(),
+		message_sender_handle: None,
 		scenegraph: Default::default(),
 		root: OnceLock::new(),
 		base_resource_prefixes: Default::default(),
 		state: OnceLock::default(),
 	});
 }
+pub fn tick_internal_client() {
+	let _ = INTERNAL_CLIENT_MESSAGE_TIMES.0.send(Instant::now());
+}
 
 pub fn get_env(pid: i32) -> Result<FxHashMap<String, String>, std::io::Error> {
 	let env = fs::read_to_string(format!("/proc/{pid}/environ"))?;
@@ -69,6 +75,7 @@ pub struct Client {
 	disconnect_status: OnceLock<Result<()>>,
 
 	id_counter: CounterU32,
+	message_last_received: watch::Receiver<Instant>,
 	pub message_sender_handle: Option<MessageSenderHandle>,
 	pub scenegraph: Arc<Scenegraph>,
 	pub root: OnceLock<Arc<Root>>,
@@ -95,6 +102,7 @@ impl Client {
 			.and_then(state)
 			.unwrap_or_else(|| Arc::new(ClientStateParsed::default()));
 
+		let (message_time_tx, message_last_received) = watch::channel(Instant::now());
 		let client = CLIENTS.add(Client {
 			pid,
 			// env,
@@ -105,6 +113,7 @@ impl Client {
 			disconnect_status: OnceLock::new(),
 
 			id_counter: CounterU32::new(256),
+			message_last_received,
 			message_sender_handle: Some(messenger_tx.handle()),
 			scenegraph: scenegraph.clone(),
 			root: OnceLock::new(),
@@ -148,6 +157,7 @@ impl Client {
 							if let Err(e) = messenger_rx.dispatch(&*scenegraph).await {
 								client.disconnect(Err(e.into()));
 							}
+							let _ = message_time_tx.send(Instant::now());
 						}
 					}
 				},
@@ -206,6 +216,11 @@ impl Client {
 			.ok_or_else(|| eyre!("{} not found", name))
 	}
 
+	pub fn unresponsive(&self) -> bool {
+		let time_since_last_message = self.message_last_received.borrow().elapsed();
+		time_since_last_message.as_millis() > 500
+	}
+
 	pub fn disconnect(&self, reason: Result<()>) {
 		let _ = self.disconnect_status.set(reason);
 		if let Some(dispatch_join_handle) = self.dispatch_join_handle.get() {
diff --git a/src/main.rs b/src/main.rs
index da63feb..5693093 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -11,7 +11,7 @@ use crate::nodes::items::camera;
 use crate::nodes::{audio, drawable, input};
 
 use clap::Parser;
-use core::client::Client;
+use core::client::{Client, tick_internal_client};
 use core::task;
 use directories::ProjectDirs;
 use objects::ServerObjects;
@@ -307,6 +307,7 @@ fn stereokit_loop(
 			Duration::from_micros(250),
 		);
 
+		tick_internal_client();
 		#[cfg(feature = "wayland")]
 		wayland.update();
 		drawable::draw(token);
diff --git a/src/nodes/input/method.rs b/src/nodes/input/method.rs
index 663f03b..c7fb645 100644
--- a/src/nodes/input/method.rs
+++ b/src/nodes/input/method.rs
@@ -159,7 +159,15 @@ impl InputMethod {
 			.iter()
 			.filter_map(Weak::upgrade)
 			.collect::<Registry<InputHandler>>();
-		self.captures.retain(|handler| sent.contains(handler));
+		self.captures.retain(|handler| {
+			!handler
+				.spatial
+				.node()
+				.and_then(|n| n.get_client())
+				.map(|c| c.unresponsive())
+				.unwrap_or(false)
+				&& sent.contains(handler)
+		});
 	}
 }
 impl InputMethodAspect for InputMethod {