From 6c2c9672e1fb6c4a3321f5035cb790c00258dac6 Mon Sep 17 00:00:00 2001 From: Donald Date: Wed, 2 Sep 2026 01:33:38 +0000 Subject: [PATCH] Initial release of herdr session switcher --- .gitignore | 3 + Makefile | 10 + README.md | 211 ++++++ bin/herdr-hub | 15 + herdr_hub/__init__.py | 3 + herdr_hub/__main__.py | 5 + herdr_hub/cli.py | 1210 ++++++++++++++++++++++++++++++++ herdr_hub/remotes.py | 434 ++++++++++++ herdr_hub/selector.py | 676 ++++++++++++++++++ scripts/install.sh | 15 + tests/__init__.py | 1 + tests/test_cli_integration.py | 136 ++++ tests/test_discovery.py | 173 +++++ tests/test_remotes.py | 121 ++++ tests/test_selector.py | 181 +++++ tests/test_tmux_integration.py | 206 ++++++ 16 files changed, 3400 insertions(+) create mode 100644 .gitignore create mode 100644 Makefile create mode 100644 README.md create mode 100755 bin/herdr-hub create mode 100644 herdr_hub/__init__.py create mode 100644 herdr_hub/__main__.py create mode 100644 herdr_hub/cli.py create mode 100644 herdr_hub/remotes.py create mode 100644 herdr_hub/selector.py create mode 100755 scripts/install.sh create mode 100644 tests/__init__.py create mode 100644 tests/test_cli_integration.py create mode 100644 tests/test_discovery.py create mode 100644 tests/test_remotes.py create mode 100644 tests/test_selector.py create mode 100644 tests/test_tmux_integration.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..e11d7d5 --- /dev/null +++ b/.gitignore @@ -0,0 +1,3 @@ +__pycache__/ +*.py[cod] +.coverage diff --git a/Makefile b/Makefile new file mode 100644 index 0000000..3b229a2 --- /dev/null +++ b/Makefile @@ -0,0 +1,10 @@ +.PHONY: doctor install test + +doctor: + ./bin/herdr-hub doctor + +install: + ./scripts/install.sh + +test: + python3 -m unittest discover -s tests -v diff --git a/README.md b/README.md new file mode 100644 index 0000000..4f39f42 --- /dev/null +++ b/README.md @@ -0,0 +1,211 @@ +# herdr-hub + +`herdr-hub` 是一个 Linux 下的 Herdr session 切换器。它让 tmux 做最外层路由， +把每个本地或 `herdr --remote` attach 放进独立 tmux window，并在 Herdr 画面顶部 +显示可点击的 Sessions 条。 + +它不是 Herdr 插件，也不会修改 Herdr core。Herdr 0.8.2 的插件 API 不能向原生 +sidebar 插入 UI，也不能让 thin client 从一个 server 原地改连另一个 server； +所以顶部状态栏横跨整个终端，而不是只嵌在 Spaces 上方。 + +## 功能 + +- 自动读取 `herdr session list --json`，导入本机正在运行的 sessions。 +- 记住使用过的远端，并通过 SSH 实时读取运行中及已停止的 Herdr sessions。 +- 在 Linux `/proc/*/cmdline` 中查找仍在运行的 `herdr --remote …` clients，并自动记住目标。 +- 首次接管已有 remote client 时安全地再 attach，不结束原进程。 +- 一个 endpoint 对应一个 tmux window，隐藏 window 中的连接继续在线。 +- 顶部 Sessions 条在窄屏下也为每个 Tab 显示 Session 优先的短名称，并支持鼠标点击、数字切换和无需 fzf 的动态 popup 选择。 +- 可隐藏不再需要的 Tab，而不停止底层本地或远端 Herdr Session。 +- dead window 留在顶栏，可原地重连。 +- 不修改 `~/.tmux.conf`，不保存 SSH 凭据或远端 Session 缓存。 + +## 依赖 + +- Linux +- Python 3.11+ +- tmux 3.4+ +- Herdr 0.8.2+ +- OpenSSH client + +检查当前环境： + +```bash +./bin/herdr-hub doctor +``` + +## 安装 + +直接从仓库运行： + +```bash +./bin/herdr-hub +``` + +或把启动器链接到 `~/.local/bin`： + +```bash +make install +herdr-hub doctor +``` + +安装脚本只创建一个指向本仓库的符号链接。移动或删除仓库后需重新安装。 + +## 第一次接管 + +如果本地和 remote Herdr clients 已经分别开在其他终端中： + +1. 保持这些 clients 运行。 +2. 从一个新的、未嵌套在 tmux 中的终端执行 `herdr-hub`。 +3. Hub 扫描现有 clients，并为每个 endpoint 创建一个新的 attach window。 +4. 确认 Hub 中可以正常操作后，手动退出原来的 clients。 + +tmux 不能把另一个终端里已经运行的进程搬进 window，所以首次接管会短暂存在两个 +clients；它们连接的是同一个 Herdr server，不会复制 server 或 panes。 + +只有一个终端时，本地 Herdr 可以先 detach，因为 `herdr session list` 仍能发现后台 +server。remote target 在 client 退出后没有本地记录，因此直接使用： + +```bash +herdr-hub open --remote orangepi +herdr-hub open --remote orangepi --session work +herdr-hub open --remote ssh://me@server:2222 --remote-keybindings server +``` + +打开本地 named session： + +```bash +herdr-hub open --session work +``` + +`open` 会把 endpoint 保存在当前 tmux hub 的 window metadata 中。远端 target 还会 +同时加入 `~/.config/herdr-hub/remotes.json`，以后可直接从 Sessions +弹窗浏览这台机器上的其他 Session。 + +## 日常操作 + +Hub 使用独立 tmux socket `herdr-hub`，不会接管普通 tmux server。外层前缀固定为 +`Ctrl+A`，Herdr 原来的 `Ctrl+B` 可以直接使用。 + +| 操作 | 快捷键 / 命令 | +| --- | --- | +| 点击切换 | 鼠标点击顶部 session | +| 打开 Session 弹窗 | 点击左上角 `Sessions` 或 `Ctrl+A g` | +| 上一个/下一个 window | `Ctrl+A p` / `Ctrl+A n` | +| 按编号切换 | `Ctrl+A 0..9` | +| 扫描新增 sessions | `Ctrl+A r` 或 `herdr-hub refresh` | +| 强制重连当前 window | `Ctrl+A R` | +| 关闭/隐藏当前 Tab | `Ctrl+A &`，确认 `y`；也可在 Sessions 弹窗中选择 | +| detach Hub | `Ctrl+A d` | +| 重新进入 Hub | `herdr-hub` | +| 停止 Hub clients | `herdr-hub stop` | + +## 远端 Session 弹窗 + +弹窗由 tmux `display-popup` 和 Hub 内置选择器组成，不依赖 fzf。它会立即显示已打开的 +window，同时在后台并行查询所有已登记远端；每台机器的结果就绪后原地更新。 + +状态与选择行为： + +| 状态 | 选择后的行为 | +| --- | --- | +| `▶` 当前 / `●` 已打开 | 跳转到已有 window | +| `×` 已断线 | 自动重连并跳转 | +| `○` 远端运行 | 新建 attach window 并跳转 | +| `■` 远端已停止 | 启动该 Session、attach 并跳转 | + +弹窗靠前位置提供“关闭当前 Tab”；它只结束该 Tab 中的 Hub client，底层 Herdr +Session 和工作现场继续运行。每个远端下方都有“新建 Session”入口。弹窗底部还可添加或管理远端；移除远端只删除 +记忆记录，不关闭已有 window，也不停止远端 Herdr server。手机端单击条目即可执行， +键盘支持上下键、Enter、输入搜索、Backspace、`Ctrl+R` 刷新和 Esc 返回/关闭。 + +也可通过 CLI 管理远端： + +```bash +herdr-hub remote list +herdr-hub remote add orangepi +herdr-hub remote add ssh://me@server:2222 --remote-keybindings server +herdr-hub remote remove orangepi +``` + +注册表默认位于 `~/.config/herdr-hub/remotes.json`（遵循 `$XDG_CONFIG_HOME`），权限为 +`0600`。它只保存 SSH target 和 remote keybinding 模式。 + +## 复制与剪贴板 + +Hub 启用 tmux 的 OSC 52 剪贴板透传。Herdr 触发的自动复制，以及 tmux 接管后的 +鼠标拖选，都会尝试写入当前 Hub 访问端的本地剪贴板。远程 session 中复制的 +文本也回到访问端，不是远程机器的剪贴板。 + +访问端的终端需要支持并允许 OSC 52。如果终端或 Web terminal 禁止了这项能力， +可按住 `Shift` 再拖选，绕过 tmux 并使用终端原生复制。 + +当访问端宽度小于 80 列时，Hub 会自动进入窄屏模式：隐藏右侧快捷键说明， +所有 Tab 都显示状态和最多 12 列的短名称，并优先显示主机：默认 Session 的本地与 +远端分别显示 `local`、`host`，非默认 Session 才显示 `local/work`、`host/work`。 +宽度小于 50 列时左上角标题会缩成 `S`，继续给 Tab 留空间。点击 `Sessions`/`S` +可打开占 90% 屏幕的完整选择器。 + +`refresh` 是 add-only：它添加新发现的 endpoint，不自动删除已有 window，也不会改变 +当前焦点。通过 `Ctrl+A &` 或弹窗关闭的 endpoint 会在当前 Hub 生命周期内记为已隐藏， +刷新和断开终端后重新进入都不会让它自动回来。要恢复它，打开 Sessions 弹窗并重新选择； +也可以执行对应的 `herdr-hub open` 命令。远端通过 SSH 目录重新选择，本地隐藏的运行中 +Session 会列在“本机”下。Hub 至少保留一个 Tab；要全部退出请使用 `herdr-hub stop`。 + +当 endpoint 连接进程退出时，window 会保留：顶栏将它标记为 `× [断开]`， +当前 pane 底部会显示退出状态和 `Ctrl+A R` 重连提示。这与正常 detach Hub +不同；仅关闭外层终端后，直接再次执行 `herdr-hub` 即可。 + +`herdr-hub stop` 只结束 tmux 中的 client processes。本地及远端 Herdr servers 和 +它们的 panes 会继续运行。 + +## CLI + +```text +herdr-hub +herdr-hub start +herdr-hub refresh +herdr-hub open [--session NAME] +herdr-hub open --remote TARGET [--session NAME] + [--remote-keybindings local|server] +herdr-hub choose +herdr-hub remote list +herdr-hub remote add TARGET [--remote-keybindings local|server] +herdr-hub remote remove TARGET +herdr-hub reconnect [--window ID] +herdr-hub close [--window ID] +herdr-hub doctor +herdr-hub stop +``` + +## 自动发现与远端查询的边界 + +零配置的首次发现只覆盖： + +- 当前用户正在运行的本地 Herdr servers； +- 当前用户、同一 PID namespace 中仍然存活的 `herdr --remote` processes。 + +发现或手动添加一次后，remote target 会被记住，不再要求本地 client 持续存活。弹窗使用 +非交互 SSH 查询 `$HOME/.local/bin/herdr session list --json`；SSH key/passphrase 应已由 +ssh-agent 处理。未知 host key、需要密码或远端不可达时会显示可重试的错误，不会卡住 +整个弹窗。首次安装或 host-key 确认仍可在实际新建的 Herdr window 中交互完成。 + +首次使用且尚未登记时，以下情况仍无法自动找回 remote target： + +- remote client 已经退出，但远端 server 仍在后台运行； +- client 位于另一个 container/PID namespace； +- `/proc` 使用了阻止读取同用户 argv 的安全配置； +- 可执行文件不叫 `herdr`。 + +这些场景可在弹窗中选择“添加远端”，或使用 `herdr-hub open --remote …`。为了避免嵌套 tmux 与前缀冲突， +`herdr-hub` 应从普通终端启动；若当前已经在另一个 tmux server 中，它会拒绝再次 +attach。 + +## 开发与测试 + +```bash +make test +``` + +测试使用伪造的 `/proc` 和 Herdr executable；tmux 集成测试使用唯一 socket 和 +`/dev/null` 配置，不接触真实 Herdr sessions 或现有 tmux server。 diff --git a/bin/herdr-hub b/bin/herdr-hub new file mode 100755 index 0000000..757dcb4 --- /dev/null +++ b/bin/herdr-hub @@ -0,0 +1,15 @@ +#!/usr/bin/env python3 +"""Repository-local launcher for herdr_hub.""" + +from pathlib import Path +import sys + + +repository_root = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(repository_root)) + +from herdr_hub.cli import main # noqa: E402 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/herdr_hub/__init__.py b/herdr_hub/__init__.py new file mode 100644 index 0000000..93c5cd2 --- /dev/null +++ b/herdr_hub/__init__.py @@ -0,0 +1,3 @@ +"""tmux-based session switcher for Herdr.""" + +__version__ = "0.2.0" diff --git a/herdr_hub/__main__.py b/herdr_hub/__main__.py new file mode 100644 index 0000000..a049ad7 --- /dev/null +++ b/herdr_hub/__main__.py @@ -0,0 +1,5 @@ +from .cli import main + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/herdr_hub/cli.py b/herdr_hub/cli.py new file mode 100644 index 0000000..00b2775 --- /dev/null +++ b/herdr_hub/cli.py @@ -0,0 +1,1210 @@ +from __future__ import annotations + +import argparse +import curses +from dataclasses import dataclass, field +import json +import os +from pathlib import Path +import re +import shlex +import shutil +import subprocess +import sys +from typing import Callable, Iterable, Sequence +from urllib.parse import quote + +from .remotes import ( + RegistryError, + RemoteEntry, + RemoteRegistry, + default_registry_path, + validate_target, +) +from .selector import CatalogWindow, SelectorError, SelectorResult, SessionSelector + + +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_LABEL_OPTION = "@herdr_hub_label" +WINDOW_ENDPOINT_OPTION = "@herdr_hub_endpoint" +WINDOW_SHORT_LABEL_OPTION = "@herdr_hub_short_label" +SESSION_DISMISSED_OPTION = "@herdr_hub_dismissed" +SESSION_NAME_RE = re.compile(r"^[A-Za-z0-9._-]{1,64}$") + + +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" + + 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") + return + validate_remote_target(self.target) + if self.remote_keybindings not in {"local", "server"}: + raise ValueError("remote keybindings must be 'local' or 'server'") + + @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", + ) -> Endpoint: + return cls( + kind="remote", + target=target, + session=session, + remote_keybindings=remote_keybindings, + ) + + @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) -> list[str]: + command = [herdr_bin] + if self.kind == "remote": + assert self.target is not None + 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 validate_session_name(name: str) -> None: + if not SESSION_NAME_RE.fullmatch(name): + raise ValueError( + "session name must be 1-64 ASCII letters, numbers, '.', '_' or '-'" + ) + + +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, + }, + 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") + if not isinstance(kind, str) or not isinstance(session, str): + raise ValueError + if not isinstance(keybindings, str): + raise ValueError + if kind == "local": + return Endpoint.local(session) + if kind == "remote" and isinstance(target, str): + return Endpoint.remote(target, session, keybindings) + 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, + 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, + run: RunFunction = subprocess.run, + ) -> None: + self.tmux_bin = tmux_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._run = run + + 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 + + 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 · & 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]) + 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", + "-E", + "-w", + "90%", + "-h", + "90%", + "-T", + "Sessions", + choose, + ], + [ + "bind-key", + "-T", + "root", + "MouseDown1StatusLeft", + "display-popup", + "-E", + "-w", + "90%", + "-h", + "90%", + "-T", + "Sessions", + choose, + ], + ["bind-key", "-T", "prefix", "R", "run-shell", "-b", reconnect], + [ + "bind-key", + "-T", + "prefix", + "&", + "confirm-before", + "-p", + "关闭当前 Session Tab？ (y/n)", + confirm_close, + ], + ] + for command in binding_commands: + self.run(command) + + self.refresh_short_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_short_labels(self) -> None: + for window in self.managed_windows(): + label = ( + window.endpoint.short_label + if window.endpoint is not None + else truncate_label(window.name, 32) + ) + self.run( + [ + "set-option", + "-w", + "-t", + window.window_id, + WINDOW_SHORT_LABEL_OPTION, + label, + ] + ) + + 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) + + def create_window(self, endpoint: Endpoint, herdr_bin: str) -> ManagedWindow: + existing = self.window_for_key(endpoint.key) + if existing is not None: + return existing + self.restore_endpoint(endpoint.key) + command = "exec " + shlex.join(endpoint.command(herdr_bin)) + result = self.run( + [ + "new-window", + "-d", + "-P", + "-F", + "#{window_id}", + "-t", + f"{self.session_name}:", + "-n", + endpoint.label, + command, + ] + ) + window_id = result.stdout.strip() + if not window_id: + raise HubError("tmux did not return the new window id") + for option, value in [ + (WINDOW_KEY_OPTION, endpoint.key), + (WINDOW_COMMAND_OPTION, command), + (WINDOW_LABEL_OPTION, endpoint.label), + (WINDOW_ENDPOINT_OPTION, serialize_endpoint(endpoint)), + (WINDOW_SHORT_LABEL_OPTION, endpoint.short_label), + ]: + self.run(["set-option", "-w", "-t", window_id, option, value]) + self.run(["set-window-option", "-t", window_id, "remain-on-exit", "on"]) + return ManagedWindow( + window_id=window_id, + key=endpoint.key, + name=endpoint.label, + dead=False, + active=False, + endpoint=endpoint, + ) + + 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 + + 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]) + + 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 + + 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]) + + 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.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) + + 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) -> TmuxHub: + tmux_bin = resolve_binary("tmux", os.environ.get("HERDR_HUB_TMUX_BIN")) + return TmuxHub( + tmux_bin=tmux_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(), + ) + + +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 remember_remote_endpoints(endpoints: Iterable[Endpoint]) -> str | None: + entries = [ + RemoteEntry(endpoint.target, endpoint.remote_keybindings) + for endpoint in endpoints + if endpoint.kind == "remote" and endpoint.target is not None + ] + if not entries: + return None + try: + make_registry().upsert_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, + ) + 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: + if created: + 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]}" + if inside_socket(hub.socket_name): + 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, + ) + 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": + warning = remember_remote_endpoints([endpoint]) + if warning: + print(f"herdr-hub: warning: {warning}", 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()) + remember_remote_endpoints( + window.endpoint for window in windows if window.endpoint is not None + ) + try: + selector = SessionSelector( + catalog_windows(windows), + make_registry(), + resolve_ssh(), + local_sessions=(endpoint.session for endpoint in local_report.endpoints), + ) + result = selector.run() + except (RegistryError, SelectorError) as error: + raise HubError(str(error)) from error + if result is None: + 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, + ) + 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}") + return 0 + if action == "add": + entry = RemoteEntry(arguments.target, arguments.remote_keybindings) + 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) + 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) + 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) + 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 + 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", + ) + + 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_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 (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()) diff --git a/herdr_hub/remotes.py b/herdr_hub/remotes.py new file mode 100644 index 0000000..5eea5b8 --- /dev/null +++ b/herdr_hub/remotes.py @@ -0,0 +1,434 @@ +from __future__ import annotations + +from collections import deque +from dataclasses import dataclass +import fcntl +import json +import os +from pathlib import Path +import subprocess +import tempfile +import time +from typing import Iterable, Sequence +from urllib.parse import SplitResult, unquote, urlsplit + + +REGISTRY_VERSION = 1 +REMOTE_LIST_SCRIPT = ( + 'if [ -x "$HOME/.local/bin/herdr" ]; then ' + 'exec "$HOME/.local/bin/herdr" session list --json; ' + "elif command -v herdr >/dev/null 2>&1; then " + "exec herdr session list --json; " + "else echo 'herdr is not installed on the remote host' >&2; exit 127; fi" +) + + +class RegistryError(RuntimeError): + """A remote registry error which is safe to show to the user.""" + + +class RemoteTargetError(ValueError): + """A remote target cannot be represented safely as an SSH command.""" + + +@dataclass(frozen=True) +class RemoteEntry: + target: str + remote_keybindings: str = "local" + + def __post_init__(self) -> None: + validate_target(self.target) + if self.remote_keybindings not in {"local", "server"}: + raise ValueError("remote keybindings must be 'local' or 'server'") + + +@dataclass(frozen=True) +class RemoteSession: + name: str + running: bool + + +@dataclass(frozen=True) +class RemoteQueryResult: + entry: RemoteEntry + sessions: tuple[RemoteSession, ...] = () + error: str | None = None + + +def validate_target(target: str) -> None: + if not target: + raise RemoteTargetError("remote target cannot be empty") + if target.startswith("-"): + raise RemoteTargetError("remote target must not start with '-'") + if any(ord(character) < 32 or ord(character) == 127 for character in target): + raise RemoteTargetError("remote target must not contain control characters") + if target.startswith("ssh://"): + _parse_ssh_url(target) + + +def _parse_ssh_url(target: str) -> SplitResult: + parsed = urlsplit(target) + if parsed.scheme != "ssh" or parsed.hostname is None: + raise RemoteTargetError("invalid ssh:// remote target") + if parsed.password is not None: + raise RemoteTargetError("ssh:// targets must not contain passwords") + if parsed.path not in {"", "/"} or parsed.query or parsed.fragment: + raise RemoteTargetError("ssh:// targets must not contain a path, query or fragment") + try: + parsed.port + except ValueError as error: + raise RemoteTargetError(f"ssh:// target has an invalid port: {error}") from error + host = parsed.hostname + if host.startswith("-") or any(character.isspace() for character in host): + raise RemoteTargetError("ssh:// target contains an invalid host") + if parsed.username is not None: + username = unquote(parsed.username) + if ( + not username + or username.startswith("-") + or any( + character.isspace() + or ord(character) < 32 + or ord(character) == 127 + for character in username + ) + ): + raise RemoteTargetError("ssh:// target contains an invalid username") + return parsed + + +def default_registry_path() -> Path: + if override := os.environ.get("HERDR_HUB_REMOTES_FILE"): + return Path(override).expanduser() + config_home = os.environ.get("XDG_CONFIG_HOME") + root = Path(config_home).expanduser() if config_home else Path.home() / ".config" + return root / "herdr-hub" / "remotes.json" + + +class RemoteRegistry: + def __init__(self, path: Path | None = None) -> None: + self.path = path or default_registry_path() + + @property + def lock_path(self) -> Path: + return self.path.with_name(f".{self.path.name}.lock") + + def load(self) -> list[RemoteEntry]: + if not self.path.exists(): + return [] + try: + payload = json.loads(self.path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as error: + raise RegistryError(f"cannot read {self.path}: {error}") from error + return _parse_registry(payload, self.path) + + def upsert(self, entry: RemoteEntry) -> bool: + def update(entries: list[RemoteEntry]) -> tuple[list[RemoteEntry], bool]: + by_target = {item.target: item for item in entries} + changed = by_target.get(entry.target) != entry + by_target[entry.target] = entry + return sorted(by_target.values(), key=lambda item: item.target), changed + + return self._update(update) + + def upsert_many(self, entries: Iterable[RemoteEntry]) -> int: + incoming = list(entries) + if not incoming: + return 0 + + def update(current: list[RemoteEntry]) -> tuple[list[RemoteEntry], int]: + by_target = {item.target: item for item in current} + changed = 0 + for entry in incoming: + if by_target.get(entry.target) != entry: + changed += 1 + by_target[entry.target] = entry + return sorted(by_target.values(), key=lambda item: item.target), changed + + return self._update(update) + + def remove(self, target: str) -> bool: + validate_target(target) + + def update(entries: list[RemoteEntry]) -> tuple[list[RemoteEntry], bool]: + remaining = [entry for entry in entries if entry.target != target] + return remaining, len(remaining) != len(entries) + + return self._update(update) + + def _update(self, callback): # type: ignore[no-untyped-def] + try: + self.path.parent.mkdir(mode=0o700, parents=True, exist_ok=True) + with self.lock_path.open("a+", encoding="utf-8") as lock: + fcntl.flock(lock.fileno(), fcntl.LOCK_EX) + entries = self.load() + updated, outcome = callback(entries) + if updated != entries: + self._write(updated) + return outcome + except RegistryError: + raise + except OSError as error: + raise RegistryError(f"cannot update {self.path}: {error}") from error + + def _write(self, entries: Sequence[RemoteEntry]) -> None: + payload = { + "version": REGISTRY_VERSION, + "remotes": [ + { + "target": entry.target, + "remote_keybindings": entry.remote_keybindings, + } + for entry in entries + ], + } + temporary_name: str | None = None + try: + with tempfile.NamedTemporaryFile( + mode="w", + encoding="utf-8", + dir=self.path.parent, + prefix=f".{self.path.name}.", + delete=False, + ) as temporary: + temporary_name = temporary.name + json.dump(payload, temporary, ensure_ascii=False, indent=2) + temporary.write("\n") + temporary.flush() + os.fsync(temporary.fileno()) + os.chmod(temporary_name, 0o600) + os.replace(temporary_name, self.path) + directory_fd = os.open(self.path.parent, os.O_RDONLY | os.O_DIRECTORY) + try: + os.fsync(directory_fd) + finally: + os.close(directory_fd) + finally: + if temporary_name is not None: + try: + Path(temporary_name).unlink(missing_ok=True) + except OSError: + pass + + +def _parse_registry(payload: object, path: Path) -> list[RemoteEntry]: + if not isinstance(payload, dict) or payload.get("version") != REGISTRY_VERSION: + raise RegistryError(f"{path} has an unsupported or missing version") + raw_entries = payload.get("remotes") + if not isinstance(raw_entries, list): + raise RegistryError(f"{path} does not contain a remotes list") + entries: list[RemoteEntry] = [] + seen: set[str] = set() + for index, raw in enumerate(raw_entries): + if not isinstance(raw, dict): + raise RegistryError(f"{path}: remotes[{index}] is not an object") + target = raw.get("target") + keybindings = raw.get("remote_keybindings", "local") + if not isinstance(target, str) or not isinstance(keybindings, str): + raise RegistryError(f"{path}: remotes[{index}] has invalid fields") + try: + entry = RemoteEntry(target, keybindings) + except ValueError as error: + raise RegistryError(f"{path}: remotes[{index}]: {error}") from error + if entry.target in seen: + raise RegistryError(f"{path}: duplicate remote target {entry.target!r}") + seen.add(entry.target) + entries.append(entry) + return sorted(entries, key=lambda entry: entry.target) + + +def ssh_destination_arguments(target: str) -> list[str]: + validate_target(target) + if not target.startswith("ssh://"): + return [target] + + parsed = _parse_ssh_url(target) + assert parsed.hostname is not None + host = parsed.hostname + arguments: list[str] = [] + if parsed.port is not None: + arguments.extend(["-p", str(parsed.port)]) + if parsed.username is not None: + username = unquote(parsed.username) + arguments.extend(["-l", username]) + arguments.append(host) + return arguments + + +def remote_session_command( + ssh_bin: str, + entry: RemoteEntry, + *, + connect_timeout: int = 4, +) -> list[str]: + return [ + ssh_bin, + "-T", + "-o", + "BatchMode=yes", + "-o", + "StrictHostKeyChecking=yes", + "-o", + "ConnectionAttempts=1", + "-o", + f"ConnectTimeout={connect_timeout}", + *ssh_destination_arguments(entry.target), + REMOTE_LIST_SCRIPT, + ] + + +def parse_remote_session_result( + entry: RemoteEntry, + returncode: int, + stdout: str, + stderr: str, +) -> RemoteQueryResult: + if returncode != 0: + detail = (stderr or stdout or f"ssh exited with status {returncode}").strip() + detail = " ".join(detail.splitlines()) + if len(detail) > 180: + detail = detail[:179] + "…" + return RemoteQueryResult(entry, error=detail) + try: + payload = json.loads(stdout) + raw_sessions = payload["sessions"] + if not isinstance(raw_sessions, list): + raise TypeError("sessions is not a list") + except (json.JSONDecodeError, KeyError, TypeError) as error: + return RemoteQueryResult(entry, error=f"invalid Herdr session JSON: {error}") + + sessions: dict[str, RemoteSession] = {} + for raw in raw_sessions: + if not isinstance(raw, dict): + continue + name = raw.get("name") + running = raw.get("running") + if isinstance(name, str) and isinstance(running, bool): + sessions.setdefault(name, RemoteSession(name, running)) + ordered = sorted( + sessions.values(), + key=lambda session: (session.name != "default", session.name), + ) + return RemoteQueryResult(entry, sessions=tuple(ordered)) + + +@dataclass +class _RemoteJob: + entry: RemoteEntry + process: subprocess.Popen[str] + started_at: float + + +class RemoteQueryPool: + """Pollable bounded pool of SSH processes for an interactive selector.""" + + def __init__( + self, + ssh_bin: str, + *, + max_parallel: int = 6, + query_timeout: float = 6.0, + connect_timeout: int = 4, + ) -> None: + self.ssh_bin = ssh_bin + self.max_parallel = max_parallel + self.query_timeout = query_timeout + self.connect_timeout = connect_timeout + self._pending: deque[RemoteEntry] = deque() + self._active: dict[str, _RemoteJob] = {} + self._completed: deque[RemoteQueryResult] = deque() + + def enqueue(self, entry: RemoteEntry) -> None: + self.cancel(entry.target) + self._pending.append(entry) + self._start_pending() + + def enqueue_many(self, entries: Iterable[RemoteEntry]) -> None: + for entry in entries: + self.cancel(entry.target) + self._pending.append(entry) + self._start_pending() + + def cancel(self, target: str) -> None: + self._pending = deque(entry for entry in self._pending if entry.target != target) + job = self._active.pop(target, None) + if job is not None: + self._stop_process(job.process) + self._start_pending() + + def is_loading(self, target: str) -> bool: + return target in self._active or any( + entry.target == target for entry in self._pending + ) + + def poll(self) -> list[RemoteQueryResult]: + finished = list(self._completed) + self._completed.clear() + now = time.monotonic() + for target, job in list(self._active.items()): + returncode = job.process.poll() + timed_out = now - job.started_at >= self.query_timeout + if returncode is None and not timed_out: + continue + if timed_out and returncode is None: + job.process.kill() + stdout, stderr = job.process.communicate() + self._active.pop(target, None) + if timed_out: + finished.append( + RemoteQueryResult( + job.entry, + error=f"remote query timed out after {self.query_timeout:g}s", + ) + ) + else: + finished.append( + parse_remote_session_result( + job.entry, + job.process.returncode, + stdout, + stderr, + ) + ) + self._start_pending() + return finished + + def close(self) -> None: + self._pending.clear() + for job in list(self._active.values()): + self._stop_process(job.process) + self._active.clear() + + def _start_pending(self) -> None: + while self._pending and len(self._active) < self.max_parallel: + entry = self._pending.popleft() + try: + process = subprocess.Popen( + remote_session_command( + self.ssh_bin, + entry, + connect_timeout=self.connect_timeout, + ), + stdin=subprocess.DEVNULL, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + except (OSError, ValueError) as error: + self._completed.append( + RemoteQueryResult(entry, error=f"cannot start remote query: {error}") + ) + continue + self._active[entry.target] = _RemoteJob(entry, process, time.monotonic()) + + @staticmethod + def _stop_process(process: subprocess.Popen[str]) -> None: + if process.poll() is not None: + process.communicate() + return + process.terminate() + try: + process.communicate(timeout=0.25) + except subprocess.TimeoutExpired: + process.kill() + process.communicate() diff --git a/herdr_hub/selector.py b/herdr_hub/selector.py new file mode 100644 index 0000000..1aefdf3 --- /dev/null +++ b/herdr_hub/selector.py @@ -0,0 +1,676 @@ +from __future__ import annotations + +from dataclasses import dataclass +import curses +import re +from typing import Iterable + +from .remotes import ( + RegistryError, + RemoteEntry, + RemoteQueryPool, + RemoteQueryResult, + RemoteRegistry, +) + + +SESSION_NAME_RE = re.compile(r"^[A-Za-z0-9._-]{1,64}$") + + +class SelectorError(RuntimeError): + """The interactive selector could not be displayed.""" + + +@dataclass(frozen=True) +class CatalogWindow: + window_id: str + key: str + label: str + dead: bool + active: bool + kind: str | None = None + target: str | None = None + session: str | None = None + + +@dataclass(frozen=True) +class SelectorResult: + kind: str + window_id: str | None = None + dead: bool = False + target: str | None = None + session: str | None = None + remote_keybindings: str = "local" + + @classmethod + def window(cls, window_id: str, *, dead: bool) -> SelectorResult: + return cls(kind="window", window_id=window_id, dead=dead) + + @classmethod + def close(cls, window_id: str) -> SelectorResult: + return cls(kind="close", window_id=window_id) + + @classmethod + def local(cls, session: str) -> SelectorResult: + return cls(kind="local", session=session) + + @classmethod + def remote( + cls, + target: str, + session: str, + remote_keybindings: str, + ) -> SelectorResult: + return cls( + kind="remote", + target=target, + session=session, + remote_keybindings=remote_keybindings, + ) + + +@dataclass(frozen=True) +class PickerRow: + key: str + kind: str + main: str + detail: str = "" + selectable: bool = True + window_id: str | None = None + dead: bool = False + target: str | None = None + session: str | None = None + remote_keybindings: str = "local" + + +def build_main_rows( + windows: Iterable[CatalogWindow], + entries: Iterable[RemoteEntry], + results: dict[str, RemoteQueryResult], + loading_targets: set[str], + *, + local_sessions: Iterable[str] = (), + registry_error: str | None = None, +) -> list[PickerRow]: + window_list = list(windows) + entry_list = sorted(entries, key=lambda entry: entry.target) + registered = {entry.target for entry in entry_list} + rows: list[PickerRow] = [] + + standalone = [ + window + for window in window_list + if window.kind == "local" + or window.target is None + or window.target not in registered + ] + if standalone: + rows.append(PickerRow("heading:open", "heading", "已打开", selectable=False)) + for window in standalone: + rows.append(_window_row(window)) + + open_local_sessions = { + window.session + for window in window_list + if window.kind == "local" and window.session is not None + } + available_local_sessions = sorted( + { + session + for session in local_sessions + if SESSION_NAME_RE.fullmatch(session) and session not in open_local_sessions + }, + key=lambda name: (name != "default", name), + ) + if available_local_sessions: + rows.append(PickerRow("heading:local", "heading", "本机", selectable=False)) + rows.extend( + PickerRow( + f"local:{session}", + "local", + f"○ {session}", + "未在 Tab 中 · 单击打开", + session=session, + ) + for session in available_local_sessions + ) + + active_window = next((window for window in window_list if window.active), None) + if active_window is not None: + can_close = len(window_list) > 1 + rows.append( + PickerRow( + "close-active", + "close", + f"− 关闭当前 Tab：{active_window.label}", + "Herdr Session 保持运行" if can_close else "至少保留一个 Session Tab", + selectable=can_close, + window_id=active_window.window_id, + ) + ) + + remote_windows = { + (window.target, window.session): window + for window in window_list + if window.kind == "remote" + and window.target is not None + and window.session is not None + } + for entry in entry_list: + rows.append( + PickerRow( + f"heading:{entry.target}", + "heading", + entry.target, + "远端", + selectable=False, + ) + ) + result = results.get(entry.target) + session_by_name = ( + { + session.name: session + for session in result.sessions + if SESSION_NAME_RE.fullmatch(session.name) + } + if result is not None and result.error is None + else {} + ) + for target, session_name in remote_windows: + if target == entry.target and session_name is not None: + session_by_name.setdefault(session_name, None) + + for session_name in sorted( + session_by_name, + key=lambda name: (name != "default", name), + ): + window = remote_windows.get((entry.target, session_name)) + if window is not None: + rows.append(_window_row(window)) + continue + session = session_by_name[session_name] + assert session is not None + rows.append( + PickerRow( + f"remote:{entry.target}:{session_name}", + "remote", + f"{'○' if session.running else '■'} {session_name}", + "远端运行 · 单击打开" if session.running else "已停止 · 单击启动", + target=entry.target, + session=session_name, + remote_keybindings=entry.remote_keybindings, + ) + ) + + if entry.target in loading_targets: + rows.append( + PickerRow( + f"loading:{entry.target}", + "loading", + "… 正在查询 Sessions", + selectable=False, + ) + ) + elif result is not None and result.error is not None: + rows.append( + PickerRow( + f"retry:{entry.target}", + "retry", + "! 查询失败 · 单击重试", + result.error, + target=entry.target, + ) + ) + elif result is not None and not session_by_name: + rows.append( + PickerRow( + f"empty:{entry.target}", + "empty", + "没有已有 Session", + selectable=False, + ) + ) + rows.append( + PickerRow( + f"new:{entry.target}", + "new", + "+ 新建 Session…", + entry.target, + target=entry.target, + remote_keybindings=entry.remote_keybindings, + ) + ) + + if registry_error is not None: + rows.append( + PickerRow( + "registry-error", + "error", + "! 远端配置不可用", + registry_error, + selectable=False, + ) + ) + else: + rows.extend( + [ + PickerRow("add-remote", "add", "+ 添加远端…", "SSH target"), + PickerRow("manage-remotes", "manage", "⚙ 管理远端…"), + ] + ) + return rows + + +def _window_row(window: CatalogWindow) -> PickerRow: + symbol = "▶" if window.active else "×" if window.dead else "●" + detail = "当前" if window.active else "已断线 · 单击重连" if window.dead else "Hub 已打开" + return PickerRow( + f"window:{window.window_id}", + "window", + f"{symbol} {window.label}", + detail, + window_id=window.window_id, + dead=window.dead, + target=window.target, + session=window.session, + ) + + +def filter_rows(rows: Iterable[PickerRow], query: str) -> list[PickerRow]: + if not query: + return list(rows) + needle = query.casefold() + return [ + row + for row in rows + if needle in f"{row.main} {row.detail}".casefold() + ] + + +class SessionSelector: + def __init__( + self, + windows: Iterable[CatalogWindow], + registry: RemoteRegistry, + ssh_bin: str, + local_sessions: Iterable[str] = (), + ) -> None: + self.windows = list(windows) + self.local_sessions = list(local_sessions) + self.registry = registry + self.pool = RemoteQueryPool(ssh_bin) + self.entries: dict[str, RemoteEntry] = {} + self.results: dict[str, RemoteQueryResult] = {} + self.registry_error: str | None = None + self.message = "" + self.query = "" + self.selected_key: str | None = None + self.scroll = 0 + self.view = "main" + self.input_value = "" + self.input_target: str | None = None + self._drawn_rows: dict[int, PickerRow] = {} + self._load_registry() + + def _load_registry(self) -> None: + try: + entries = self.registry.load() + except RegistryError as error: + self.registry_error = str(error) + return + self.registry_error = None + self.entries = {entry.target: entry for entry in entries} + self.pool.enqueue_many(entries) + + def run(self) -> SelectorResult | None: + try: + return curses.wrapper(self._curses_main) + except curses.error as error: + raise SelectorError(f"could not initialize the terminal selector: {error}") from error + finally: + self.pool.close() + + def _curses_main(self, screen) -> SelectorResult | None: # type: ignore[no-untyped-def] + screen.keypad(True) + screen.timeout(100) + try: + curses.curs_set(0) + except curses.error: + pass + try: + curses.use_default_colors() + except curses.error: + pass + curses.mouseinterval(150) + curses.mousemask(curses.ALL_MOUSE_EVENTS) + + while True: + for result in self.pool.poll(): + self.results[result.entry.target] = result + rows = self._rows() + filtered = filter_rows(rows, self.query) if self.view == "main" else rows + self._normalize_selection(filtered) + self._draw(screen, filtered) + try: + key = screen.get_wch() + except curses.error: + continue + + outcome = self._handle_key(key, filtered) + if isinstance(outcome, SelectorResult) or outcome is False: + return outcome if isinstance(outcome, SelectorResult) else None + + def _rows(self) -> list[PickerRow]: + if self.view == "manage": + rows = [PickerRow("manage-heading", "heading", "管理远端", selectable=False)] + rows.extend( + PickerRow( + f"remove:{entry.target}", + "remove", + f"− {entry.target}", + "单击后确认移除", + target=entry.target, + ) + for entry in sorted(self.entries.values(), key=lambda item: item.target) + ) + if len(rows) == 1: + rows.append(PickerRow("no-remotes", "empty", "没有已登记远端", selectable=False)) + rows.append(PickerRow("manage-back", "back", "← 返回 Sessions")) + return rows + if self.view == "confirm-remove": + target = self.input_target or "" + return [ + PickerRow("confirm-heading", "heading", f"移除 {target}？", "不会关闭已有窗口或远端 server", selectable=False), + PickerRow("confirm-yes", "confirm-yes", "移除", target=target), + PickerRow("confirm-no", "back", "取消"), + ] + if self.view in {"input-add", "input-new"}: + return [] + return build_main_rows( + self.windows, + self.entries.values(), + self.results, + {target for target in self.entries if self.pool.is_loading(target)}, + local_sessions=self.local_sessions, + registry_error=self.registry_error, + ) + + def _handle_key(self, key, rows: list[PickerRow]): # type: ignore[no-untyped-def] + if self.view in {"input-add", "input-new"}: + return self._handle_input_key(key) + if key == curses.KEY_MOUSE: + return self._handle_mouse(rows) + if key in {curses.KEY_UP}: + self._move(rows, -1) + return None + if key in {curses.KEY_DOWN}: + self._move(rows, 1) + return None + if key == curses.KEY_PPAGE: + self._move(rows, -8) + return None + if key == curses.KEY_NPAGE: + self._move(rows, 8) + return None + if key in {curses.KEY_ENTER, "\n", "\r"}: + return self._activate(self._selected_row(rows)) + if key == "\x12" and self.view == "main": + self.results.clear() + self.pool.enqueue_many(self.entries.values()) + self.message = "正在刷新所有远端" + return None + if key == "\x1b": + if self.view != "main": + self._show_main() + return None + if self.query: + self.query = "" + self.selected_key = None + return None + return False + if key in {curses.KEY_BACKSPACE, "\x7f", "\b"} and self.view == "main": + self.query = self.query[:-1] + self.selected_key = None + return None + if isinstance(key, str) and key.isprintable() and self.view == "main": + if len(self.query) < 128: + self.query += key + self.selected_key = None + return None + return None + + def _handle_input_key(self, key): # type: ignore[no-untyped-def] + if key == "\x1b": + self._show_main() + return None + if key in {curses.KEY_ENTER, "\n", "\r"}: + value = self.input_value.strip() + if self.view == "input-add": + return self._finish_add(value) + return self._finish_new(value) + if key in {curses.KEY_BACKSPACE, "\x7f", "\b"}: + self.input_value = self.input_value[:-1] + return None + if isinstance(key, str) and key.isprintable() and len(self.input_value) < 512: + self.input_value += key + return None + + def _handle_mouse(self, rows: list[PickerRow]): + try: + _, _, y, _, state = curses.getmouse() + except curses.error: + return None + if state & getattr(curses, "BUTTON4_PRESSED", 0): + self._move(rows, -3) + return None + if state & getattr(curses, "BUTTON5_PRESSED", 0): + self._move(rows, 3) + return None + clicked = ( + getattr(curses, "BUTTON1_CLICKED", 0) + | getattr(curses, "BUTTON1_DOUBLE_CLICKED", 0) + ) + if state & clicked: + row = self._drawn_rows.get(y) + if row is not None and row.selectable: + self.selected_key = row.key + return self._activate(row) + return None + + def _activate(self, row: PickerRow | None): + if row is None: + return None + if row.kind == "window": + assert row.window_id is not None + return SelectorResult.window(row.window_id, dead=row.dead) + if row.kind == "close": + assert row.window_id is not None + return SelectorResult.close(row.window_id) + if row.kind == "local": + assert row.session is not None + return SelectorResult.local(row.session) + if row.kind == "remote": + assert row.target is not None and row.session is not None + return SelectorResult.remote( + row.target, + row.session, + row.remote_keybindings, + ) + if row.kind == "retry": + assert row.target is not None + self.results.pop(row.target, None) + self.pool.enqueue(self.entries[row.target]) + self.message = f"正在重试 {row.target}" + elif row.kind == "new": + self.view = "input-new" + self.input_target = row.target + self.input_value = "" + self.message = "" + elif row.kind == "add": + self.view = "input-add" + self.input_target = None + self.input_value = "" + self.message = "" + elif row.kind == "manage": + self.view = "manage" + self.query = "" + self.selected_key = None + elif row.kind == "remove": + self.view = "confirm-remove" + self.input_target = row.target + self.selected_key = "confirm-no" + elif row.kind == "confirm-yes": + self._finish_remove(row.target or "") + elif row.kind == "back": + self._show_main() + return None + + def _finish_add(self, target: str): + try: + entry = RemoteEntry(target) + self.registry.upsert(entry) + except (ValueError, RegistryError) as error: + self.message = str(error) + return None + self.entries[target] = entry + self.results.pop(target, None) + self.pool.enqueue(entry) + self._show_main() + self.message = f"已添加 {target}，正在查询" + return None + + def _finish_new(self, session: str): + if not SESSION_NAME_RE.fullmatch(session): + self.message = "Session 名称需为 1-64 位字母、数字、. _ -" + return None + target = self.input_target + assert target is not None + entry = self.entries[target] + return SelectorResult.remote(target, session, entry.remote_keybindings) + + def _finish_remove(self, target: str) -> None: + try: + self.registry.remove(target) + except RegistryError as error: + self.message = str(error) + self._show_main() + return + self.pool.cancel(target) + self.entries.pop(target, None) + self.results.pop(target, None) + self._show_main() + self.message = f"已移除 {target}；已有 window 保持不变" + + def _show_main(self) -> None: + self.view = "main" + self.input_target = None + self.input_value = "" + self.selected_key = None + self.scroll = 0 + + def _normalize_selection(self, rows: list[PickerRow]) -> None: + selectable = [row for row in rows if row.selectable] + if not selectable: + self.selected_key = None + self.scroll = 0 + return + if self.selected_key not in {row.key for row in selectable}: + self.selected_key = selectable[0].key + + def _selected_row(self, rows: list[PickerRow]) -> PickerRow | None: + return next((row for row in rows if row.key == self.selected_key), None) + + def _move(self, rows: list[PickerRow], delta: int) -> None: + selectable = [row for row in rows if row.selectable] + if not selectable: + return + current = next( + (index for index, row in enumerate(selectable) if row.key == self.selected_key), + 0, + ) + next_index = max(0, min(len(selectable) - 1, current + delta)) + self.selected_key = selectable[next_index].key + + def _draw(self, screen, rows: list[PickerRow]) -> None: # type: ignore[no-untyped-def] + screen.erase() + height, width = screen.getmaxyx() + self._drawn_rows = {} + if height < 8 or width < 28: + self._safe_add(screen, 0, 0, "终端太小，请旋转屏幕", curses.A_BOLD) + screen.refresh() + return + + if self.view in {"input-add", "input-new"}: + self._draw_input(screen, height, width) + return + try: + curses.curs_set(0) + except curses.error: + pass + + title = "Sessions" if self.view == "main" else "管理远端" + self._safe_add(screen, 1, 2, title, curses.A_BOLD) + if self.view == "main": + self._safe_add(screen, 2, 2, f"搜索: {self.query}", curses.A_DIM) + if self.message: + self._safe_add(screen, 3, 2, self.message, curses.A_DIM) + + list_top = 4 + list_height = max(1, height - list_top - 2) + selected_index = next( + (index for index, row in enumerate(rows) if row.key == self.selected_key), + 0, + ) + if selected_index < self.scroll: + self.scroll = selected_index + elif selected_index >= self.scroll + list_height: + self.scroll = selected_index - list_height + 1 + self.scroll = max(0, min(self.scroll, max(0, len(rows) - list_height))) + + for offset, row in enumerate(rows[self.scroll : self.scroll + list_height]): + y = list_top + offset + self._drawn_rows[y] = row + selected = row.key == self.selected_key + attribute = curses.A_REVERSE if selected else curses.A_NORMAL + if row.kind == "heading": + attribute |= curses.A_BOLD + elif not row.selectable: + attribute |= curses.A_DIM + detail_room = max(0, width // 2 - 2) + main_room = max(8, width - detail_room - 6) + main = row.main[:main_room] + detail = row.detail[:detail_room] + line = f" {main:<{main_room}} {detail}" + self._safe_add(screen, y, 0, line, attribute) + + footer = "单击/Enter 打开 · 输入搜索 · Ctrl+R 刷新 · Esc 关闭" + if self.view != "main": + footer = "单击/Enter 选择 · Esc 返回" + self._safe_add(screen, height - 1, 1, footer, curses.A_DIM) + screen.refresh() + + def _draw_input(self, screen, height: int, width: int) -> None: # type: ignore[no-untyped-def] + is_add = self.view == "input-add" + title = "添加远端" if is_add else f"在 {self.input_target} 新建 Session" + prompt = "SSH target" if is_add else "Session 名称" + self._safe_add(screen, 1, 2, title, curses.A_BOLD) + self._safe_add(screen, 3, 2, f"{prompt}: {self.input_value}") + if self.message: + self._safe_add(screen, 5, 2, self.message, curses.A_DIM) + self._safe_add(screen, height - 1, 1, "Enter 确认 · Esc 返回", curses.A_DIM) + try: + curses.curs_set(1) + screen.move(3, min(width - 2, len(prompt) + len(self.input_value) + 4)) + except curses.error: + pass + screen.refresh() + + @staticmethod + def _safe_add(screen, y: int, x: int, text: str, attribute: int = curses.A_NORMAL) -> None: # type: ignore[no-untyped-def] + height, width = screen.getmaxyx() + if y < 0 or y >= height or x >= width: + return + try: + screen.addnstr(y, x, text, max(0, width - x - 1), attribute) + except curses.error: + pass diff --git a/scripts/install.sh b/scripts/install.sh new file mode 100755 index 0000000..fcae0d4 --- /dev/null +++ b/scripts/install.sh @@ -0,0 +1,15 @@ +#!/bin/sh +set -eu + +script_dir=$(CDPATH= cd -- "$(dirname -- "$0")" && pwd) +repository_root=$(CDPATH= cd -- "$script_dir/.." && pwd) +install_dir=${XDG_BIN_HOME:-"$HOME/.local/bin"} + +mkdir -p "$install_dir" +ln -sfn "$repository_root/bin/herdr-hub" "$install_dir/herdr-hub" + +printf 'installed %s\n' "$install_dir/herdr-hub" +case ":${PATH:-}:" in + *:"$install_dir":*) ;; + *) printf 'note: add %s to PATH\n' "$install_dir" ;; +esac diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..578a465 --- /dev/null +++ b/tests/__init__.py @@ -0,0 +1 @@ +"""Test suite for herdr-hub.""" diff --git a/tests/test_cli_integration.py b/tests/test_cli_integration.py new file mode 100644 index 0000000..dc31deb --- /dev/null +++ b/tests/test_cli_integration.py @@ -0,0 +1,136 @@ +from __future__ import annotations + +import os +from pathlib import Path +import shutil +import subprocess +import tempfile +import unittest +import uuid + +from herdr_hub.cli import TmuxHub + + +@unittest.skipUnless(shutil.which("tmux"), "tmux is required") +class CliIntegrationTests(unittest.TestCase): + def test_remote_registry_cli_round_trip(self) -> None: + repository_root = Path(__file__).resolve().parents[1] + launcher = repository_root / "bin" / "herdr-hub" + with tempfile.TemporaryDirectory() as directory: + environment = { + **os.environ, + "HERDR_HUB_REMOTES_FILE": str(Path(directory) / "remotes.json"), + } + add = subprocess.run( + [ + str(launcher), + "remote", + "add", + "orangepi", + "--remote-keybindings", + "server", + ], + cwd=repository_root, + env=environment, + capture_output=True, + text=True, + check=False, + ) + self.assertEqual(add.returncode, 0, add.stderr) + listing = subprocess.run( + [str(launcher), "remote", "list"], + cwd=repository_root, + env=environment, + capture_output=True, + text=True, + check=False, + ) + self.assertEqual(listing.stdout.strip(), "orangepi\tserver") + remove = subprocess.run( + [str(launcher), "remote", "remove", "orangepi"], + cwd=repository_root, + env=environment, + capture_output=True, + text=True, + check=False, + ) + self.assertEqual(remove.returncode, 0, remove.stderr) + self.assertIn("removed orangepi", remove.stdout) + + def test_repository_launcher_refreshes_and_stops_isolated_hub(self) -> None: + repository_root = Path(__file__).resolve().parents[1] + launcher = repository_root / "bin" / "herdr-hub" + socket_name = f"herdr-hub-cli-test-{os.getpid()}-{uuid.uuid4().hex[:8]}" + socket_path = Path( + os.environ.get("TMUX_TMPDIR", f"/tmp/tmux-{os.getuid()}") + ) / socket_name + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + proc_root = root / "proc" + proc_root.mkdir() + fake_herdr = root / "herdr" + fake_herdr.write_text( + "#!/bin/sh\n" + "if [ \"${1:-}\" = \"session\" ]; then\n" + " printf '%s\\n' '{\"sessions\":[{\"name\":\"default\",\"running\":true}]}'\n" + " exit 0\n" + "fi\n" + "exec sleep 30\n", + encoding="utf-8", + ) + fake_herdr.chmod(0o755) + environment = { + **os.environ, + "HERDR_HUB_HERDR_BIN": str(fake_herdr), + "HERDR_HUB_TMUX_BIN": shutil.which("tmux") or "tmux", + "HERDR_HUB_SOCKET": socket_name, + "HERDR_HUB_SESSION": "hub", + "HERDR_HUB_TMUX_CONFIG": "/dev/null", + "HERDR_HUB_PROC_ROOT": str(proc_root), + } + + refresh = subprocess.run( + [str(launcher), "refresh"], + cwd=repository_root, + env=environment, + capture_output=True, + text=True, + check=False, + ) + self.assertEqual(refresh.returncode, 0, refresh.stderr) + self.assertIn("added 1 session", refresh.stdout) + + hub = TmuxHub( + tmux_bin=shutil.which("tmux") or "tmux", + socket_name=socket_name, + session_name="hub", + config_file="/dev/null", + script_path=str(launcher), + ) + + def cleanup_hub() -> None: + if hub.exists(): + hub.stop() + socket_path.unlink(missing_ok=True) + + self.addCleanup(cleanup_hub) + self.assertEqual( + [window.key for window in hub.managed_windows()], + ["local:default"], + ) + + stop = subprocess.run( + [str(launcher), "stop"], + cwd=repository_root, + env=environment, + capture_output=True, + text=True, + check=False, + ) + self.assertEqual(stop.returncode, 0, stop.stderr) + self.assertIn("Herdr servers were left running", stop.stdout) + self.assertFalse(hub.exists()) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_discovery.py b/tests/test_discovery.py new file mode 100644 index 0000000..a9998f4 --- /dev/null +++ b/tests/test_discovery.py @@ -0,0 +1,173 @@ +from __future__ import annotations + +import json +import os +from pathlib import Path +import subprocess +import tempfile +import unittest + +from herdr_hub.cli import ( + DEFAULT_SESSION_NAME, + Endpoint, + decode_cmdline, + dedupe_endpoints, + deserialize_endpoint, + discover_local_sessions, + discover_remote_processes, + parse_remote_argv, + serialize_endpoint, +) + + +class EndpointTests(unittest.TestCase): + def test_local_and_remote_commands(self) -> None: + self.assertEqual(Endpoint.local().command("herdr"), ["herdr"]) + self.assertEqual( + Endpoint.local("work").command("/bin/herdr"), + ["/bin/herdr", "--session", "work"], + ) + self.assertEqual( + Endpoint.remote("orangepi", "work", "server").command("herdr"), + [ + "herdr", + "--remote", + "orangepi", + "--remote-keybindings", + "server", + "--session", + "work", + ], + ) + + def test_identity_ignores_remote_keybinding_mode(self) -> None: + local = Endpoint.remote("orangepi", remote_keybindings="local") + server = Endpoint.remote("orangepi", remote_keybindings="server") + self.assertEqual(local.key, server.key) + self.assertEqual(dedupe_endpoints([local, server]), [local]) + + def test_short_labels_put_the_host_first_and_omit_default(self) -> None: + self.assertEqual(Endpoint.local().short_label, "local") + self.assertEqual(Endpoint.local("work").short_label, "local/work") + self.assertEqual(Endpoint.remote("orangepi").short_label, "orangepi") + self.assertEqual(Endpoint.remote("orangepi", "work").short_label, "orangepi/work") + + def test_default_local_sorts_first(self) -> None: + endpoints = dedupe_endpoints( + [ + Endpoint.remote("z"), + Endpoint.local("work"), + Endpoint.local(), + Endpoint.remote("a"), + ] + ) + self.assertEqual( + [endpoint.key for endpoint in endpoints], + ["local:default", "local:work", "remote:a:default", "remote:z:default"], + ) + + def test_endpoint_metadata_and_legacy_command_round_trip(self) -> None: + endpoint = Endpoint.remote("ssh://me@host:2222", "work", "server") + self.assertEqual(deserialize_endpoint(serialize_endpoint(endpoint)), endpoint) + self.assertEqual( + deserialize_endpoint( + "", + "exec /usr/local/bin/herdr --remote orangepi --session work", + ), + Endpoint.remote("orangepi", "work"), + ) + + +class RemoteArgvTests(unittest.TestCase): + def test_separate_arguments(self) -> None: + endpoint = parse_remote_argv( + [ + "/home/me/.local/bin/herdr", + "--remote", + "orangepi", + "--session", + "work", + "--remote-keybindings", + "server", + ] + ) + self.assertEqual(endpoint, Endpoint.remote("orangepi", "work", "server")) + + def test_equal_arguments_and_arbitrary_order(self) -> None: + endpoint = parse_remote_argv( + ["herdr", "--session=work", "--remote=ssh://me@host:2222"] + ) + self.assertEqual(endpoint, Endpoint.remote("ssh://me@host:2222", "work")) + + def test_default_session(self) -> None: + endpoint = parse_remote_argv(["herdr", "--remote", "orangepi"]) + self.assertIsNotNone(endpoint) + assert endpoint is not None + self.assertEqual(endpoint.session, DEFAULT_SESSION_NAME) + + def test_non_herdr_and_malformed_commands_are_ignored(self) -> None: + self.assertIsNone(parse_remote_argv(["bash", "-c", "herdr --remote host"])) + self.assertIsNone(parse_remote_argv(["herdr", "--remote"])) + self.assertIsNone(parse_remote_argv(["herdr", "--remote", "-bad"])) + + def test_cmdline_is_nul_delimited(self) -> None: + self.assertEqual( + decode_cmdline(b"/bin/herdr\0--remote\0host name\0"), + ["/bin/herdr", "--remote", "host name"], + ) + + +class DiscoveryTests(unittest.TestCase): + def test_local_discovery_returns_only_running_sessions(self) -> None: + payload = { + "sessions": [ + {"name": "default", "running": True}, + {"name": "work", "running": True}, + {"name": "old", "running": False}, + ] + } + + def fake_run(*args, **kwargs): # type: ignore[no-untyped-def] + return subprocess.CompletedProcess(args[0], 0, json.dumps(payload), "") + + report = discover_local_sessions("herdr", run=fake_run) + self.assertEqual( + [endpoint.key for endpoint in report.endpoints], + ["local:default", "local:work"], + ) + self.assertEqual(report.warnings, []) + + def test_invalid_local_json_is_a_warning(self) -> None: + def fake_run(*args, **kwargs): # type: ignore[no-untyped-def] + return subprocess.CompletedProcess(args[0], 0, "not json", "") + + report = discover_local_sessions("herdr", run=fake_run) + self.assertEqual(report.endpoints, []) + self.assertEqual(len(report.warnings), 1) + + def test_proc_discovery_reads_exact_argv_and_deduplicates(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + commands = { + "101": ["/usr/bin/herdr", "--remote", "orangepi"], + "102": ["herdr", "--remote=orangepi"], + "103": ["herdr", "--remote", "prod", "--session", "work"], + "104": ["bash", "-c", "herdr --remote ignored"], + } + for process_id, argv in commands.items(): + process = root / process_id + process.mkdir() + (process / "cmdline").write_bytes( + b"\0".join(os.fsencode(argument) for argument in argv) + b"\0" + ) + + report = discover_remote_processes(root, user_id=os.getuid()) + + self.assertEqual( + [endpoint.key for endpoint in report.endpoints], + ["remote:orangepi:default", "remote:prod:work"], + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_remotes.py b/tests/test_remotes.py new file mode 100644 index 0000000..928cbf9 --- /dev/null +++ b/tests/test_remotes.py @@ -0,0 +1,121 @@ +from __future__ import annotations + +from pathlib import Path +import stat +import tempfile +import time +import unittest + +from herdr_hub.remotes import ( + RegistryError, + RemoteEntry, + RemoteQueryPool, + RemoteRegistry, + RemoteTargetError, + parse_remote_session_result, + remote_session_command, + ssh_destination_arguments, +) + + +class RemoteRegistryTests(unittest.TestCase): + def test_round_trip_update_remove_and_permissions(self) -> None: + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "config" / "remotes.json" + registry = RemoteRegistry(path) + self.assertEqual(registry.load(), []) + self.assertTrue(registry.upsert(RemoteEntry("orangepi"))) + self.assertFalse(registry.upsert(RemoteEntry("orangepi"))) + self.assertTrue(registry.upsert(RemoteEntry("prod", "server"))) + self.assertEqual( + registry.load(), + [RemoteEntry("orangepi"), RemoteEntry("prod", "server")], + ) + self.assertEqual(stat.S_IMODE(path.stat().st_mode), 0o600) + self.assertTrue(registry.remove("orangepi")) + self.assertFalse(registry.remove("missing")) + self.assertEqual(registry.load(), [RemoteEntry("prod", "server")]) + + def test_corrupt_registry_is_not_overwritten(self) -> None: + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "remotes.json" + path.write_text("not json", encoding="utf-8") + registry = RemoteRegistry(path) + with self.assertRaises(RegistryError): + registry.upsert(RemoteEntry("orangepi")) + self.assertEqual(path.read_text(encoding="utf-8"), "not json") + + +class RemoteSshTests(unittest.TestCase): + def test_alias_and_url_targets_become_safe_ssh_arguments(self) -> None: + self.assertEqual(ssh_destination_arguments("orangepi"), ["orangepi"]) + self.assertEqual( + ssh_destination_arguments("ssh://me@example.com:2222"), + ["-p", "2222", "-l", "me", "example.com"], + ) + command = remote_session_command("/usr/bin/ssh", RemoteEntry("host name")) + self.assertEqual(command[0], "/usr/bin/ssh") + self.assertIn("host name", command) + self.assertNotIn("sh", command[:1]) + with self.assertRaises(RemoteTargetError): + ssh_destination_arguments("ssh://me:secret@example.com") + with self.assertRaises(RemoteTargetError): + ssh_destination_arguments("ssh://example.com/a/path") + + def test_remote_session_json_keeps_running_and_stopped(self) -> None: + result = parse_remote_session_result( + RemoteEntry("orangepi"), + 0, + '{"sessions":[' + '{"name":"old","running":false},' + '{"name":"default","running":true}]}', + "", + ) + self.assertIsNone(result.error) + self.assertEqual( + [(session.name, session.running) for session in result.sessions], + [("default", True), ("old", False)], + ) + + def test_query_pool_reports_results_and_spawn_failures(self) -> None: + with tempfile.TemporaryDirectory() as directory: + fake_ssh = Path(directory) / "ssh" + fake_ssh.write_text( + "#!/bin/sh\n" + "printf '%s\\n' '{\"sessions\":[{\"name\":\"work\",\"running\":true}]}'\n", + encoding="utf-8", + ) + fake_ssh.chmod(0o755) + pool = RemoteQueryPool(str(fake_ssh), query_timeout=1) + self.addCleanup(pool.close) + pool.enqueue(RemoteEntry("orangepi")) + results = [] + deadline = time.monotonic() + 2 + while not results and time.monotonic() < deadline: + results = pool.poll() + time.sleep(0.01) + self.assertEqual(results[0].sessions[0].name, "work") + + missing = RemoteQueryPool(str(Path(directory) / "missing")) + self.addCleanup(missing.close) + missing.enqueue(RemoteEntry("orangepi")) + failure = missing.poll() + self.assertEqual(len(failure), 1) + self.assertIn("cannot start", failure[0].error or "") + + def test_query_pool_times_out_without_blocking_other_results(self) -> None: + with tempfile.TemporaryDirectory() as directory: + slow_ssh = Path(directory) / "ssh" + slow_ssh.write_text("#!/bin/sh\nexec sleep 5\n", encoding="utf-8") + slow_ssh.chmod(0o755) + pool = RemoteQueryPool(str(slow_ssh), query_timeout=0.05) + self.addCleanup(pool.close) + pool.enqueue(RemoteEntry("slow")) + time.sleep(0.08) + results = pool.poll() + self.assertEqual(len(results), 1) + self.assertIn("timed out", results[0].error or "") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_selector.py b/tests/test_selector.py new file mode 100644 index 0000000..37e5e7a --- /dev/null +++ b/tests/test_selector.py @@ -0,0 +1,181 @@ +from __future__ import annotations + +import os +from pathlib import Path +import pty +import select +import signal +import tempfile +import time +import unittest + +from herdr_hub.remotes import RemoteEntry, RemoteQueryResult, RemoteRegistry, RemoteSession +from herdr_hub.selector import ( + CatalogWindow, + PickerRow, + SessionSelector, + SelectorResult, + build_main_rows, + filter_rows, +) + + +class SelectorCatalogTests(unittest.TestCase): + def test_rows_merge_open_running_stopped_and_new_sessions(self) -> None: + windows = [ + CatalogWindow( + "@1", + "local:default", + "L default", + False, + True, + kind="local", + session="default", + ), + CatalogWindow( + "@2", + "remote:orangepi:work", + "R orangepi/work", + True, + False, + kind="remote", + target="orangepi", + session="work", + ), + ] + result = RemoteQueryResult( + RemoteEntry("orangepi"), + sessions=( + RemoteSession("default", True), + RemoteSession("old", False), + RemoteSession("work", True), + ), + ) + rows = build_main_rows( + windows, + [RemoteEntry("orangepi")], + {"orangepi": result}, + set(), + local_sessions=["default", "draft"], + ) + by_key = {row.key: row for row in rows} + self.assertEqual(by_key["window:@1"].main, "▶ L default") + self.assertEqual(by_key["window:@2"].main, "× R orangepi/work") + self.assertIn("remote:orangepi:default", by_key) + self.assertIn("remote:orangepi:old", by_key) + self.assertIn("new:orangepi", by_key) + self.assertEqual(by_key["local:draft"].main, "○ draft") + self.assertEqual(by_key["close-active"].window_id, "@1") + + def test_loading_error_and_filter_rows(self) -> None: + entry = RemoteEntry("prod") + loading = build_main_rows([], [entry], {}, {"prod"}) + self.assertTrue(any(row.kind == "loading" for row in loading)) + failed = build_main_rows( + [], + [entry], + {"prod": RemoteQueryResult(entry, error="offline")}, + set(), + ) + self.assertTrue(any(row.kind == "retry" for row in failed)) + matches = filter_rows(failed, "offline") + self.assertEqual([row.kind for row in matches], ["retry"]) + + +class SelectorActionTests(unittest.TestCase): + def test_dead_window_reconnect_action_and_new_session_action(self) -> None: + with tempfile.TemporaryDirectory() as directory: + registry = RemoteRegistry(Path(directory) / "remotes.json") + registry.upsert(RemoteEntry("orangepi", "server")) + selector = SessionSelector( + [CatalogWindow("@2", "key", "remote", True, False)], + registry, + str(Path(directory) / "missing-ssh"), + ) + self.addCleanup(selector.pool.close) + window_result = selector._activate( # type: ignore[attr-defined] + next(row for row in selector._rows() if row.kind == "window") # type: ignore[attr-defined] + ) + self.assertEqual(window_result, SelectorResult.window("@2", dead=True)) + + close_result = selector._activate( # type: ignore[attr-defined] + PickerRow("close", "close", "close", window_id="@2") + ) + self.assertEqual(close_result, SelectorResult.close("@2")) + + local_result = selector._activate( # type: ignore[attr-defined] + PickerRow("local", "local", "local", session="draft") + ) + self.assertEqual(local_result, SelectorResult.local("draft")) + + selector.view = "input-new" + selector.input_target = "orangepi" + new_result = selector._finish_new("work") # type: ignore[attr-defined] + self.assertEqual( + new_result, + SelectorResult.remote("orangepi", "work", "server"), + ) + self.assertIsNone(selector._finish_new("bad name")) # type: ignore[attr-defined] + + @unittest.skipUnless(hasattr(pty, "fork"), "a Unix PTY is required") + def test_curses_selector_handles_single_click_and_esc(self) -> None: + with tempfile.TemporaryDirectory() as directory: + pid, descriptor = pty.fork() + if pid == 0: + os.environ["TERM"] = "xterm-256color" + selector = SessionSelector( + [], + RemoteRegistry(Path(directory) / "remotes.json"), + "/missing/ssh", + ) + try: + result = selector.run() + os._exit(0 if result is None else 2) + except BaseException: + os._exit(3) + + status: int | None = None + try: + deadline = time.monotonic() + 3 + output = b"" + while time.monotonic() < deadline and b"Sessions" not in output: + readable, _, _ = select.select([descriptor], [], [], 0.1) + if readable: + output += os.read(descriptor, 4096) + self.assertIn(b"Sessions", output) + # Empty registries render "+ 添加远端" as the first selectable + # row at terminal row 5. A single SGR mouse click must activate it. + os.write(descriptor, b"\x1b[<0;5;5M\x1b[<0;5;5m") + opened_input = b"" + input_deadline = time.monotonic() + 2 + while time.monotonic() < input_deadline and b"SSH target" not in opened_input: + readable, _, _ = select.select([descriptor], [], [], 0.1) + if readable: + opened_input += os.read(descriptor, 4096) + self.assertIn(b"SSH target", opened_input) + os.write(descriptor, b"\x1b\x1b") + while time.monotonic() < deadline: + completed, wait_status = os.waitpid(pid, os.WNOHANG) + if completed == pid: + status = wait_status + break + time.sleep(0.05) + self.assertIsNotNone(status, "selector did not exit after Esc") + assert status is not None + self.assertTrue(os.WIFEXITED(status)) + self.assertEqual(os.WEXITSTATUS(status), 0) + finally: + try: + os.close(descriptor) + except OSError: + pass + if status is None: + try: + os.kill(pid, signal.SIGKILL) + except ProcessLookupError: + pass + os.waitpid(pid, 0) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_tmux_integration.py b/tests/test_tmux_integration.py new file mode 100644 index 0000000..444ba2d --- /dev/null +++ b/tests/test_tmux_integration.py @@ -0,0 +1,206 @@ +from __future__ import annotations + +import os +from pathlib import Path +import shutil +import tempfile +import time +import unittest +import uuid + +from herdr_hub.cli import ( + BOOTSTRAP_WINDOW_NAME, + Endpoint, + HubError, + TmuxHub, + WINDOW_SHORT_LABEL_OPTION, +) + + +@unittest.skipUnless(shutil.which("tmux"), "tmux is required") +class TmuxIntegrationTests(unittest.TestCase): + def setUp(self) -> None: + self.temporary_directory = tempfile.TemporaryDirectory() + self.addCleanup(self.temporary_directory.cleanup) + root = Path(self.temporary_directory.name) + self.fake_herdr = root / "herdr" + self.fake_herdr.write_text( + "#!/bin/sh\n" + "if [ \"${1:-}\" = \"--session\" ] && [ \"${2:-}\" = \"dead\" ]; then\n" + " exit 0\n" + "fi\n" + "exec sleep 30\n", + encoding="utf-8", + ) + self.fake_herdr.chmod(0o755) + self.socket_name = f"herdr-hub-test-{os.getpid()}-{uuid.uuid4().hex[:8]}" + self.socket_path = Path( + os.environ.get("TMUX_TMPDIR", f"/tmp/tmux-{os.getuid()}") + ) / self.socket_name + self.hub = TmuxHub( + tmux_bin=shutil.which("tmux") or "tmux", + socket_name=self.socket_name, + session_name="hub", + config_file="/dev/null", + script_path="/dev/null", + ) + self.addCleanup(self._stop_hub) + + def _stop_hub(self) -> None: + try: + self.hub.stop() + except Exception: + pass + self.socket_path.unlink(missing_ok=True) + + def test_configures_isolated_server_and_adds_without_duplicates(self) -> None: + self.assertTrue(self.hub.ensure_server()) + endpoints = [Endpoint.local(), Endpoint.remote("orangepi")] + added = self.hub.add_missing(endpoints, str(self.fake_herdr)) + self.assertEqual(len(added), 2) + self.assertEqual(self.hub.add_missing(endpoints, str(self.fake_herdr)), []) + + windows = self.hub.managed_windows() + self.assertEqual( + {window.key for window in windows}, + {"local:default", "remote:orangepi:default"}, + ) + remote_window = next( + window for window in windows if window.key == "remote:orangepi:default" + ) + self.assertEqual(remote_window.endpoint, Endpoint.remote("orangepi")) + all_names = self.hub.run( + ["list-windows", "-t", "hub", "-F", "#{window_name}"] + ).stdout.splitlines() + self.assertNotIn(BOOTSTRAP_WINDOW_NAME, all_names) + + prefix = self.hub.run(["show-options", "-gv", "prefix"]).stdout.strip() + status = self.hub.run(["show-options", "-gv", "status"]).stdout.strip() + position = self.hub.run( + ["show-options", "-gv", "status-position"] + ).stdout.strip() + mouse = self.hub.run(["show-options", "-gv", "mouse"]).stdout.strip() + clipboard = self.hub.run( + ["show-options", "-sv", "set-clipboard"] + ).stdout.strip() + status_right = self.hub.run( + ["show-options", "-gv", "status-right"] + ).stdout.strip() + dead_format = self.hub.run( + ["show-options", "-wgv", "remain-on-exit-format"] + ).stdout.strip() + inactive_window_format = self.hub.run( + ["show-options", "-wgv", "window-status-format"] + ).stdout.strip() + current_window_format = self.hub.run( + ["show-options", "-wgv", "window-status-current-format"] + ).stdout.strip() + status_left_binding = self.hub.run( + ["list-keys", "-T", "root", "MouseDown1StatusLeft"] + ).stdout.strip() + close_binding = self.hub.run( + ["list-keys", "-T", "prefix", "&"] + ).stdout.strip() + self.assertEqual(prefix, "C-a") + self.assertEqual(status, "on") + self.assertEqual(position, "top") + self.assertEqual(mouse, "on") + self.assertEqual(clipboard, "on") + self.assertIn("连接已断开", status_right) + self.assertIn("C-a R", status_right) + self.assertIn("连接已断开", dead_format) + self.assertIn("Ctrl+A R", dead_format) + self.assertIn("client_width", inactive_window_format) + self.assertIn(WINDOW_SHORT_LABEL_OPTION, inactive_window_format) + self.assertNotIn("#I", inactive_window_format) + self.assertIn("client_width", current_window_format) + self.assertIn(WINDOW_SHORT_LABEL_OPTION, current_window_format) + self.assertIn("display-popup", status_left_binding) + self.assertIn("confirm-before", close_binding) + self.assertIn(" close --window ", close_binding) + + local_short_label = self.hub.run( + [ + "show-options", + "-wv", + "-t", + next(window.window_id for window in windows if window.key == "local:default"), + WINDOW_SHORT_LABEL_OPTION, + ] + ).stdout.strip() + remote_short_label = self.hub.run( + [ + "show-options", + "-wv", + "-t", + remote_window.window_id, + WINDOW_SHORT_LABEL_OPTION, + ] + ).stdout.strip() + rendered_local_short_label = self.hub.run( + [ + "display-message", + "-p", + "-t", + next(window.window_id for window in windows if window.key == "local:default"), + "#{=12:@herdr_hub_short_label}", + ] + ).stdout.strip() + self.assertEqual(local_short_label, "local") + self.assertEqual(remote_short_label, "orangepi") + self.assertEqual(rendered_local_short_label, "local") + + def test_dismissed_window_stays_hidden_until_explicitly_reopened(self) -> None: + self.hub.ensure_server() + endpoints = [Endpoint.local(), Endpoint.local("work")] + self.hub.add_missing(endpoints, str(self.fake_herdr)) + + default_window = self.hub.window_for_key("local:default") + assert default_window is not None + self.hub.dismiss_window(default_window.window_id) + self.assertEqual(self.hub.dismissed_keys(), {"local:default"}) + self.assertIsNone(self.hub.window_for_key("local:default")) + self.assertEqual(self.hub.add_missing(endpoints, str(self.fake_herdr)), []) + + reopened = self.hub.activate_endpoint(Endpoint.local(), str(self.fake_herdr)) + self.assertEqual(reopened.key, "local:default") + self.assertEqual(self.hub.dismissed_keys(), set()) + + work_window = self.hub.window_for_key("local:work") + assert work_window is not None + self.hub.dismiss_window(work_window.window_id) + with self.assertRaises(HubError): + self.hub.dismiss_window(reopened.window_id) + + def test_dead_window_is_retained_and_can_be_respawned(self) -> None: + self.hub.ensure_server() + window = self.hub.create_window(Endpoint.local("dead"), str(self.fake_herdr)) + self.hub.remove_bootstrap() + + deadline = time.monotonic() + 3 + dead = False + while time.monotonic() < deadline: + current = self.hub.window_for_key("local:dead") + dead = current.dead if current is not None else False + if dead: + break + time.sleep(0.05) + self.assertTrue(dead) + + rendered_status = self.hub.run( + [ + "display-message", + "-p", + "-t", + window.window_id, + "#{?pane_dead,× #{window_name} [断开],● #{window_name}}", + ] + ).stdout.strip() + self.assertEqual(rendered_status, f"× {window.name} [断开]") + + self.hub.reconnect(window.window_id) + self.assertIsNotNone(self.hub.window_for_key("local:dead")) + + +if __name__ == "__main__": + unittest.main()