This repository has no description
12 kB
366 lines
1#!/usr/bin/env python3
2"""
3Workspace Manager Daemon
4
5Tracks the current sway workspace and provides:
6- Current workspace name and directory via Unix socket
7- Workspace change events to subscribers
8
9Config: ~/workspaces/<name>.json with at minimum {"directory": "..."}
10System config (from Nix) takes precedence over user config.
11"""
12
13from __future__ import annotations
14
15import asyncio
16import json
17import os
18import signal
19import sys
20from pathlib import Path
21from typing import Any
22
23# XDG runtime directory for socket and state
24RUNTIME_DIR = Path(os.environ.get("XDG_RUNTIME_DIR", f"/run/user/{os.getuid()}"))
25SOCKET_DIR = RUNTIME_DIR / "workspace-manager"
26SOCKET_PATH = SOCKET_DIR / "socket"
27STATE_PATH = SOCKET_DIR / "state.json"
28
29# System config directory (declarative, takes precedence)
30_system_config_dir_env = os.environ.get("WORKSPACE_MANAGER_SYSTEM_CONFIG_DIR")
31SYSTEM_CONFIG_DIR: Path | None = (
32 Path(_system_config_dir_env).expanduser() if _system_config_dir_env else None
33)
34
35# User config directory (fallback)
36_config_dir_env = os.environ.get("WORKSPACE_MANAGER_CONFIG_DIR")
37CONFIG_DIR = (
38 Path(_config_dir_env).expanduser()
39 if _config_dir_env
40 else Path.home() / "workspaces"
41)
42
43_default_dir_env = os.environ.get("WORKSPACE_MANAGER_DEFAULT_DIR")
44DEFAULT_DIRECTORY = (
45 Path(_default_dir_env).expanduser() if _default_dir_env else Path.home()
46)
47
48
49def load_config_file(config_file: Path) -> dict[str, Any] | None:
50 """Load a workspace config file, returning None if it doesn't exist or fails."""
51 if not config_file.exists():
52 return None
53 try:
54 with open(config_file) as f:
55 return json.load(f) # type: ignore[no-any-return]
56 except (json.JSONDecodeError, OSError) as e:
57 print(f"Warning: Failed to load config {config_file}: {e}", file=sys.stderr)
58 return None
59
60
61def get_workspace_directory(name: str) -> Path:
62 """Get the directory for a workspace by merging user and system config."""
63 merged: dict[str, Any] = {}
64
65 user_data = load_config_file(CONFIG_DIR / f"{name}.json")
66 if user_data is not None:
67 merged.update(user_data)
68
69 if SYSTEM_CONFIG_DIR is not None:
70 system_data = load_config_file(SYSTEM_CONFIG_DIR / f"{name}.json")
71 if system_data is not None:
72 merged.update(system_data)
73
74 if "directory" in merged:
75 return Path(str(merged["directory"])).expanduser()
76 return DEFAULT_DIRECTORY
77
78
79class WorkspaceManager:
80 """Tracks current workspace state."""
81
82 def __init__(self) -> None:
83 self.current_workspace: str | None = None
84
85 def get_current_directory(self) -> Path:
86 if self.current_workspace:
87 return get_workspace_directory(self.current_workspace)
88 return DEFAULT_DIRECTORY
89
90 def handle_workspace_change(
91 self, old_workspace: str | None, new_workspace: str
92 ) -> None:
93 self.current_workspace = new_workspace
94 self._write_state()
95
96 def _write_state(self) -> None:
97 state: dict[str, Any] = {
98 "workspace": self.current_workspace,
99 "directory": str(self.get_current_directory()),
100 }
101 try:
102 STATE_PATH.parent.mkdir(parents=True, exist_ok=True)
103 with open(STATE_PATH, "w") as f:
104 json.dump(state, f)
105 except OSError as e:
106 print(f"Failed to write state file: {e}", file=sys.stderr)
107
108
109manager = WorkspaceManager()
110
111# Subscribers: list of (StreamWriter, Event) tuples
112subscribers: list[tuple[asyncio.StreamWriter, asyncio.Event]] = []
113
114
115async def notify_subscribers(old_workspace: str | None, new_workspace: str) -> None:
116 """Notify all subscribers of a workspace change."""
117 event = (
118 json.dumps(
119 {
120 "event": "workspace_change",
121 "old": old_workspace,
122 "new": new_workspace,
123 }
124 ).encode()
125 + b"\n"
126 )
127 for writer, done in subscribers:
128 try:
129 writer.write(event)
130 await writer.drain()
131 except (OSError, ConnectionResetError):
132 done.set()
133
134
135async def subscribe_sway() -> None:
136 """Subscribe to sway IPC workspace events."""
137 print("Starting sway workspace subscription...", file=sys.stderr)
138
139 while True:
140 try:
141 proc = await asyncio.create_subprocess_exec(
142 "swaymsg",
143 "-t",
144 "subscribe",
145 "-m",
146 '["workspace"]',
147 stdout=asyncio.subprocess.PIPE,
148 stderr=asyncio.subprocess.PIPE,
149 )
150
151 if proc.stdout is None:
152 print("Failed to get stdout from swaymsg", file=sys.stderr)
153 await asyncio.sleep(5)
154 continue
155
156 async for line in proc.stdout:
157 try:
158 event: dict[str, Any] = json.loads(line.decode().strip())
159 if event.get("change") == "focus":
160 current: dict[str, Any] = event.get("current", {})
161 old: dict[str, Any] | None = event.get("old")
162 new_name: str | None = current.get("name")
163 old_name: str | None = old.get("name") if old else None
164
165 if new_name:
166 print(
167 f"Workspace change: {old_name} -> {new_name}",
168 file=sys.stderr,
169 )
170 manager.handle_workspace_change(old_name, new_name)
171 await notify_subscribers(old_name, new_name)
172 except json.JSONDecodeError as e:
173 print(f"Failed to parse sway event: {e}", file=sys.stderr)
174
175 await proc.wait()
176 print("swaymsg exited, retrying in 5 seconds...", file=sys.stderr)
177 await asyncio.sleep(5)
178
179 except FileNotFoundError:
180 print("swaymsg not found, retrying in 10 seconds...", file=sys.stderr)
181 await asyncio.sleep(10)
182 except Exception as e:
183 print(
184 f"Error in sway subscription: {e}, retrying in 5 seconds...",
185 file=sys.stderr,
186 )
187 await asyncio.sleep(5)
188
189
190def load_user_config(workspace: str) -> dict[str, Any]:
191 """Load user workspace config, returning empty dict if not found."""
192 config_file = CONFIG_DIR / f"{workspace}.json"
193 return load_config_file(config_file) or {}
194
195
196def save_user_config(workspace: str, updates: dict[str, Any]) -> bool:
197 """Merge updates into user workspace config and write to disk."""
198 config_file = CONFIG_DIR / f"{workspace}.json"
199 try:
200 existing = load_user_config(workspace)
201 existing.update(updates)
202 CONFIG_DIR.mkdir(parents=True, exist_ok=True)
203 with open(config_file, "w") as f:
204 json.dump(existing, f, indent=2)
205 return True
206 except OSError as e:
207 print(f"Failed to save config for {workspace}: {e}", file=sys.stderr)
208 return False
209
210
211async def handle_client(
212 reader: asyncio.StreamReader, writer: asyncio.StreamWriter
213) -> None:
214 """Handle a client connection to the Unix socket."""
215 try:
216 data = await reader.read(65536)
217 raw = data.decode().strip()
218
219 response: dict[str, Any]
220
221 # Try JSON command first, fall back to simple text commands
222 msg: dict[str, Any] | None = None
223 try:
224 msg = json.loads(raw)
225 except json.JSONDecodeError:
226 pass
227
228 if msg is not None:
229 cmd = msg.get("cmd", "")
230 if cmd == "save-tabs":
231 workspace = msg.get("workspace", "")
232 tabs: list[Any] = msg.get("tabs", [])
233 if workspace:
234 ok = save_user_config(workspace, {"tabs": tabs})
235 response = {"ok": ok}
236 else:
237 response = {"ok": False, "error": "No workspace specified"}
238 elif cmd == "get-tabs":
239 workspace = msg.get("workspace", "")
240 if workspace:
241 config = load_user_config(workspace)
242 response = {"tabs": config.get("tabs", [])}
243 else:
244 response = {"tabs": []}
245 else:
246 response = {"error": f"Unknown JSON command: {cmd}"}
247 elif raw == "subscribe":
248 done = asyncio.Event()
249 subscribers.append((writer, done))
250 print("New subscriber connected", file=sys.stderr)
251 try:
252 await done.wait()
253 except asyncio.CancelledError:
254 pass
255 finally:
256 subscribers[:] = [(w, d) for w, d in subscribers if w is not writer]
257 print("Subscriber disconnected", file=sys.stderr)
258 return
259 elif raw == "dir":
260 response = {"directory": str(manager.get_current_directory())}
261 elif raw == "workspace":
262 response = {"workspace": manager.current_workspace}
263 elif raw == "status":
264 response = {
265 "workspace": manager.current_workspace,
266 "directory": str(manager.get_current_directory()),
267 "running": True,
268 }
269 else:
270 response = {"error": f"Unknown command: {raw}"}
271
272 writer.write(json.dumps(response).encode() + b"\n")
273 await writer.drain()
274 except Exception as e:
275 print(f"Error handling client: {e}", file=sys.stderr)
276 finally:
277 writer.close()
278 await writer.wait_closed()
279
280
281async def socket_server() -> None:
282 """Run the Unix socket server."""
283 SOCKET_DIR.mkdir(parents=True, exist_ok=True)
284
285 if SOCKET_PATH.exists():
286 SOCKET_PATH.unlink()
287
288 server = await asyncio.start_unix_server(handle_client, path=str(SOCKET_PATH))
289 os.chmod(SOCKET_PATH, 0o600)
290
291 print(f"Socket server listening on {SOCKET_PATH}", file=sys.stderr)
292
293 async with server:
294 await server.serve_forever()
295
296
297async def get_current_workspace() -> str | None:
298 """Get the current workspace from sway."""
299 try:
300 proc = await asyncio.create_subprocess_exec(
301 "swaymsg",
302 "-t",
303 "get_workspaces",
304 stdout=asyncio.subprocess.PIPE,
305 stderr=asyncio.subprocess.PIPE,
306 )
307 stdout, _ = await proc.communicate()
308 workspaces: list[dict[str, Any]] = json.loads(stdout.decode())
309 for ws in workspaces:
310 if ws.get("focused"):
311 name = ws.get("name")
312 return str(name) if name is not None else None
313 except Exception as e:
314 print(f"Failed to get current workspace: {e}", file=sys.stderr)
315 return None
316
317
318async def main() -> None:
319 """Main entry point."""
320 CONFIG_DIR.mkdir(parents=True, exist_ok=True)
321
322 if SYSTEM_CONFIG_DIR:
323 print(f"System config dir: {SYSTEM_CONFIG_DIR}", file=sys.stderr)
324 print(f"User config dir: {CONFIG_DIR}", file=sys.stderr)
325
326 # Register sway rule to hide Firefox anchor window before Firefox starts
327 try:
328 proc = await asyncio.create_subprocess_exec(
329 "swaymsg",
330 'for_window [title="^workspace-anchor"] move scratchpad',
331 )
332 await proc.wait()
333 print("Registered sway rule for anchor window", file=sys.stderr)
334 except Exception as e:
335 print(f"Failed to register anchor sway rule: {e}", file=sys.stderr)
336
337 initial_workspace = await get_current_workspace()
338 if initial_workspace:
339 manager.current_workspace = initial_workspace
340 manager._write_state()
341 print(f"Initial workspace: {initial_workspace}", file=sys.stderr)
342
343 loop = asyncio.get_event_loop()
344
345 def shutdown(sig: signal.Signals) -> None:
346 print(f"Received {sig.name}, shutting down...", file=sys.stderr)
347 for task in asyncio.all_tasks(loop):
348 task.cancel()
349
350 for sig in (signal.SIGTERM, signal.SIGINT):
351 loop.add_signal_handler(sig, shutdown, sig)
352
353 try:
354 await asyncio.gather(
355 subscribe_sway(),
356 socket_server(),
357 )
358 except asyncio.CancelledError:
359 print("Shutdown complete", file=sys.stderr)
360 finally:
361 if SOCKET_PATH.exists():
362 SOCKET_PATH.unlink()
363
364
365if __name__ == "__main__":
366 asyncio.run(main())