This repository has no description
20 kB
571 lines
1#!/usr/bin/env python3
2"""
3Workspace Manager Daemon
4
5A compositor-agnostic workspace manager that:
6- Subscribes to sway IPC for workspace change events
7- Runs user-defined hooks on workspace enter/leave
8- Provides current workspace and directory info via Unix socket
9
10Config merging: user config is loaded first, then system config (declarative) overrides per-field
11"""
12
13from __future__ import annotations
14
15import asyncio
16import json
17import os
18import signal
19import subprocess
20import sys
21from pathlib import Path
22from typing import Any
23
24# XDG runtime directory for socket and state
25RUNTIME_DIR = Path(os.environ.get("XDG_RUNTIME_DIR", f"/run/user/{os.getuid()}"))
26SOCKET_DIR = RUNTIME_DIR / "workspace-manager"
27SOCKET_PATH = SOCKET_DIR / "socket"
28STATE_PATH = SOCKET_DIR / "state.json"
29
30# System config directory (declarative, takes precedence)
31_system_config_dir_env = os.environ.get("WORKSPACE_MANAGER_SYSTEM_CONFIG_DIR")
32SYSTEM_CONFIG_DIR: Path | None = (
33 Path(_system_config_dir_env).expanduser() if _system_config_dir_env else None
34)
35
36# User config directory (fallback)
37_config_dir_env = os.environ.get("WORKSPACE_MANAGER_CONFIG_DIR")
38CONFIG_DIR = (
39 Path(_config_dir_env).expanduser()
40 if _config_dir_env
41 else Path.home() / "workspaces"
42)
43
44_default_dir_env = os.environ.get("WORKSPACE_MANAGER_DEFAULT_DIR")
45DEFAULT_DIRECTORY = (
46 Path(_default_dir_env).expanduser() if _default_dir_env else Path.home()
47)
48
49TERMINAL_COMMAND = os.environ.get("WORKSPACE_MANAGER_TERMINAL", "kitty")
50
51_on_new_handler_env = os.environ.get("WORKSPACE_MANAGER_ON_NEW_HANDLER")
52ON_NEW_HANDLER: str | None = _on_new_handler_env if _on_new_handler_env else None
53
54
55def load_config_file(config_file: Path) -> dict[str, Any] | None:
56 """Load a workspace config file, returning None if it doesn't exist or fails."""
57 if not config_file.exists():
58 return None
59 try:
60 with open(config_file) as f:
61 return json.load(f) # type: ignore[no-any-return]
62 except (json.JSONDecodeError, OSError) as e:
63 print(f"Warning: Failed to load config {config_file}: {e}", file=sys.stderr)
64 return None
65
66
67class WorkspaceConfig:
68 """Configuration for a workspace.
69
70 Merges user config with system config (declarative).
71 System config fields take precedence over user config fields.
72 """
73
74 def __init__(self, name: str) -> None:
75 self.name = name
76 self.directory: Path = DEFAULT_DIRECTORY
77 self.on_enter: list[str] = []
78 self.on_leave: list[str] = []
79 self.on_create: list[str] = []
80 self._load()
81
82 def _load(self) -> None:
83 # Load user config first (base)
84 user_config_file = CONFIG_DIR / f"{self.name}.json"
85 user_data = load_config_file(user_config_file)
86
87 # Load system config (declarative, takes precedence)
88 system_data: dict[str, Any] | None = None
89 if SYSTEM_CONFIG_DIR is not None:
90 system_config_file = SYSTEM_CONFIG_DIR / f"{self.name}.json"
91 system_data = load_config_file(system_config_file)
92
93 # Merge: start with user config, override with system config
94 merged: dict[str, Any] = {}
95 if user_data is not None:
96 merged.update(user_data)
97 if system_data is not None:
98 merged.update(system_data)
99
100 # Apply merged config
101 if "directory" in merged:
102 self.directory = Path(str(merged["directory"])).expanduser()
103 if "on_enter" in merged:
104 self.on_enter = list(merged["on_enter"])
105 if "on_leave" in merged:
106 self.on_leave = list(merged["on_leave"])
107 if "on_create" in merged:
108 self.on_create = list(merged["on_create"])
109
110
111class WorkspaceManager:
112 """Manages workspace state and configuration"""
113
114 def __init__(self) -> None:
115 self.current_workspace: str | None = None
116 self.configs: dict[str, WorkspaceConfig] = {}
117 self.created_workspaces: set[str] = set()
118
119 def get_config(self, workspace: str) -> WorkspaceConfig:
120 """Get or create config for a workspace (configs are loaded fresh each time)"""
121 # Always reload config to pick up changes
122 self.configs[workspace] = WorkspaceConfig(workspace)
123 return self.configs[workspace]
124
125 def get_current_directory(self) -> Path:
126 """Get the directory for the current workspace"""
127 if self.current_workspace:
128 config = self.get_config(self.current_workspace)
129 return config.directory
130 return DEFAULT_DIRECTORY
131
132 def _tmux_session_exists(self, session_name: str) -> bool:
133 """Check if a tmux session already exists"""
134 try:
135 result = subprocess.run(
136 ["tmux", "has-session", "-t", session_name],
137 capture_output=True,
138 )
139 return result.returncode == 0
140 except FileNotFoundError:
141 print("Warning: tmux not found", file=sys.stderr)
142 return False
143
144 async def _run_on_create(self, workspace: str, config: WorkspaceConfig) -> None:
145 """Start on_create commands each in their own tmux session + terminal.
146
147 For each command, creates a tmux session named <workspace>-start<N>
148 running the command, then opens a terminal window attached to it.
149 If the terminal is closed, reattach with: tmux attach -t <workspace>-start<N>
150 """
151 if not config.on_create:
152 return
153
154 work_dir = str(config.directory)
155
156 for i, cmd in enumerate(config.on_create, start=1):
157 session_name = f"{workspace}-start{i}"
158
159 # Skip if tmux session already exists (e.g. from a previous run)
160 if self._tmux_session_exists(session_name):
161 print(
162 f"Tmux session '{session_name}' already exists, skipping",
163 file=sys.stderr,
164 )
165 continue
166
167 # Create a detached tmux session running the command
168 try:
169 subprocess.run(
170 [
171 "tmux",
172 "new-session",
173 "-d",
174 "-s",
175 session_name,
176 "-c",
177 work_dir,
178 cmd,
179 ],
180 check=True,
181 )
182 print(
183 f"Created tmux session '{session_name}': {cmd}",
184 file=sys.stderr,
185 )
186 except (OSError, subprocess.CalledProcessError) as e:
187 print(
188 f"Failed to create tmux session '{session_name}': {e}",
189 file=sys.stderr,
190 )
191 continue
192
193 # Open a terminal window attached to the tmux session
194 try:
195 subprocess.Popen(
196 [
197 TERMINAL_COMMAND,
198 "-T",
199 session_name,
200 "tmux",
201 "attach",
202 "-t",
203 session_name,
204 ],
205 start_new_session=True,
206 )
207 print(
208 f"Opened terminal for tmux session '{session_name}'",
209 file=sys.stderr,
210 )
211 except OSError as e:
212 print(
213 f"Failed to open terminal for '{session_name}': {e}",
214 file=sys.stderr,
215 )
216
217 def _workspace_has_windows(self, workspace: str) -> bool:
218 """Check if a workspace already has windows (via sway tree)."""
219 try:
220 result = subprocess.run(
221 ["swaymsg", "-t", "get_tree", "-r"],
222 capture_output=True,
223 text=True,
224 timeout=2,
225 )
226 if result.returncode != 0:
227 return False
228 tree: dict[str, Any] = json.loads(result.stdout)
229 # Find the workspace node and check for app containers
230 for output in tree.get("nodes", []):
231 for ws in output.get("nodes", []):
232 if ws.get("name") == workspace:
233
234 def has_apps(node: dict[str, Any]) -> bool:
235 if node.get("type") == "con" and (
236 node.get("app_id") or node.get("pid")
237 ):
238 return True
239 for child in node.get("nodes", []) + node.get(
240 "floating_nodes", []
241 ):
242 if has_apps(child):
243 return True
244 return False
245
246 return has_apps(ws)
247 except (OSError, json.JSONDecodeError, subprocess.TimeoutExpired) as e:
248 print(f"Failed to check workspace windows: {e}", file=sys.stderr)
249 return False
250
251 async def _ask_on_new(self, workspace: str) -> bool:
252 """Ask the user whether to restore a workspace via the on-new handler.
253
254 Returns True if on_create should be run, False to skip.
255 If no handler is configured, always returns True (auto-restore).
256 Skips the prompt if the workspace already has windows.
257 """
258 if self._workspace_has_windows(workspace):
259 return True
260 if ON_NEW_HANDLER is None:
261 return True
262 try:
263 proc = await asyncio.create_subprocess_exec(
264 ON_NEW_HANDLER,
265 workspace,
266 stdout=asyncio.subprocess.PIPE,
267 stderr=asyncio.subprocess.PIPE,
268 )
269 _, stderr = await proc.communicate()
270 if stderr:
271 print(
272 f"on-new handler stderr: {stderr.decode().strip()}",
273 file=sys.stderr,
274 )
275 # Exit 0 = restore, anything else = skip
276 return proc.returncode == 0
277 except FileNotFoundError:
278 print(
279 f"on-new handler not found: {ON_NEW_HANDLER}",
280 file=sys.stderr,
281 )
282 return True
283 except Exception as e:
284 print(f"on-new handler error: {e}", file=sys.stderr)
285 return True
286
287 async def handle_workspace_change(
288 self, old_workspace: str | None, new_workspace: str
289 ) -> bool:
290 """Handle a workspace focus change event.
291
292 Returns True if the workspace should be restored, False to start fresh.
293 """
294 # Run on_leave hooks for old workspace
295 if old_workspace:
296 old_config = self.get_config(old_workspace)
297 for hook in old_config.on_leave:
298 await self._run_hook(hook, "on_leave", old_workspace)
299
300 # Update current workspace
301 self.current_workspace = new_workspace
302
303 # Run on_create hooks if this is the first time entering the workspace
304 should_restore = True
305 new_config = self.get_config(new_workspace)
306 if new_workspace not in self.created_workspaces:
307 self.created_workspaces.add(new_workspace)
308 should_restore = await self._ask_on_new(new_workspace)
309 if should_restore:
310 await self._run_on_create(new_workspace, new_config)
311
312 # Run on_enter hooks for new workspace
313 for hook in new_config.on_enter:
314 await self._run_hook(hook, "on_enter", new_workspace)
315
316 # Write state file
317 self._write_state()
318
319 return should_restore
320
321 async def _run_hook(self, command: str, hook_type: str, workspace: str) -> None:
322 """Run a hook command in the background"""
323 try:
324 # Run in background, don't wait
325 subprocess.Popen(
326 command,
327 shell=True,
328 stdout=subprocess.DEVNULL,
329 stderr=subprocess.DEVNULL,
330 start_new_session=True,
331 )
332 print(f"Ran {hook_type} hook for {workspace}: {command}", file=sys.stderr)
333 except OSError as e:
334 print(
335 f"Failed to run {hook_type} hook for {workspace}: {e}",
336 file=sys.stderr,
337 )
338
339 def _write_state(self) -> None:
340 """Write current state to state file"""
341 state: dict[str, Any] = {
342 "workspace": self.current_workspace,
343 "directory": str(self.get_current_directory()),
344 }
345 try:
346 STATE_PATH.parent.mkdir(parents=True, exist_ok=True)
347 with open(STATE_PATH, "w") as f:
348 json.dump(state, f)
349 except OSError as e:
350 print(f"Failed to write state file: {e}", file=sys.stderr)
351
352
353manager = WorkspaceManager()
354
355# Subscribers: list of (StreamWriter, Event) tuples
356subscribers: list[tuple[asyncio.StreamWriter, asyncio.Event]] = []
357
358
359async def notify_subscribers(
360 old_workspace: str | None, new_workspace: str, *, restore: bool = True
361) -> None:
362 """Notify all subscribers of a workspace change."""
363 event = (
364 json.dumps(
365 {
366 "event": "workspace_change",
367 "old": old_workspace,
368 "new": new_workspace,
369 "restore": restore,
370 }
371 ).encode()
372 + b"\n"
373 )
374 for writer, done in subscribers:
375 try:
376 writer.write(event)
377 await writer.drain()
378 except (OSError, ConnectionResetError):
379 done.set()
380
381
382async def subscribe_sway() -> None:
383 """Subscribe to sway IPC workspace events"""
384 print("Starting sway workspace subscription...", file=sys.stderr)
385
386 while True:
387 try:
388 proc = await asyncio.create_subprocess_exec(
389 "swaymsg",
390 "-t",
391 "subscribe",
392 "-m",
393 '["workspace"]',
394 stdout=asyncio.subprocess.PIPE,
395 stderr=asyncio.subprocess.PIPE,
396 )
397
398 if proc.stdout is None:
399 print("Failed to get stdout from swaymsg", file=sys.stderr)
400 await asyncio.sleep(5)
401 continue
402
403 async for line in proc.stdout:
404 try:
405 event: dict[str, Any] = json.loads(line.decode().strip())
406 if event.get("change") == "focus":
407 current: dict[str, Any] = event.get("current", {})
408 old: dict[str, Any] | None = event.get("old")
409 new_name: str | None = current.get("name")
410 old_name: str | None = old.get("name") if old else None
411
412 if new_name:
413 print(
414 f"Workspace change: {old_name} -> {new_name}",
415 file=sys.stderr,
416 )
417 restore = await manager.handle_workspace_change(
418 old_name, new_name
419 )
420 await notify_subscribers(
421 old_name, new_name, restore=restore
422 )
423 except json.JSONDecodeError as e:
424 print(f"Failed to parse sway event: {e}", file=sys.stderr)
425
426 # If swaymsg exits, wait a bit and retry
427 await proc.wait()
428 print("swaymsg exited, retrying in 5 seconds...", file=sys.stderr)
429 await asyncio.sleep(5)
430
431 except FileNotFoundError:
432 print("swaymsg not found, retrying in 10 seconds...", file=sys.stderr)
433 await asyncio.sleep(10)
434 except Exception as e:
435 print(
436 f"Error in sway subscription: {e}, retrying in 5 seconds...",
437 file=sys.stderr,
438 )
439 await asyncio.sleep(5)
440
441
442async def handle_client(
443 reader: asyncio.StreamReader, writer: asyncio.StreamWriter
444) -> None:
445 """Handle a client connection to the Unix socket"""
446 try:
447 data = await reader.read(1024)
448 command = data.decode().strip()
449
450 response: dict[str, Any]
451 if command == "subscribe":
452 # Keep connection open and send workspace change events.
453 # Disconnection is detected in notify_subscribers when write fails.
454 done = asyncio.Event()
455 subscribers.append((writer, done))
456 print("New subscriber connected", file=sys.stderr)
457 try:
458 await done.wait()
459 except asyncio.CancelledError:
460 pass
461 finally:
462 subscribers[:] = [(w, d) for w, d in subscribers if w is not writer]
463 print("Subscriber disconnected", file=sys.stderr)
464 return
465 elif command == "dir":
466 response = {"directory": str(manager.get_current_directory())}
467 elif command == "workspace":
468 response = {"workspace": manager.current_workspace}
469 elif command == "status":
470 response = {
471 "workspace": manager.current_workspace,
472 "directory": str(manager.get_current_directory()),
473 "running": True,
474 }
475 else:
476 response = {"error": f"Unknown command: {command}"}
477
478 writer.write(json.dumps(response).encode() + b"\n")
479 await writer.drain()
480 except Exception as e:
481 print(f"Error handling client: {e}", file=sys.stderr)
482 finally:
483 writer.close()
484 await writer.wait_closed()
485
486
487async def socket_server() -> None:
488 """Run the Unix socket server for CLI queries"""
489 # Ensure socket directory exists
490 SOCKET_DIR.mkdir(parents=True, exist_ok=True)
491
492 # Remove stale socket if exists
493 if SOCKET_PATH.exists():
494 SOCKET_PATH.unlink()
495
496 server = await asyncio.start_unix_server(handle_client, path=str(SOCKET_PATH))
497
498 # Make socket accessible
499 os.chmod(SOCKET_PATH, 0o600)
500
501 print(f"Socket server listening on {SOCKET_PATH}", file=sys.stderr)
502
503 async with server:
504 await server.serve_forever()
505
506
507async def get_current_workspace() -> str | None:
508 """Get the current workspace from sway"""
509 try:
510 proc = await asyncio.create_subprocess_exec(
511 "swaymsg",
512 "-t",
513 "get_workspaces",
514 stdout=asyncio.subprocess.PIPE,
515 stderr=asyncio.subprocess.PIPE,
516 )
517 stdout, _ = await proc.communicate()
518 workspaces: list[dict[str, Any]] = json.loads(stdout.decode())
519 for ws in workspaces:
520 if ws.get("focused"):
521 name = ws.get("name")
522 return str(name) if name is not None else None
523 except Exception as e:
524 print(f"Failed to get current workspace: {e}", file=sys.stderr)
525 return None
526
527
528async def main() -> None:
529 """Main entry point"""
530 # Ensure user config directory exists (system config is read-only from nix store)
531 CONFIG_DIR.mkdir(parents=True, exist_ok=True)
532
533 # Log config directories
534 if SYSTEM_CONFIG_DIR:
535 print(f"System config dir: {SYSTEM_CONFIG_DIR}", file=sys.stderr)
536 print(f"User config dir: {CONFIG_DIR}", file=sys.stderr)
537
538 # Get initial workspace
539 initial_workspace = await get_current_workspace()
540 if initial_workspace:
541 manager.current_workspace = initial_workspace
542 manager._write_state()
543 print(f"Initial workspace: {initial_workspace}", file=sys.stderr)
544
545 # Handle shutdown gracefully
546 loop = asyncio.get_event_loop()
547
548 def shutdown(sig: signal.Signals) -> None:
549 print(f"Received {sig.name}, shutting down...", file=sys.stderr)
550 for task in asyncio.all_tasks(loop):
551 task.cancel()
552
553 for sig in (signal.SIGTERM, signal.SIGINT):
554 loop.add_signal_handler(sig, shutdown, sig)
555
556 # Run both tasks concurrently
557 try:
558 await asyncio.gather(
559 subscribe_sway(),
560 socket_server(),
561 )
562 except asyncio.CancelledError:
563 print("Shutdown complete", file=sys.stderr)
564 finally:
565 # Cleanup socket
566 if SOCKET_PATH.exists():
567 SOCKET_PATH.unlink()
568
569
570if __name__ == "__main__":
571 asyncio.run(main())