Fix screenshot lock handling and sticky service startup edge cases

README.md

@@ -58,10 +58,10 @@ spawn-sh-at-startup "niri-scripts/support-sticky-floating"
 
 binds {
     // To take a screenshot, use whatever keybind you want
-    Super+S { spawn-sh "niri-scripts/screeenshot <PATH_TO_SAVE_SCREENSHOT>"; }
+    Super+S { spawn-sh "niri-scripts/screenshot <PATH_TO_SAVE_SCREENSHOT>"; }
 
     // To take a screenshot then annotate it, use whatever keybind you want
-    Super+Ctrl+S { spawn-sh "niri-scripts/screeenshot <PATH_TO_SAVE_SCREENSHOT> --annotate"; }
+    Super+Ctrl+S { spawn-sh "niri-scripts/screenshot <PATH_TO_SAVE_SCREENSHOT> --annotate"; }
 
     // To make the focused window sticky, use whatever keybind you want
     Super+Shift+S { spawn-sh "niri-scripts/toggle-sticky"; }

auto-consume-new-windows

@@ -0,0 +1,246 @@
+#!/usr/bin/env scriptisto
+// vim: ft=rust shiftwidth=2 softtabstop=2
+
+// TASK:
+// #region meta
+// scriptisto-begin
+// script_src: src/main.rs
+// build_cmd: >
+//  cargo clippy --color=always &&
+//  cargo build --release --color=always && strip ./target/release/niri_auto_consume
+// target_bin: ./target/release/niri_auto_consume
+// files:
+//  - path: Cargo.toml
+//    content: |
+//     package = { name = "niri_auto_consume", version = "0.1.0", edition = "2024"}
+//     [dependencies]
+//     serde = { version = "1.0.219", features = [ "derive" ] }
+//     serde_json = "1.0.140"
+// scriptisto-end
+// #endregion
+
+#![deny(clippy::unwrap_used)]
+
+use std::collections::HashSet;
+use std::io::{
+	BufRead,
+	BufReader,
+	Write,
+};
+use std::os::unix::net::{UnixListener, UnixStream};
+use std::process::{
+	Command,
+	Stdio,
+};
+use std::sync::{
+	LazyLock,
+	RwLock,
+};
+
+type Result<T> = std::result::Result<T, Box<dyn std::error::Error>>;
+
+static AUTO_CONSUME_ENABLED: LazyLock<RwLock<bool>> = LazyLock::new(|| RwLock::new(true));
+static PROCESSED_WINDOWS: LazyLock<RwLock<HashSet<u64>>> = LazyLock::new(|| RwLock::new(HashSet::new()));
+const IPC_FILE: &str = "/tmp/niri-auto-consume.sock";
+
+#[derive(Debug, serde::Deserialize)]
+struct WindowOpenedOrChanged {
+	id: u64,
+	workspace_id: u64,
+	is_floating: bool,
+}
+
+#[derive(Debug, serde::Deserialize)]
+enum Event {
+	WindowOpenedOrChanged {
+		window: WindowOpenedOrChanged,
+	},
+	WindowClosed {
+		id: u64,
+	},
+}
+
+fn main() -> Result<()> {
+	if std::env::args().any(|e| e == "toggle") {
+		toggle_auto_consume();
+		return Ok(());
+	}
+
+	if std::env::args().any(|e| e == "status") {
+		show_status();
+		return Ok(());
+	}
+
+	std::fs::remove_file(IPC_FILE).ok();
+
+	std::thread::spawn(socket_loop);
+
+	let mut child = Command::new("niri")
+		.arg("msg")
+		.arg("--json")
+		.arg("event-stream")
+		.stdout(Stdio::piped())
+		.spawn()?;
+
+	let stdout = child.stdout.take().expect("Failed to capture stdout");
+
+	let reader = BufReader::new(stdout);
+
+	for line in reader.lines() {
+		match line {
+			Ok(line_content) => {
+				let payload = serde_json::from_str::<Event>(&line_content);
+				match payload {
+					Ok(Event::WindowOpenedOrChanged { window }) => {
+						on_window_opened(window);
+					}
+					Ok(Event::WindowClosed { id }) => {
+						on_window_closed(id);
+					}
+					_ => (),
+				}
+			}
+			Err(e) => {
+				eprintln!("Error reading output: {}", e);
+			}
+		}
+	}
+
+	child.wait()?;
+
+	Ok(())
+}
+
+fn toggle_auto_consume() {
+	// Send toggle command to the running daemon via socket
+	if let Ok(mut stream) = std::os::unix::net::UnixStream::connect(IPC_FILE) {
+		use std::io::{Write, Read};
+		stream.write_all(b"toggle\n").expect("Failed to send toggle command");
+		stream.flush().expect("Failed to flush toggle command");
+
+		// Read and print the response
+		let mut response = String::new();
+		stream.read_to_string(&mut response).expect("Failed to read response");
+		print!("{}", response);
+	} else {
+		eprintln!("Failed to connect to auto-consume daemon. Is it running?");
+	}
+}
+
+fn show_status() {
+	// Send status command to the running daemon via socket
+	if let Ok(mut stream) = std::os::unix::net::UnixStream::connect(IPC_FILE) {
+		use std::io::{Write, Read};
+		stream.write_all(b"status\n").expect("Failed to send status command");
+		stream.flush().expect("Failed to flush status command");
+
+		// Read and print the response
+		let mut response = String::new();
+		stream.read_to_string(&mut response).expect("Failed to read response");
+		print!("{}", response);
+	} else {
+		eprintln!("Failed to connect to auto-consume daemon. Is it running?");
+	}
+}
+
+fn socket_loop() {
+	let listener = UnixListener::bind(IPC_FILE).expect("We couldn't make the listener");
+
+	for stream in listener.incoming().flatten() {
+		let mut reader = BufReader::new(stream.try_clone().expect("Failed to clone stream"));
+		let mut writer = reader.get_mut().try_clone().expect("Failed to get writer");
+		let mut line = String::new();
+		if reader.read_line(&mut line).is_ok() {
+			let command = line.trim();
+			match command {
+				"toggle" => {
+					let response = handle_toggle();
+					let _ = writer.write_all(response.as_bytes());
+					let _ = writer.flush();
+				}
+				"status" => {
+					let response = handle_status();
+					let _ = writer.write_all(response.as_bytes());
+					let _ = writer.flush();
+				}
+				_ => {
+					let response = format!("Unknown command: {}\n", command);
+					let _ = writer.write_all(response.as_bytes());
+					let _ = writer.flush();
+				}
+			}
+		}
+	}
+}
+
+fn handle_toggle() -> String {
+	let mut enabled = AUTO_CONSUME_ENABLED.write().expect("Failed to acquire write lock");
+	*enabled = !*enabled;
+
+	if *enabled {
+		"Auto consume enabled\n".to_string()
+	} else {
+		"Auto consume disabled\n".to_string()
+	}
+}
+
+fn handle_status() -> String {
+	let enabled = AUTO_CONSUME_ENABLED.read().expect("Failed to acquire read lock");
+	if *enabled {
+		"Auto consume: enabled\n".to_string()
+	} else {
+		"Auto consume: disabled\n".to_string()
+	}
+}
+
+fn consume_window(window_id: u64) -> Result<()> {
+	let output = Command::new("niri")
+		.args(["msg", "action", "consume-or-expel-window-left"])
+		.arg("--id")
+		.arg(window_id.to_string())
+		.output()?;
+
+	if !output.status.success() {
+		let stderr = String::from_utf8_lossy(&output.stderr);
+		return Err(format!("Failed to consume window {}: {}", window_id, stderr).into());
+	}
+
+	Ok(())
+}
+
+fn on_window_opened(window: WindowOpenedOrChanged) {
+	let enabled = {
+		*AUTO_CONSUME_ENABLED.read().expect("Failed to acquire read lock")
+	};
+
+	if !enabled {
+		return;
+	}
+
+	// Check if we've already processed this window
+	let already_processed = {
+		let mut processed = PROCESSED_WINDOWS.write().expect("Failed to acquire write lock");
+		if processed.contains(&window.id) {
+			true
+		} else {
+			processed.insert(window.id);
+			false
+		}
+	};
+
+	// If already processed, ignore this event
+	if already_processed {
+		return;
+	}
+
+	// Consume the window automatically regardless of floating state
+	if let Err(e) = consume_window(window.id) {
+		eprintln!("Failed to consume window {}: {}", window.id, e);
+	}
+}
+
+fn on_window_closed(window_id: u64) {
+	// Remove window from processed set when it closes to prevent memory leaks
+	let mut processed = PROCESSED_WINDOWS.write().expect("Failed to acquire write lock");
+	processed.remove(&window_id);
+}