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