from __future__ import annotations import argparse from contextlib import contextmanager import curses from dataclasses import dataclass, field, replace import fcntl from functools import wraps import hashlib import json import os from pathlib import Path import shlex import shutil import stat import subprocess import sys import tempfile import threading from typing import Callable, Iterable, Mapping, Sequence from urllib.parse import quote from .aliases import AliasRegistry, AliasRegistryError, default_alias_registry_path from .domain import validate_session_name from .remotes import ( RegistryError, RemoteEntry, RemoteRegistry, default_registry_path, delete_remote_session, ssh_destination_arguments, stop_remote_session, validate_target, ) from .selector import ( CatalogWindow, DeleteSessionOutcome, SelectorError, SelectorResult, SessionSelector, StopSessionOutcome, ) DEFAULT_SOCKET_NAME = "herdr-hub" DEFAULT_TMUX_SESSION = "hub" DEFAULT_SESSION_NAME = "default" BOOTSTRAP_WINDOW_NAME = "__herdr_hub_bootstrap__" WINDOW_KEY_OPTION = "@herdr_hub_key" WINDOW_COMMAND_OPTION = "@herdr_hub_command" WINDOW_ENDPOINT_OPTION = "@herdr_hub_endpoint" WINDOW_SHORT_LABEL_OPTION = "@herdr_hub_short_label" SESSION_DISMISSED_OPTION = "@herdr_hub_dismissed" class HubError(RuntimeError): """A user-facing failure.""" @dataclass(frozen=True) class Endpoint: kind: str session: str = DEFAULT_SESSION_NAME target: str | None = None remote_keybindings: str = "local" remote_transport: str = "native" def __post_init__(self) -> None: if self.kind not in {"local", "remote"}: raise ValueError(f"unsupported endpoint kind: {self.kind}") validate_session_name(self.session) if self.kind == "local": if self.target is not None: raise ValueError("a local endpoint cannot have a remote target") if self.remote_transport != "native": raise ValueError("a local endpoint cannot use a remote transport") return validate_remote_target(self.target) if self.remote_keybindings not in {"local", "server"}: raise ValueError("remote keybindings must be 'local' or 'server'") if self.remote_transport not in {"native", "ssh"}: raise ValueError("remote transport must be 'native' or 'ssh'") @classmethod def local(cls, session: str = DEFAULT_SESSION_NAME) -> Endpoint: return cls(kind="local", session=session) @classmethod def remote( cls, target: str, session: str = DEFAULT_SESSION_NAME, remote_keybindings: str = "local", remote_transport: str = "native", ) -> Endpoint: return cls( kind="remote", target=target, session=session, remote_keybindings=remote_keybindings, remote_transport=remote_transport, ) @property def key(self) -> str: if self.kind == "local": return f"local:{quote(self.session, safe='')}" assert self.target is not None return f"remote:{quote(self.target, safe='')}:{quote(self.session, safe='')}" @property def label(self) -> str: if self.kind == "local": return truncate_label(f"L {self.session}") assert self.target is not None return truncate_label(f"R {display_target(self.target)}/{self.session}") @property def short_label(self) -> str: if self.kind == "local": label = "local" if self.session != DEFAULT_SESSION_NAME: label += f"/{self.session}" return truncate_label(label, 32) assert self.target is not None label = display_target(self.target) if self.session != DEFAULT_SESSION_NAME: label += f"/{self.session}" return truncate_label(label, 32) def command(self, herdr_bin: str, ssh_bin: str = "ssh") -> list[str]: command = [herdr_bin] if self.kind == "remote": assert self.target is not None if self.remote_transport == "ssh": remote_command = ["herdr"] if self.session != DEFAULT_SESSION_NAME: remote_command.extend(["--session", self.session]) return [ ssh_bin, "-tt", "-o", "ServerAliveInterval=15", "-o", "ServerAliveCountMax=3", *ssh_destination_arguments(self.target), shlex.join(remote_command), ] command.extend(["--remote", self.target]) if self.remote_keybindings != "local": command.extend(["--remote-keybindings", self.remote_keybindings]) if self.session != DEFAULT_SESSION_NAME: command.extend(["--session", self.session]) return command @dataclass class DiscoveryReport: endpoints: list[Endpoint] = field(default_factory=list) warnings: list[str] = field(default_factory=list) def extend(self, other: DiscoveryReport) -> None: self.endpoints.extend(other.endpoints) self.warnings.extend(other.warnings) @dataclass(frozen=True) class ManagedWindow: window_id: str key: str name: str dead: bool active: bool endpoint: Endpoint | None = None RunFunction = Callable[..., subprocess.CompletedProcess[str]] def serialized_hub_mutation(method): # type: ignore[no-untyped-def] @wraps(method) def wrapper(self, *args, **kwargs): # type: ignore[no-untyped-def] with self.mutation_lock(): return method(self, *args, **kwargs) return wrapper def validate_remote_target(target: str | None) -> None: validate_target(target or "") def truncate_label(label: str, limit: int = 48) -> str: clean = "".join( character if ord(character) >= 32 and ord(character) != 127 else "?" for character in label ) if len(clean) <= limit: return clean return clean[: limit - 1] + "…" def display_target(target: str) -> str: if target.startswith("ssh://"): return target.removeprefix("ssh://") return target def endpoint_sort_key(endpoint: Endpoint) -> tuple[int, str, str]: if endpoint.kind == "local": session_rank = "" if endpoint.session == DEFAULT_SESSION_NAME else endpoint.session return (0, "", session_rank) return (1, endpoint.target or "", endpoint.session) def dedupe_endpoints(endpoints: Iterable[Endpoint]) -> list[Endpoint]: by_key: dict[str, Endpoint] = {} for endpoint in endpoints: by_key.setdefault(endpoint.key, endpoint) return sorted(by_key.values(), key=endpoint_sort_key) def discover_local_sessions( herdr_bin: str, *, run: RunFunction = subprocess.run, ) -> DiscoveryReport: try: result = run( [herdr_bin, "session", "list", "--json"], capture_output=True, text=True, timeout=5, check=False, ) except (OSError, subprocess.SubprocessError) as error: return DiscoveryReport(warnings=[f"local Herdr session discovery failed: {error}"]) if result.returncode != 0: detail = (result.stderr or result.stdout or "unknown error").strip() return DiscoveryReport( warnings=[f"local Herdr session discovery failed: {detail}"] ) try: payload = json.loads(result.stdout) sessions = payload["sessions"] except (json.JSONDecodeError, KeyError, TypeError) as error: return DiscoveryReport( warnings=[f"local Herdr returned invalid session JSON: {error}"] ) endpoints: list[Endpoint] = [] for item in sessions: if not isinstance(item, dict) or item.get("running") is not True: continue name = item.get("name") if not isinstance(name, str): continue try: endpoints.append(Endpoint.local(name)) except ValueError: continue return DiscoveryReport(endpoints=dedupe_endpoints(endpoints)) def parse_endpoint_argv(argv: Sequence[str]) -> Endpoint | None: if not argv or Path(argv[0]).name != "herdr": return None target: str | None = None session = DEFAULT_SESSION_NAME keybindings = "local" index = 1 while index < len(argv): argument = argv[index] if argument == "--": break if argument == "--remote": if index + 1 >= len(argv): return None target = argv[index + 1] index += 2 continue if argument.startswith("--remote="): target = argument.split("=", 1)[1] index += 1 continue if argument == "--session": if index + 1 >= len(argv): return None session = argv[index + 1] index += 2 continue if argument.startswith("--session="): session = argument.split("=", 1)[1] index += 1 continue if argument == "--remote-keybindings": if index + 1 >= len(argv): return None keybindings = argv[index + 1] index += 2 continue if argument.startswith("--remote-keybindings="): keybindings = argument.split("=", 1)[1] index += 1 continue index += 1 try: return ( Endpoint.remote(target, session, keybindings) if target is not None else Endpoint.local(session) ) except ValueError: return None def parse_remote_argv(argv: Sequence[str]) -> Endpoint | None: endpoint = parse_endpoint_argv(argv) return endpoint if endpoint is not None and endpoint.kind == "remote" else None def serialize_endpoint(endpoint: Endpoint) -> str: return json.dumps( { "version": 1, "kind": endpoint.kind, "session": endpoint.session, "target": endpoint.target, "remote_keybindings": endpoint.remote_keybindings, "remote_transport": endpoint.remote_transport, }, ensure_ascii=False, separators=(",", ":"), ) def deserialize_endpoint(metadata: str, command: str = "") -> Endpoint | None: if metadata: try: payload = json.loads(metadata) if not isinstance(payload, dict) or payload.get("version") != 1: raise ValueError kind = payload.get("kind") session = payload.get("session") target = payload.get("target") keybindings = payload.get("remote_keybindings", "local") transport = payload.get("remote_transport", "native") if not isinstance(kind, str) or not isinstance(session, str): raise ValueError if not isinstance(keybindings, str): raise ValueError if not isinstance(transport, str): raise ValueError if kind == "local": return Endpoint.local(session) if kind == "remote" and isinstance(target, str): return Endpoint.remote(target, session, keybindings, transport) raise ValueError except (json.JSONDecodeError, TypeError, ValueError): pass if not command: return None try: arguments = shlex.split(command) except ValueError: return None if arguments and arguments[0] == "exec": arguments = arguments[1:] return parse_endpoint_argv(arguments) def decode_cmdline(raw: bytes) -> list[str]: return [ os.fsdecode(argument) for argument in raw.rstrip(b"\0").split(b"\0") if argument ] def discover_remote_processes( proc_root: Path = Path("/proc"), *, user_id: int | None = None, ) -> DiscoveryReport: if user_id is None: user_id = os.getuid() if not proc_root.is_dir(): return DiscoveryReport( warnings=[f"remote discovery unavailable: {proc_root} is not mounted"] ) endpoints: list[Endpoint] = [] unreadable = 0 try: entries = sorted( (entry for entry in proc_root.iterdir() if entry.name.isdigit()), key=lambda entry: int(entry.name), ) except OSError as error: return DiscoveryReport(warnings=[f"remote process discovery failed: {error}"]) for entry in entries: cmdline_path = entry / "cmdline" try: if cmdline_path.stat().st_uid != user_id: continue raw = cmdline_path.read_bytes() except (FileNotFoundError, ProcessLookupError): continue except (PermissionError, OSError): unreadable += 1 continue endpoint = parse_remote_argv(decode_cmdline(raw)) if endpoint is not None: endpoints.append(endpoint) warnings: list[str] = [] if unreadable: warnings.append(f"could not inspect {unreadable} same-user process command line(s)") return DiscoveryReport(endpoints=dedupe_endpoints(endpoints), warnings=warnings) def discover_endpoints(herdr_bin: str, proc_root: Path = Path("/proc")) -> DiscoveryReport: report = discover_local_sessions(herdr_bin) report.extend(discover_remote_processes(proc_root)) report.endpoints = dedupe_endpoints(report.endpoints) return report class TmuxHub: def __init__( self, *, tmux_bin: str, ssh_bin: str = "ssh", socket_name: str = DEFAULT_SOCKET_NAME, session_name: str = DEFAULT_TMUX_SESSION, config_file: str | None = None, script_path: str | None = None, python_bin: str | None = None, aliases: Mapping[str, str] | None = None, run: RunFunction = subprocess.run, ) -> None: self.tmux_bin = tmux_bin self.ssh_bin = ssh_bin self.socket_name = socket_name self.session_name = session_name self.config_file = config_file self.script_path = script_path or str(Path(sys.argv[0]).absolute()) self.python_bin = python_bin or sys.executable self.aliases = dict(aliases or {}) self._run = run self._local_mutation_lock = threading.RLock() self._mutation_depth = 0 @property def mutation_lock_path(self) -> Path: identity = f"{os.getuid()}\0{self.socket_name}\0{self.session_name}" digest = hashlib.sha256(identity.encode("utf-8")).hexdigest()[:20] roots = [] if configured_runtime := os.environ.get("XDG_RUNTIME_DIR"): roots.append(Path(configured_runtime).expanduser()) temporary_root = Path(tempfile.gettempdir()) if temporary_root not in roots: roots.append(temporary_root) failures: list[str] = [] for runtime_root in roots: lock_directory = runtime_root / f"herdr-hub-{os.getuid()}" try: lock_directory.mkdir(mode=0o700, parents=False, exist_ok=True) metadata = lock_directory.lstat() if ( not stat.S_ISDIR(metadata.st_mode) or metadata.st_uid != os.getuid() ): failures.append(f"unsafe directory {lock_directory}") continue os.chmod(lock_directory, 0o700) return lock_directory / f"{digest}.lock" except OSError as error: failures.append(f"{lock_directory}: {error}") raise HubError("could not prepare Hub lock directory: " + "; ".join(failures)) @contextmanager def mutation_lock(self): # type: ignore[no-untyped-def] """Serialize state-changing Hub commands across CLI processes.""" with self._local_mutation_lock: if self._mutation_depth: self._mutation_depth += 1 try: yield finally: self._mutation_depth -= 1 return descriptor = os.open( self.mutation_lock_path, os.O_CREAT | os.O_RDWR, 0o600, ) try: os.fchmod(descriptor, 0o600) fcntl.flock(descriptor, fcntl.LOCK_EX) self._mutation_depth = 1 try: yield finally: self._mutation_depth = 0 fcntl.flock(descriptor, fcntl.LOCK_UN) finally: os.close(descriptor) def base_command(self) -> list[str]: command = [self.tmux_bin, "-L", self.socket_name] if self.config_file is not None: command.extend(["-f", self.config_file]) return command def run( self, arguments: Sequence[str], *, check: bool = True, capture_output: bool = True, text: bool = True, ) -> subprocess.CompletedProcess[str]: try: result = self._run( [*self.base_command(), *arguments], capture_output=capture_output, text=text, check=False, ) except OSError as error: raise HubError(f"could not run tmux: {error}") from error if check and result.returncode != 0: detail = (result.stderr or result.stdout or "tmux command failed").strip() raise HubError(detail) return result def exists(self) -> bool: result = self.run( ["has-session", "-t", self.session_name], check=False, ) return result.returncode == 0 @serialized_hub_mutation def ensure_server(self) -> bool: if self.exists(): self.configure() return False bootstrap = shlex.join(["sleep", "infinity"]) self.run( [ "new-session", "-d", "-s", self.session_name, "-n", BOOTSTRAP_WINDOW_NAME, bootstrap, ] ) self.configure() return True def configure(self) -> None: option_commands = [ ["set-option", "-s", "set-clipboard", "on"], ["set-option", "-g", "prefix", "C-a"], ["set-option", "-g", "status", "on"], ["set-option", "-g", "status-position", "top"], ["set-option", "-g", "status-interval", "2"], ["set-option", "-g", "mouse", "on"], ["set-option", "-g", "focus-events", "on"], ["set-option", "-g", "status-left-length", "18"], ["set-option", "-g", "status-right-length", "42"], [ "set-option", "-g", "status-left", "#{?#{<:#{client_width},50}," "#[bold#,fg=colour81] S #[default]," "#[bold#,fg=colour81] Sessions #[default]}", ], [ "set-option", "-g", "status-right", "#{?#{<:#{client_width},80},," "#[dim]#{?pane_dead,连接已断开 · C-a R 重连," "C-a g choose · q close · r refresh}#[default]}", ], ["set-option", "-g", "status-style", "bg=colour235,fg=colour250"], ["set-window-option", "-g", "window-status-separator", ""], [ "set-window-option", "-g", "window-status-format", "#{?#{<:#{client_width},80}," " #{?pane_dead,×,○} #{=12:@herdr_hub_short_label} ," " #{?pane_dead,× #W [断开],○ #W} }", ], [ "set-window-option", "-g", "window-status-current-format", "#[bold,bg=colour31,fg=colour231]" "#{?#{<:#{client_width},80}," " #{?pane_dead,×,●} #{=12:@herdr_hub_short_label} ," " #{?pane_dead,× #W [断开],● #W} }" "#[default]", ], ["set-window-option", "-g", "automatic-rename", "off"], ["set-window-option", "-g", "allow-rename", "off"], ["set-window-option", "-g", "remain-on-exit", "on"], [ "set-window-option", "-g", "remain-on-exit-format", "#[bold,fg=colour196]连接已断开#[default]" "（退出状态 #{pane_dead_status}） · 按 Ctrl+A R 重连", ], ] for command in option_commands: self.run(command) refresh = shlex.join([self.python_bin, self.script_path, "refresh"]) choose = shlex.join([self.python_bin, self.script_path, "choose"]) reconnect = shlex.join([self.python_bin, self.script_path, "reconnect"]) close = shlex.join( [self.python_bin, self.script_path, "close", "--window", "#{window_id}"] ) confirm_close = shlex.join(["run-shell", "-b", close]) # Remove tmux's default/legacy '&' binding when reconfiguring both new # and already-running Hub servers. self.run(["unbind-key", "-T", "prefix", "&"], check=False) self.run(["unbind-key", "-T", "prefix", "Q"], check=False) binding_commands = [ ["bind-key", "-T", "prefix", "C-a", "send-prefix"], ["bind-key", "-T", "prefix", "r", "run-shell", "-b", refresh], [ "bind-key", "-T", "prefix", "g", "display-popup", "-EE", "-w", "90%", "-h", "90%", "-T", "Sessions", choose, ], [ "bind-key", "-T", "root", "MouseDown1StatusLeft", "display-popup", "-EE", "-w", "90%", "-h", "90%", "-T", "Sessions", choose, ], ["bind-key", "-T", "prefix", "R", "run-shell", "-b", reconnect], [ "bind-key", "-T", "prefix", "q", "confirm-before", "-p", "关闭当前 Session Tab？ (y/n)", confirm_close, ], ] for command in binding_commands: self.run(command) self.refresh_labels() def managed_windows(self) -> list[ManagedWindow]: if not self.exists(): return [] line_format = ( f"#{{window_id}}\t#{{{WINDOW_KEY_OPTION}}}\t#{{window_name}}\t" f"#{{pane_dead}}\t#{{window_active}}\t#{{{WINDOW_ENDPOINT_OPTION}}}\t" f"#{{{WINDOW_COMMAND_OPTION}}}" ) result = self.run( ["list-windows", "-t", self.session_name, "-F", line_format] ) windows: list[ManagedWindow] = [] for line in result.stdout.splitlines(): fields = line.split("\t", 6) if len(fields) != 7 or not fields[1]: continue windows.append( ManagedWindow( window_id=fields[0], key=fields[1], name=fields[2], dead=fields[3] == "1", active=fields[4] == "1", endpoint=deserialize_endpoint(fields[5], fields[6]), ) ) return windows def window_for_key(self, key: str) -> ManagedWindow | None: return next((window for window in self.managed_windows() if window.key == key), None) def refresh_labels(self) -> None: for window in self.managed_windows(): self._apply_window_label(window) def _apply_window_label(self, window: ManagedWindow) -> None: alias = self.aliases.get(window.key) display_label = truncate_label(alias) if alias else window.name short_label = truncate_label(alias, 32) if alias else truncate_label(window.name, 32) if window.endpoint is not None: display_label = truncate_label(alias) if alias else window.endpoint.label short_label = ( truncate_label(alias, 32) if alias else window.endpoint.short_label ) if window.name != display_label: self.run(["rename-window", "-t", window.window_id, display_label]) self.run( [ "set-option", "-w", "-t", window.window_id, WINDOW_SHORT_LABEL_OPTION, short_label, ] ) @serialized_hub_mutation def apply_alias(self, key: str, label: str) -> ManagedWindow | None: if label: self.aliases[key] = label else: self.aliases.pop(key, None) window = self.window_for_key(key) if window is not None: self._apply_window_label(window) return self.window_for_key(key) def dismissed_keys(self) -> set[str]: result = self.run( [ "show-options", "-v", "-t", self.session_name, SESSION_DISMISSED_OPTION, ], check=False, ) if result.returncode != 0 or not result.stdout.strip(): return set() try: payload = json.loads(result.stdout) except json.JSONDecodeError: return set() if not isinstance(payload, list): return set() return {key for key in payload if isinstance(key, str) and key} def set_dismissed_keys(self, keys: Iterable[str]) -> None: value = json.dumps(sorted(set(keys)), ensure_ascii=False, separators=(",", ":")) self.run( [ "set-option", "-t", self.session_name, SESSION_DISMISSED_OPTION, value, ] ) def restore_endpoint(self, key: str) -> None: dismissed = self.dismissed_keys() if key not in dismissed: return dismissed.remove(key) self.set_dismissed_keys(dismissed) @serialized_hub_mutation def create_window(self, endpoint: Endpoint, herdr_bin: str) -> ManagedWindow: existing = self.window_for_key(endpoint.key) if existing is not None: return existing was_dismissed = endpoint.key in self.dismissed_keys() self.restore_endpoint(endpoint.key) command = "exec " + shlex.join(endpoint.command(herdr_bin, self.ssh_bin)) alias = self.aliases.get(endpoint.key) display_label = truncate_label(alias) if alias else endpoint.label short_label = truncate_label(alias, 32) if alias else endpoint.short_label result = self.run( [ "new-window", "-d", "-P", "-F", "#{window_id}", "-t", f"{self.session_name}:", "-n", display_label, command, ] ) window_id = result.stdout.strip() if not window_id: if was_dismissed: self.set_dismissed_keys({*self.dismissed_keys(), endpoint.key}) raise HubError("tmux did not return the new window id") try: for option, value in [ (WINDOW_KEY_OPTION, endpoint.key), (WINDOW_COMMAND_OPTION, command), (WINDOW_ENDPOINT_OPTION, serialize_endpoint(endpoint)), (WINDOW_SHORT_LABEL_OPTION, short_label), ]: self.run(["set-option", "-w", "-t", window_id, option, value]) self.run(["set-window-option", "-t", window_id, "remain-on-exit", "on"]) except HubError: self.run(["kill-window", "-t", window_id], check=False) if was_dismissed: self.set_dismissed_keys({*self.dismissed_keys(), endpoint.key}) raise return ManagedWindow( window_id=window_id, key=endpoint.key, name=display_label, dead=False, active=False, endpoint=endpoint, ) @serialized_hub_mutation def add_missing(self, endpoints: Sequence[Endpoint], herdr_bin: str) -> list[ManagedWindow]: existing_keys = {window.key for window in self.managed_windows()} dismissed_keys = self.dismissed_keys() added: list[ManagedWindow] = [] for endpoint in endpoints: if endpoint.key in existing_keys or endpoint.key in dismissed_keys: continue window = self.create_window(endpoint, herdr_bin) existing_keys.add(endpoint.key) added.append(window) if self.managed_windows(): self.remove_bootstrap() return added @serialized_hub_mutation def remove_bootstrap(self) -> None: result = self.run( [ "list-windows", "-t", self.session_name, "-F", "#{window_id}\t#{window_name}", ], check=False, ) if result.returncode != 0: return for line in result.stdout.splitlines(): fields = line.split("\t", 1) if len(fields) == 2 and fields[1] == BOOTSTRAP_WINDOW_NAME: self.run(["kill-window", "-t", fields[0]], check=False) def select_window(self, window_id: str) -> None: self.run(["select-window", "-t", window_id]) @serialized_hub_mutation def dismiss_window(self, window_id: str | None = None) -> ManagedWindow: window_id = window_id or self.current_window_id() windows = self.managed_windows() window = next((item for item in windows if item.window_id == window_id), None) if window is None or window.endpoint is None: raise HubError("the selected window is not managed by herdr-hub") if len(windows) <= 1: raise HubError("至少保留一个 Session Tab；可用 herdr-hub stop 退出整个 Hub") previous = self.dismissed_keys() self.set_dismissed_keys({*previous, window.key}) result = self.run(["kill-window", "-t", window.window_id], check=False) if result.returncode != 0: self.set_dismissed_keys(previous) detail = (result.stderr or result.stdout or "tmux command failed").strip() raise HubError(detail) return window def select_primary(self, endpoints: Sequence[Endpoint]) -> None: preferred = next( (endpoint for endpoint in endpoints if endpoint.key == Endpoint.local().key), endpoints[0] if endpoints else None, ) if preferred is None: return window = self.window_for_key(preferred.key) if window is not None: self.select_window(window.window_id) def current_window_id(self) -> str: arguments = ["display-message", "-p"] if pane_id := os.environ.get("TMUX_PANE"): arguments.extend(["-t", pane_id]) arguments.append("#{window_id}") result = self.run(arguments) window_id = result.stdout.strip() if not window_id: raise HubError("could not identify the current hub window") return window_id @serialized_hub_mutation def reconnect(self, window_id: str | None = None) -> None: window_id = window_id or self.current_window_id() command_result = self.run( ["show-options", "-w", "-v", "-t", window_id, WINDOW_COMMAND_OPTION], check=False, ) command = command_result.stdout.strip() if command_result.returncode != 0 or not command: raise HubError("the selected window is not managed by herdr-hub") self.run(["respawn-window", "-k", "-t", window_id, command]) @serialized_hub_mutation def replace_endpoint( self, window: ManagedWindow, endpoint: Endpoint, herdr_bin: str, ) -> ManagedWindow: command = "exec " + shlex.join(endpoint.command(herdr_bin, self.ssh_bin)) alias = self.aliases.get(endpoint.key) short_label = truncate_label(alias, 32) if alias else endpoint.short_label replacements = [ (WINDOW_COMMAND_OPTION, command), (WINDOW_ENDPOINT_OPTION, serialize_endpoint(endpoint)), (WINDOW_SHORT_LABEL_OPTION, short_label), ] previous: dict[str, str] = {} for option, _ in replacements: result = self.run( ["show-options", "-w", "-v", "-t", window.window_id, option], check=False, ) if result.returncode == 0: previous[option] = result.stdout.rstrip("\n") try: for option, value in replacements: self.run( ["set-option", "-w", "-t", window.window_id, option, value] ) self.run(["respawn-window", "-k", "-t", window.window_id, command]) except HubError: for option, value in previous.items(): self.run( ["set-option", "-w", "-t", window.window_id, option, value], check=False, ) raise replaced = self.window_for_key(endpoint.key) if replaced is None: raise HubError("the selected Session window disappeared during reconnect") self._apply_window_label(replaced) return self.window_for_key(endpoint.key) or replaced @serialized_hub_mutation def activate_endpoint(self, endpoint: Endpoint, herdr_bin: str) -> ManagedWindow: window = self.window_for_key(endpoint.key) if window is None: window = self.create_window(endpoint, herdr_bin) self.remove_bootstrap() elif window.endpoint is not None and ( window.endpoint.remote_transport != endpoint.remote_transport or window.endpoint.remote_keybindings != endpoint.remote_keybindings ): window = self.replace_endpoint(window, endpoint, herdr_bin) elif window.dead: self.reconnect(window.window_id) self.select_window(window.window_id) return window def attach(self) -> None: if inside_any_tmux(): if inside_socket(self.socket_name): return raise HubError( "already inside another tmux server; detach from it before attaching to herdr-hub" ) result = self.run( ["attach-session", "-t", self.session_name], check=False, capture_output=False, text=True, ) if result.returncode != 0: raise HubError("could not attach to the herdr-hub tmux session") def display_message(self, message: str) -> None: self.run(["display-message", message], check=False) @serialized_hub_mutation def stop(self) -> None: if not self.exists(): return self.run(["kill-server"]) def inside_any_tmux() -> bool: return bool(os.environ.get("TMUX")) def inside_socket(socket_name: str) -> bool: value = os.environ.get("TMUX", "") socket_path = value.split(",", 1)[0] return bool(socket_path) and Path(socket_path).name == socket_name def resolve_binary(name: str, override: str | None = None) -> str: if override: return override resolved = shutil.which(name) if resolved is None: raise HubError(f"required executable not found on PATH: {name}") return resolved def invocation_script_path() -> str: invoked = Path(sys.argv[0]).absolute() if invoked.name == "__main__.py": repository_launcher = Path(__file__).resolve().parents[1] / "bin" / "herdr-hub" return str(repository_launcher) return str(invoked) def make_hub( arguments: argparse.Namespace, *, load_aliases: bool = True, ) -> TmuxHub: tmux_bin = resolve_binary("tmux", os.environ.get("HERDR_HUB_TMUX_BIN")) ssh_bin = os.environ.get("HERDR_HUB_SSH_BIN") or shutil.which("ssh") or "ssh" aliases = make_alias_registry().load() if load_aliases else {} return TmuxHub( tmux_bin=tmux_bin, ssh_bin=ssh_bin, socket_name=os.environ.get("HERDR_HUB_SOCKET", DEFAULT_SOCKET_NAME), session_name=os.environ.get("HERDR_HUB_SESSION", DEFAULT_TMUX_SESSION), config_file=os.environ.get("HERDR_HUB_TMUX_CONFIG"), script_path=invocation_script_path(), aliases=aliases, ) def resolve_herdr() -> str: return resolve_binary("herdr", os.environ.get("HERDR_HUB_HERDR_BIN")) def resolve_ssh() -> str: return resolve_binary("ssh", os.environ.get("HERDR_HUB_SSH_BIN")) def configured_proc_root() -> Path: return Path(os.environ.get("HERDR_HUB_PROC_ROOT", "/proc")) def make_registry() -> RemoteRegistry: return RemoteRegistry(default_registry_path()) def make_alias_registry() -> AliasRegistry: return AliasRegistry(default_alias_registry_path()) def remember_remote_endpoints(endpoints: Iterable[Endpoint]) -> str | None: entries = [ RemoteEntry( endpoint.target, endpoint.remote_keybindings, endpoint.remote_transport, ) for endpoint in endpoints if endpoint.kind == "remote" and endpoint.target is not None ] if not entries: return None try: make_registry().remember_many(entries) except RegistryError as error: return f"remote registry: {error}" return None def catalog_windows(windows: Iterable[ManagedWindow]) -> list[CatalogWindow]: return [ CatalogWindow( window_id=window.window_id, key=window.key, label=window.name, dead=window.dead, active=window.active, kind=window.endpoint.kind if window.endpoint is not None else None, target=window.endpoint.target if window.endpoint is not None else None, session=window.endpoint.session if window.endpoint is not None else None, default_label=( window.endpoint.label if window.endpoint is not None else window.name ), remote_transport=( window.endpoint.remote_transport if window.endpoint is not None else "native" ), ) for window in windows ] def command_start(arguments: argparse.Namespace) -> int: herdr_bin = resolve_herdr() report = discover_endpoints(herdr_bin, configured_proc_root()) if warning := remember_remote_endpoints(report.endpoints): report.warnings.append(warning) hub = make_hub(arguments) created = hub.ensure_server() added = hub.add_missing(report.endpoints, herdr_bin) windows = hub.managed_windows() if not windows: hub.stop() detail = "; ".join(report.warnings) suffix = f" ({detail})" if detail else "" raise HubError( "no running Herdr sessions or live --remote clients were found; " f"use 'herdr-hub open' to create one{suffix}" ) if created: hub.select_primary(report.endpoints) if added and not created: hub.display_message(f"herdr-hub: added {len(added)} session(s)") hub.attach() return 0 def command_refresh(arguments: argparse.Namespace) -> int: herdr_bin = resolve_herdr() report = discover_endpoints(herdr_bin, configured_proc_root()) if warning := remember_remote_endpoints(report.endpoints): report.warnings.append(warning) hub = make_hub(arguments) hub.ensure_server() added = hub.add_missing(report.endpoints, herdr_bin) message = f"herdr-hub: added {len(added)} session(s)" if report.warnings: message += f"; {report.warnings[0]}" has_windows = bool(hub.managed_windows()) if not has_windows: hub.stop() if inside_socket(hub.socket_name) and has_windows: hub.display_message(message) else: print(message) return 0 def command_open(arguments: argparse.Namespace) -> int: herdr_bin = resolve_herdr() try: endpoint = ( Endpoint.remote( arguments.remote, arguments.session, arguments.remote_keybindings, arguments.transport, ) if arguments.remote is not None else Endpoint.local(arguments.session) ) except ValueError as error: raise HubError(str(error)) from error if endpoint.kind == "remote": assert endpoint.target is not None try: make_registry().upsert( RemoteEntry( endpoint.target, endpoint.remote_keybindings, endpoint.remote_transport, ) ) except RegistryError as error: print(f"herdr-hub: warning: remote registry: {error}", file=sys.stderr) hub = make_hub(arguments) hub.ensure_server() hub.activate_endpoint(endpoint, herdr_bin) hub.attach() return 0 def command_choose(arguments: argparse.Namespace) -> int: hub = make_hub(arguments) if not hub.exists(): raise HubError("herdr-hub is not running") windows = hub.managed_windows() local_report = discover_local_sessions(resolve_herdr()) registry_warning = remember_remote_endpoints( window.endpoint for window in windows if window.endpoint is not None ) alias_registry = make_alias_registry() remote_registry = make_registry() ssh_bin = resolve_ssh() def close_window(window_id: str) -> list[CatalogWindow]: hub.dismiss_window(window_id) return catalog_windows(hub.managed_windows()) def rename_session(endpoint_key: str, alias: str) -> None: alias_registry.set(endpoint_key, alias) if hub.apply_alias(endpoint_key, alias) is None: raise HubError("the selected Session is no longer open") def delete_session( target: str, session: str, window_id: str | None, ) -> DeleteSessionOutcome: delete_remote_session(ssh_bin, RemoteEntry(target), session) endpoint_key = Endpoint.remote(target, session).key warnings: list[str] = [] tab_closed = False if window_id is not None: try: hub.dismiss_window(window_id) tab_closed = True except HubError as error: warnings.append(f"Tab 未关闭：{error}") try: alias_registry.set(endpoint_key, "") except AliasRegistryError as error: warnings.append(f"显示名称清理失败：{error}") hub.aliases.pop(endpoint_key, None) try: hub.restore_endpoint(endpoint_key) except HubError as error: warnings.append(f"隐藏记录清理失败：{error}") return DeleteSessionOutcome( warning="；".join(warnings) or None, windows=tuple(catalog_windows(hub.managed_windows())), tab_closed=tab_closed, ) def stop_session( window_id: str, target: str, session: str, ) -> StopSessionOutcome: stop_remote_session(ssh_bin, RemoteEntry(target), session) warning: str | None = None tab_closed = False try: hub.dismiss_window(window_id) tab_closed = True except HubError as error: warning = f"Session 已停止，但 Tab 未关闭：{error}" updated = catalog_windows(hub.managed_windows()) if not tab_closed: updated = [ replace(window, dead=True) if window.window_id == window_id else window for window in updated ] return StopSessionOutcome(tuple(updated), tab_closed, warning) try: selector = SessionSelector( catalog_windows(windows), remote_registry, ssh_bin, local_sessions=(endpoint.session for endpoint in local_report.endpoints), aliases=hub.aliases, rename_handler=rename_session, close_window_handler=close_window, delete_session_handler=delete_session, stop_session_handler=stop_session, ) startup_warnings = [*local_report.warnings] if registry_warning: startup_warnings.append(registry_warning) if startup_warnings: selector.message = "; ".join(startup_warnings) result = selector.run() except (RegistryError, SelectorError) as error: raise HubError(str(error)) from error if result is None: return 0 if result.kind == "rename": assert result.endpoint_key is not None and result.alias is not None rename_session(result.endpoint_key, result.alias) window = hub.window_for_key(result.endpoint_key) assert window is not None action = "renamed" if result.alias else "restored the default name for" hub.display_message(f"herdr-hub: {action} {window.name}") return 0 if result.kind == "close": assert result.window_id is not None hub.dismiss_window(result.window_id) return 0 if result.kind == "window": assert result.window_id is not None if result.dead: hub.reconnect(result.window_id) hub.select_window(result.window_id) return 0 if result.kind == "local": assert result.session is not None hub.activate_endpoint(Endpoint.local(result.session), resolve_herdr()) return 0 assert result.target is not None and result.session is not None endpoint = Endpoint.remote( result.target, result.session, result.remote_keybindings, result.remote_transport, ) try: remote_registry.upsert( RemoteEntry( endpoint.target, endpoint.remote_keybindings, endpoint.remote_transport, ) ) except RegistryError as error: hub.display_message(f"herdr-hub: remote registry: {error}") hub.activate_endpoint(endpoint, resolve_herdr()) return 0 def command_remote(arguments: argparse.Namespace) -> int: registry = make_registry() action = arguments.remote_command if action == "list": entries = registry.load() if not entries: print("no registered remote targets") return 0 for entry in entries: print( f"{entry.target}\t{entry.remote_keybindings}\t{entry.remote_transport}" ) return 0 if action == "add": entry = RemoteEntry( arguments.target, arguments.remote_keybindings, arguments.transport, ) changed = registry.upsert(entry) print(f"{'registered' if changed else 'already registered'} {entry.target}") return 0 if action == "remove": removed = registry.remove(arguments.target) print(f"{'removed' if removed else 'not registered'} {arguments.target}") return 0 raise HubError("use 'herdr-hub remote list|add|remove'") def command_reconnect(arguments: argparse.Namespace) -> int: hub = make_hub(arguments, load_aliases=False) if not hub.exists(): raise HubError("herdr-hub is not running") hub.reconnect(arguments.window) if inside_socket(hub.socket_name): hub.display_message("herdr-hub: reconnecting current session") return 0 def command_close(arguments: argparse.Namespace) -> int: hub = make_hub(arguments, load_aliases=False) if not hub.exists(): raise HubError("herdr-hub is not running") try: window = hub.dismiss_window(arguments.window) except HubError as error: if inside_socket(hub.socket_name): hub.display_message(f"herdr-hub: {error}") return 1 raise message = f"herdr-hub: closed {window.name}; Herdr Session is still running" if inside_socket(hub.socket_name): hub.display_message(message) else: print(message) return 0 def command_stop(arguments: argparse.Namespace) -> int: hub = make_hub(arguments, load_aliases=False) if hub.exists(): hub.stop() print("stopped herdr-hub; Herdr servers were left running") else: print("herdr-hub is not running") return 0 def command_doctor(arguments: argparse.Namespace) -> int: failures = 0 print(f"python: {sys.version.split()[0]} ({sys.executable})") print(f"curses: {getattr(curses, 'version', 'available')}") for binary in ["tmux", "herdr", "ssh"]: override = os.environ.get(f"HERDR_HUB_{binary.upper()}_BIN") path = shutil.which(binary) if override is None else override if path is None: print(f"{binary}: missing") failures += 1 continue version_args = [path, "-V"] if binary in {"tmux", "ssh"} else [path, "--version"] result = subprocess.run( version_args, capture_output=True, text=True, check=False, ) version = (result.stdout or result.stderr).strip() print(f"{binary}: {version or path}") registry = make_registry() try: registered = registry.load() print(f"remotes: {len(registered)} registered ({registry.path})") except RegistryError as error: print(f"remotes: invalid ({error})") failures += 1 alias_registry = make_alias_registry() try: aliases = alias_registry.load() print(f"aliases: {len(aliases)} configured ({alias_registry.path})") except AliasRegistryError as error: print(f"aliases: invalid ({error})") failures += 1 proc_root = configured_proc_root() print(f"proc: {proc_root} ({'available' if proc_root.is_dir() else 'missing'})") try: herdr_bin = resolve_herdr() except HubError as error: print(f"discovery: unavailable ({error})") return 1 report = discover_endpoints(herdr_bin, proc_root) print(f"discovered: {len(report.endpoints)} endpoint(s)") for endpoint in report.endpoints: print(f" {endpoint.label} [{endpoint.key}]") for warning in report.warnings: print(f"warning: {warning}") return failures def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser( prog="herdr-hub", description="Switch local and live --remote Herdr sessions through a tmux top bar.", ) subparsers = parser.add_subparsers(dest="command") subparsers.add_parser("start", help="scan, create and attach to the hub") subparsers.add_parser("refresh", help="add newly discovered sessions") open_parser = subparsers.add_parser("open", help="open a session directly in the hub") open_parser.add_argument("--remote", metavar="TARGET") open_parser.add_argument("--session", default=DEFAULT_SESSION_NAME) open_parser.add_argument( "--remote-keybindings", choices=["local", "server"], default="local", ) open_parser.add_argument( "--transport", choices=["native", "ssh"], default="native", help="remote attach transport; Windows targets require ssh", ) subparsers.add_parser("choose", help="browse and open local or remote sessions") remote_parser = subparsers.add_parser("remote", help="manage remembered remote targets") remote_subparsers = remote_parser.add_subparsers(dest="remote_command") remote_subparsers.add_parser("list", help="list remembered remote targets") remote_add_parser = remote_subparsers.add_parser("add", help="remember a remote target") remote_add_parser.add_argument("target", metavar="TARGET") remote_add_parser.add_argument( "--remote-keybindings", choices=["local", "server"], default="local", ) remote_add_parser.add_argument( "--transport", choices=["native", "ssh"], default="native", help="preferred attach transport; use ssh for Windows targets", ) remote_remove_parser = remote_subparsers.add_parser( "remove", help="forget a remote target" ) remote_remove_parser.add_argument("target", metavar="TARGET") reconnect_parser = subparsers.add_parser( "reconnect", help="respawn a dead managed window" ) reconnect_parser.add_argument("--window") close_parser = subparsers.add_parser( "close", help="hide a managed window without stopping its Herdr session" ) close_parser.add_argument("--window") subparsers.add_parser("doctor", help="check dependencies and discovery") subparsers.add_parser("stop", help="stop the hub clients, leaving Herdr servers running") return parser def main(argv: Sequence[str] | None = None) -> int: parser = build_parser() arguments = parser.parse_args(argv) command = arguments.command or "start" handlers = { "start": command_start, "refresh": command_refresh, "open": command_open, "choose": command_choose, "remote": command_remote, "reconnect": command_reconnect, "close": command_close, "doctor": command_doctor, "stop": command_stop, } try: return handlers[command](arguments) except (AliasRegistryError, HubError, RegistryError, ValueError) as error: print(f"herdr-hub: {error}", file=sys.stderr) return 1 except KeyboardInterrupt: return 130 if __name__ == "__main__": raise SystemExit(main())