This repository has no description
7.7 kB
219 lines
1#!/usr/bin/env python3
2"""Native messaging host for the Workspace Tabs Firefox extension.
3
4Bridges Firefox extension with workspace-manager daemon and sway.
5Tab persistence goes through the daemon (single writer to workspace files).
6"""
7
8from __future__ import annotations
9
10import json
11import os
12import select
13import socket
14import struct
15import subprocess
16import sys
17from pathlib import Path
18from typing import Any
19
20RUNTIME_DIR = Path(os.environ.get("XDG_RUNTIME_DIR", f"/run/user/{os.getuid()}"))
21SOCKET_PATH = RUNTIME_DIR / "workspace-manager" / "socket"
22
23
24def read_message() -> dict[str, Any] | None:
25 """Read a native messaging message from stdin."""
26 raw = sys.stdin.buffer.read(4)
27 if not raw or len(raw) < 4:
28 return None
29 length = struct.unpack("=I", raw)[0]
30 result: dict[str, Any] = json.loads(sys.stdin.buffer.read(length))
31 return result
32
33
34def write_message(msg: dict[str, Any]) -> None:
35 """Write a native messaging message to stdout."""
36 data = json.dumps(msg).encode()
37 sys.stdout.buffer.write(struct.pack("=I", len(data)) + data)
38 sys.stdout.buffer.flush()
39
40
41def run(cmd: list[str], timeout: float = 2) -> str:
42 """Run a command and return stripped stdout, or empty string on failure."""
43 try:
44 r = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout) # noqa: S603
45 except (OSError, subprocess.TimeoutExpired):
46 return ""
47 return r.stdout.strip() if r.returncode == 0 else ""
48
49
50def daemon_query(msg: dict[str, Any]) -> dict[str, Any]:
51 """Send a JSON command to the workspace-manager daemon and return the response."""
52 try:
53 sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
54 sock.connect(str(SOCKET_PATH))
55 sock.settimeout(5.0)
56 sock.sendall(json.dumps(msg).encode())
57 # Shutdown write side so daemon knows we're done sending
58 sock.shutdown(socket.SHUT_WR)
59 data = b""
60 while True:
61 chunk = sock.recv(65536)
62 if not chunk:
63 break
64 data += chunk
65 sock.close()
66 if data:
67 result: dict[str, Any] = json.loads(data.strip())
68 return result
69 except (OSError, json.JSONDecodeError) as e:
70 print(f"Daemon query failed: {e}", file=sys.stderr)
71 return {}
72
73
74def get_firefox_windows() -> list[dict[str, str]] | None:
75 """Get all Firefox windows and their workspaces from sway tree.
76
77 Returns None if the sway tree could not be queried (so callers can
78 distinguish "no Firefox windows" from "we don't know"). Otherwise
79 returns the list of Firefox windows, possibly empty.
80 """
81 raw = run(["swaymsg", "-t", "get_tree", "-r"], timeout=5)
82 if not raw:
83 return None
84 try:
85 tree: dict[str, Any] = json.loads(raw)
86 except (json.JSONDecodeError, ValueError):
87 return None
88
89 results: list[dict[str, str]] = []
90
91 def traverse(node: dict[str, Any], workspace: str) -> None:
92 node_type = node.get("type", "")
93 if node_type == "workspace":
94 workspace = node.get("name", workspace)
95 app_id = node.get("app_id", "") or ""
96 name = node.get("name", "") or ""
97 # Skip the hidden workspace-manager anchor window — it must never be
98 # mapped to a real workspace, or it gets saved as a phantom tab.
99 if (
100 "firefox" in app_id.lower()
101 and node_type == "con"
102 and not name.startswith("workspace-anchor")
103 ):
104 results.append(
105 {
106 "workspace": workspace,
107 "windowTitle": name,
108 }
109 )
110 for child in node.get("nodes", []) + node.get("floating_nodes", []):
111 traverse(child, workspace)
112
113 traverse(tree, "")
114 return results
115
116
117def handle(msg: dict[str, Any]) -> dict[str, Any]:
118 """Handle a single message from the extension."""
119 mid: int = msg.get("id", 0)
120 cmd = msg.get("cmd", "")
121 if cmd == "current_workspace":
122 return {"id": mid, "workspace": run(["workspace-manager", "workspace"])}
123 if cmd == "list_workspaces":
124 try:
125 ws: list[dict[str, Any]] = json.loads(
126 run(["swaymsg", "-t", "get_workspaces", "-r"])
127 )
128 return {"id": mid, "workspaces": [w["name"] for w in ws]}
129 except (json.JSONDecodeError, KeyError):
130 return {"id": mid, "workspaces": []}
131 if cmd == "goto_workspace":
132 run(["goto-workspace", msg.get("name", "")])
133 return {"id": mid, "ok": True}
134 if cmd == "move_to_workspace":
135 run(["swaymsg", "move", "container", "to", "workspace", msg.get("name", "")])
136 return {"id": mid, "ok": True}
137 if cmd == "map_windows":
138 windows = get_firefox_windows()
139 if windows is None:
140 return {"id": mid, "error": "Failed to query sway tree"}
141 return {"id": mid, "windows": windows}
142 if cmd == "save_tabs":
143 workspace = msg.get("workspace", "")
144 tabs = msg.get("tabs", [])
145 resp = daemon_query({"cmd": "save-tabs", "workspace": workspace, "tabs": tabs})
146 return {"id": mid, "ok": resp.get("ok", False)}
147 if cmd == "get_tabs":
148 workspace = msg.get("workspace", "")
149 resp = daemon_query({"cmd": "get-tabs", "workspace": workspace})
150 return {"id": mid, "tabs": resp.get("tabs", [])}
151 return {"id": mid, "error": f"Unknown: {cmd}"}
152
153
154def connect_subscriber() -> socket.socket | None:
155 """Connect to workspace-manager daemon as a subscriber."""
156 try:
157 sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
158 sock.connect(str(SOCKET_PATH))
159 sock.sendall(b"subscribe")
160 sock.setblocking(False)
161 return sock
162 except OSError as e:
163 print(f"Failed to subscribe to workspace-manager: {e}", file=sys.stderr)
164 return None
165
166
167def main() -> None:
168 """Main loop: handle Firefox messages and workspace-manager events."""
169 sub_sock = connect_subscriber()
170 stdin_fd = sys.stdin.buffer.fileno()
171 sub_buf = b""
172
173 while True:
174 read_fds = [stdin_fd]
175 if sub_sock is not None:
176 read_fds.append(sub_sock.fileno())
177
178 try:
179 readable, _, _ = select.select(read_fds, [], [], 10.0)
180 except (ValueError, OSError):
181 break
182
183 if not readable and sub_sock is None:
184 sub_sock = connect_subscriber()
185 continue
186
187 for fd in readable:
188 if fd == stdin_fd:
189 msg = read_message()
190 if msg is None:
191 return
192 write_message(handle(msg))
193 elif sub_sock is not None and fd == sub_sock.fileno():
194 try:
195 data = sub_sock.recv(4096)
196 if not data:
197 sub_sock.close()
198 sub_sock = connect_subscriber()
199 continue
200 sub_buf += data
201 while b"\n" in sub_buf:
202 line, sub_buf = sub_buf.split(b"\n", 1)
203 event: dict[str, Any] = json.loads(line)
204 write_message(
205 {
206 "id": 0,
207 "event": "workspace_change",
208 "workspace": event.get("new", ""),
209 }
210 )
211 except (OSError, json.JSONDecodeError) as e:
212 print(f"Subscriber error: {e}", file=sys.stderr)
213 if sub_sock is not None:
214 sub_sock.close()
215 sub_sock = None
216
217
218if __name__ == "__main__":
219 main()