"""Herdr command menu, updates, and conversation forks.""" from __future__ import annotations from concurrent.futures import ThreadPoolExecutor import curses from dataclasses import dataclass import json import os from pathlib import Path import shlex import subprocess import sys import tempfile import threading import time from typing import TYPE_CHECKING, Callable, Sequence import uuid from .remotes import ( RemoteEntry, clip_utf8_tail, decode_remote_output, parse_remote_session_result, remote_control_command, remote_session_command, ) from .text import wrap_display if TYPE_CHECKING: from .cli import Endpoint # A dedicated command inside the Hub's existing Ctrl+A prefix table. CONTROL_KEY = "h" CONTROL_SHORTCUT = "C-a h" ROUTING_ENV = ( "HERDR_ENV", "HERDR_SOCKET_PATH", "HERDR_CLIENT_SOCKET_PATH", "HERDR_SESSION", "HERDR_PANE_ID", "HERDR_TAB_ID", "HERDR_WORKSPACE_ID", "HERDR_ACTIVE_PANE_ID", "HERDR_ACTIVE_TAB_ID", "HERDR_ACTIVE_WORKSPACE_ID", ) class ControlError(RuntimeError): """A diagnostic safe to display inside the command menu.""" class ControlCancelled(RuntimeError): """A read-only control request cancelled while closing the popup.""" def update_failure(header: str, output: str) -> str: """Put the final diagnostic ahead of setup chatter in small popups.""" lines = [line.strip() for line in output.splitlines() if line.strip()] diagnostic = next((line for line in reversed(lines) if line.startswith("curl:")), "") if not diagnostic: diagnostic = next((line for line in reversed(lines) if not line.startswith("没有自动执行普通重启")), "无错误输出") return f"{header}\n{diagnostic}\n\n完整命令输出(可翻页):\n{output}" def object_field(data: dict, name: str) -> dict: value = data.get(name) if not isinstance(value, dict): raise ControlError(f"Herdr 返回缺少 {name} 对象(需要 Herdr 0.9+)") return value def string_field(data: dict, name: str) -> str: value = data.get(name) if not isinstance(value, str) or not value.strip() or "\0" in value: raise ControlError(f"Herdr 返回缺少有效 {name}(需要 Herdr 0.9+)") return value def result_object(payload: object) -> dict: if not isinstance(payload, dict): raise ControlError("Herdr 返回的 JSON 不是对象") if payload.get("error"): raise ControlError(f"Herdr: {json.dumps(payload['error'], ensure_ascii=False)}") return object_field(payload, "result") if "result" in payload else payload @dataclass(frozen=True) class ForkSource: workspace_id: str pane_id: str tab_id: str cwd: str kind: str session_id: str workspace_label: str existing_names: frozenset[str] @property def arguments(self) -> list[str]: if self.kind == "codex": return ["fork", self.session_id] return ["--resume", self.session_id, "--fork-session"] def parse_snapshot(payload: object) -> ForkSource: result = result_object(payload) snapshot = object_field(result, "snapshot") if "snapshot" in result else result workspace_id = string_field(snapshot, "focused_workspace_id") pane_id = string_field(snapshot, "focused_pane_id") panes = snapshot.get("panes") if not isinstance(panes, list): raise ControlError("Herdr snapshot 缺少 panes") matches = [p for p in panes if isinstance(p, dict) and p.get("pane_id") == pane_id] if len(matches) != 1: raise ControlError("找不到唯一的当前 Agent pane") pane = matches[0] if pane.get("workspace_id") != workspace_id: raise ControlError("当前 pane 与 Workspace 不一致,请重新打开菜单") kind = pane.get("agent") if not kind: raise ControlError("当前 pane 没有运行 Agent") if kind not in ("codex", "claude"): raise ControlError(f"暂不支持 Fork {kind};支持 Codex / Claude") session = pane.get("agent_session") if not isinstance(session, dict): raise ControlError( f"当前 {kind.capitalize()} 尚未向 Herdr 上报会话 ID,暂时无法 Fork。\n" f"在目标主机检查 herdr integration status;确认集成已加载后按 Ctrl+R 重读。" ) if session.get("kind") != "id" or session.get("agent") != kind: raise ControlError("当前 Agent 尚未报告可 Fork 的 conversation ID") session_id = string_field(session, "value") if session_id.startswith("-") or any(ord(c) < 32 for c in session_id): raise ControlError("Agent conversation ID 无效") if kind == "codex": try: uuid.UUID(session_id) except ValueError as error: raise ControlError("Codex conversation ID 不是有效 UUID") from error agents = snapshot.get("agents", []) if not isinstance(agents, list): raise ControlError("Herdr snapshot 的 agents 无效") names = frozenset(a["name"] for a in agents if isinstance(a, dict) and isinstance(a.get("name"), str)) workspaces = snapshot.get("workspaces", []) label = workspace_id if isinstance(workspaces, list): for workspace in workspaces: if isinstance(workspace, dict) and workspace.get("workspace_id") == workspace_id: label = workspace.get("label") or workspace_id return ForkSource(workspace_id, pane_id, string_field(pane, "tab_id"), string_field(pane, "cwd"), kind, session_id, str(label), names) class HerdrControl: def __init__(self, endpoint: Endpoint, herdr_bin: str, ssh_bin: str = "ssh") -> None: self.endpoint = endpoint self.herdr_bin = herdr_bin self.ssh_bin = ssh_bin self.platform: str | None = None def _execute( self, command: Sequence[str], timeout: float, cancel_event: threading.Event | None = None, progress: Callable[[str], None] | None = None, ) -> subprocess.CompletedProcess: environment = {k: v for k, v in os.environ.items() if k not in ROUTING_ENV} if progress is not None: # A file avoids pipe deadlocks and lets us show output before exit, # including progress indicators that do not end in a newline. try: with tempfile.TemporaryFile() as output: process = subprocess.Popen(command, stdout=output, stderr=subprocess.STDOUT, stdin=subprocess.DEVNULL, env=environment) deadline = time.monotonic() + timeout started = time.monotonic() last_output = started last_notice = started position = 0 captured = b"" try: while True: finished = process.poll() is not None chunk = os.pread(output.fileno(), 65536, position) if chunk: position += len(chunk) # Bound memory at a character boundary; the # chunk may also end mid-sequence, so decode # it as a still-growing stream. captured = clip_utf8_tail(captured + chunk, 65536) progress( decode_remote_output(captured, partial=True).replace("\r", "\n") ) last_output = time.monotonic() now = time.monotonic() if not finished and now - last_notice >= 1: progress( decode_remote_output(captured, partial=True).replace("\r", "\n") + f"\n已执行 {int(now - started)} 秒 · {int(now - last_output)} 秒无新输出" ) last_notice = now if finished and not chunk: return subprocess.CompletedProcess(command, process.returncode, captured, b"") if time.monotonic() >= deadline: raise ControlError( f"命令超时({timeout:g}s);操作可能已完成,请检查 Herdr\n" + decode_remote_output(captured) ) if not chunk: time.sleep(0.1) finally: if process.poll() is None: process.kill() process.wait() except OSError as error: raise ControlError(f"无法执行 Herdr 命令: {error}") from error if cancel_event is not None and cancel_event.is_set(): raise ControlCancelled() if cancel_event is not None: try: process = subprocess.Popen( command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, stdin=subprocess.DEVNULL, env=environment, ) except OSError as error: raise ControlError(f"无法执行 Herdr 命令: {error}") from error deadline = time.monotonic() + timeout while True: if cancel_event.wait(0.025): process.terminate() try: process.communicate(timeout=0.25) except subprocess.TimeoutExpired: process.kill() process.communicate() raise ControlCancelled() remaining = deadline - time.monotonic() if remaining <= 0: process.kill() process.communicate() raise ControlError( f"命令超时({timeout:g}s);操作可能已完成,请检查 Herdr" ) try: stdout, stderr = process.communicate( timeout=min(0.05, remaining) ) except subprocess.TimeoutExpired: continue return subprocess.CompletedProcess( command, process.returncode, stdout, stderr, ) try: return subprocess.run(command, capture_output=True, check=False, timeout=timeout, env=environment, stdin=subprocess.DEVNULL) except subprocess.TimeoutExpired as error: raise ControlError(f"命令超时({timeout:g}s);操作可能已完成,请检查 Herdr") from error except OSError as error: raise ControlError(f"无法执行 Herdr 命令: {error}") from error def _remote_platform( self, entry: RemoteEntry, cancel_event: threading.Event | None = None, ) -> str: if self.platform is None: for fallback in (False, True): result = self._execute( remote_session_command( self.ssh_bin, entry, posix_fallback=fallback, ), 8, cancel_event, ) if result.returncode != 127 or fallback: break missing = decode_remote_output(result.stderr) + decode_remote_output(result.stdout) if result.returncode != 0 and any(marker in missing.lower() for marker in ( "not recognized", "不是内部或外部命令", "commandnotfoundexception", )): # A Windows shell can lack Herdr on PATH while the standard # per-user executable is installed. This probe is read-only. result = self._execute(remote_control_command(self.ssh_bin, entry, ["session", "list", "--json"], platform="windows"), 8, cancel_event) parsed = parse_remote_session_result(entry, result.returncode, decode_remote_output(result.stdout), decode_remote_output(result.stderr)) if parsed.error: raise ControlError(parsed.error) if parsed.platform not in ("windows", "posix"): raise ControlError("无法从远端 Session 路径识别平台,不能安全执行控制命令") self.platform = parsed.platform return self.platform def call( self, arguments: Sequence[str], *, timeout: float = 10, cancel_event: threading.Event | None = None, ) -> dict: argv = ["--session", self.endpoint.session, *arguments] result = self._execute(self._command(argv, cancel_event), timeout, cancel_event) stdout, stderr = decode_remote_output(result.stdout), decode_remote_output(result.stderr) if result.returncode != 0: raise ControlError((stderr or stdout or f"Herdr exited {result.returncode}").strip()) try: return result_object(json.loads(stdout)) except json.JSONDecodeError as error: raise ControlError(f"Herdr 返回无效 JSON: {error}") from error def _command( self, argv: Sequence[str], cancel_event: threading.Event | None = None, ) -> list[str]: if self.endpoint.kind == "remote": entry = RemoteEntry(self.endpoint.target or "", self.endpoint.remote_keybindings, self.endpoint.remote_transport) return remote_control_command(self.ssh_bin, entry, argv, platform=self._remote_platform(entry, cancel_event), unset_environment=ROUTING_ENV) return [self.herdr_bin, *argv] def update(self, progress: Callable[[str], None] | None = None) -> str: """Run the host updater once; its output is text, not control API JSON.""" if progress: progress("正在连接更新目标…") command = self._command(["update"]) if self.endpoint.kind == "remote": command = self._update_proxy(command, progress) result = self._execute(command, timeout=300, progress=progress) stdout, stderr = decode_remote_output(result.stdout), decode_remote_output(result.stderr) output = "\n".join(part.strip() for part in (stdout, stderr) if part.strip()) if result.returncode != 0: raise ControlError(update_failure(f"更新失败(退出码 {result.returncode})", output)) return "更新命令已完成。重新连接可使用更新后的客户端。\n" + output def _update_proxy( self, command: list[str], progress: Callable[[str], None] | None, ) -> list[str]: """Load proxy settings on the remote host, never from the Hub host.""" if self.platform == "windows": entry = RemoteEntry(self.endpoint.target or "", self.endpoint.remote_keybindings, self.endpoint.remote_transport) return remote_control_command( self.ssh_bin, entry, ["update"], platform="windows", unset_environment=ROUTING_ENV, load_proxy_environment=True, ) if self.platform != "posix": return command if progress: progress("正在检测远端自身的代理配置…") helper = Path(__file__).with_name("proxy_env.py").read_text(encoding="utf-8") result = list(command) result[-1] = ( "if command -v python3 >/dev/null 2>&1; then exec python3 -u -c " + shlex.quote(helper) + " " + shlex.quote(command[-1]) + "; else echo '远端没有 python3,沿用远端现有代理环境及 curl 配置'; " + command[-1] + "; fi" ) return result def update_and_restart(self, progress: Callable[[str], None]) -> str: """Update and verify a live server handoff on the selected POSIX session.""" helper = str(Path(__file__).with_name("restart.py")) if self.endpoint.kind == "local": command = [sys.executable, helper, self.herdr_bin, self.endpoint.session] else: entry = RemoteEntry(self.endpoint.target or "", self.endpoint.remote_keybindings, self.endpoint.remote_transport) progress("正在检查远端热切换支持…") if self._remote_platform(entry) != "posix": raise ControlError("Windows 不支持热切换;请使用“更新 Herdr”,再安排普通重启。") # Reuse the SSH destination/options, but run the standalone helper # on the host that owns the Unix socket. No socket crosses SSH. command = remote_control_command(self.ssh_bin, entry, [], platform="posix") script = Path(helper).read_text(encoding="utf-8") command[-1] = ( "".join(f"unset {key}; " for key in ROUTING_ENV) + 'if [ -x "$HOME/.local/bin/herdr" ]; then hub_herdr="$HOME/.local/bin/herdr"; ' 'else hub_herdr=$(command -v herdr) || exit 127; fi; ' + "exec python3 -u -c " + shlex.quote(script) + ' "$hub_herdr" ' + shlex.quote(self.endpoint.session) ) command = self._update_proxy(command, progress) result = self._execute(command, timeout=400, progress=progress) output = decode_remote_output(result.stdout) if result.returncode != 0: raise ControlError(update_failure("更新/热切换失败;未自动执行普通重启。", output)) return output def source( self, cancel_event: threading.Event | None = None, ) -> ForkSource: return parse_snapshot( self.call(["api", "snapshot"], cancel_event=cancel_event) ) def fork(self, source: ForkSource, progress: Callable[[str], None] = lambda _: None) -> str: name = f"{source.kind}-fork-{uuid.uuid4().hex[:12]}" while name in source.existing_names: name = f"{source.kind}-fork-{uuid.uuid4().hex[:12]}" progress("正在同一 Workspace 创建新 Tab…") # Never delete either tab: a timed-out request may already have succeeded. tab_id = "未知(请检查 Workspace)" try: created = self.call(["tab", "create", "--workspace", source.workspace_id, "--cwd", source.cwd, "--no-focus"]) tab_id = string_field(object_field(created, "tab"), "tab_id") root_pane = object_field(created, "root_pane") pane_id = string_field(root_pane, "pane_id") if tab_id == source.tab_id or pane_id == source.pane_id: raise ControlError("返回的 Tab / root pane 与原 Tab / pane 相同,已拒绝启动") if (root_pane.get("workspace_id") != source.workspace_id or root_pane.get("tab_id") != tab_id): raise ControlError("新 root pane 的 Workspace / Tab 不一致,已拒绝启动") progress(f"正在启动 {source.kind} Fork,等待就绪…") started = self.call(["agent", "start", name, "--kind", source.kind, "--pane", pane_id, "--", *source.arguments], timeout=40) agent = object_field(started, "agent") if (agent.get("pane_id") != pane_id or agent.get("agent") != source.kind or agent.get("interactive_ready") is False or agent.get("launch_pending") is True or agent.get("agent_status") in ("blocked", "error")): raise ControlError("新 Agent 未确认就绪: " + json.dumps(started, ensure_ascii=False)) progress("Fork 已就绪,正在切换 Tab…") self.call(["tab", "focus", tab_id]) except ControlError as error: raise ControlError(f"{error}\n原 Tab 保留;新增 Tab {tab_id} 未自动关闭。请检查后再重试。") from error return tab_id class ControlMenu: """Small action menu; all subprocess work stays off the curses thread.""" ACTIONS = ( ( "fork", "Fork 当前会话到新 Tab", "保留当前 Tab,在同一 Workspace 和工作目录继续", ), ("update", "更新 Herdr", "更新当前主机的 Herdr(使用已配置的更新渠道)"), ("update-restart", "更新并热重启 Herdr", "实验性热切换,保留进程;成功后自动重连当前 Tab"), ) def __init__(self, control: HerdrControl) -> None: self.control = control self.reconnect_requested = False def run(self, screen) -> int: # type: ignore[no-untyped-def] import queue messages: queue.SimpleQueue[str] = queue.SimpleQueue() screen.keypad(True) screen.timeout(100) try: curses.set_escdelay(25) curses.curs_set(0) curses.mousemask(curses.BUTTON1_CLICKED | curses.BUTTON1_PRESSED) except (AttributeError, curses.error): pass source: ForkSource | None = None source_error = "" status = "" selected = 0 update_attempted = False pending_update: str | None = None status_offset = 0 target = "本机" if self.control.endpoint.kind == "local" else self.control.endpoint.target cancel_loading = threading.Event() with ThreadPoolExecutor(max_workers=1) as executor: future = executor.submit(self.control.source, cancel_loading) phase = "load" while True: while not messages.empty(): status = messages.get() if future is not None and future.done(): # The worker is finished: discard queued progress before # publishing its final result so stale output cannot replace it. while not messages.empty(): messages.get() status_offset = 0 try: value = future.result() if phase == "load": source = value source_error = "" status = "" elif phase == "update": status = value elif phase == "update-restart": self.reconnect_requested = True return 0 else: return 0 except ControlCancelled: if not pending_update: return 0 except Exception as error: status = str(error) if phase == "load": source_error = status if phase == "fork": source = None # Avoid accidental duplicate forks after ambiguous failure. source_error = status future = None if pending_update and future is None: phase = pending_update pending_update = None update_attempted = True method = self.control.update if phase == "update" else self.control.update_and_restart future = executor.submit(method, messages.put) status = f"正在更新 {target} 的 Herdr…" status_offset = 0 screen.erase() height, width = screen.getmaxyx() compact = height < 16 if source is not None: context = ( f"当前:{source.kind.capitalize()} · " f"{source.workspace_label}" ) elif future is not None and phase == "load": context = "正在读取当前页面…" else: context = "无法读取当前页面" if not compact: self._add(screen, 1, 2, context, curses.A_BOLD) heading = context if compact and selected == 0 else f"更新目标:{target}" self._add(screen, 0 if compact else 2, 2, heading, curses.A_BOLD) action_top = 1 if compact else 3 action_stride = 1 if compact else 2 for index, (action, label, detail) in enumerate(self.ACTIONS): y = action_top + index * action_stride enabled = ( source is not None and future is None if action == "fork" else not update_attempted and (future is None or phase == "load") ) attribute = ( curses.A_REVERSE if enabled and index == selected else curses.A_NORMAL if enabled else curses.A_DIM ) if action == "fork" and source is None: label += "(读取中)" if phase == "load" and future is not None else "(不可用)" self._add(screen, y, 2, f" {label} ", attribute) if not compact: self._add(screen, y + 1, 4, detail, curses.A_DIM) status_top = action_top + len(self.ACTIONS) * action_stride + (0 if compact else 1) status_lines = wrap_display(status, max(1, width - 5)) if status else [] visible_lines = max(1, height - status_top - 1) status_offset = min(status_offset, max(0, len(status_lines) - visible_lines)) if future is not None and phase in ("update", "update-restart"): status_offset = max(0, len(status_lines) - visible_lines) for offset, line in enumerate(status_lines[status_offset:status_offset + visible_lines]): self._add(screen, status_top + offset, 2, line, curses.A_DIM) footer = ( "操作执行中,请等待结果" if future is not None and phase != "load" else "Enter 执行 · Ctrl+R 重读 · PgUp/PgDn 结果 · Esc 关闭" ) self._add(screen, height - 1, 2, footer, curses.A_DIM) screen.refresh() key = screen.getch() if key == 27: if future is not None and phase == "load": cancel_loading.set() return 0 if future is None: return 0 status = "操作执行中,请等待结果;关闭菜单不会取消已发送的操作。" elif key == 18 and future is None: source = None source_error = "" cancel_loading = threading.Event() future = executor.submit(self.control.source, cancel_loading) phase = "load" status = "正在重新读取当前页面…" status_offset = 0 elif key in (curses.KEY_PPAGE, curses.KEY_NPAGE): status_offset = max(0, status_offset + ( visible_lines if key == curses.KEY_NPAGE else -visible_lines )) elif key in (curses.KEY_UP, curses.KEY_DOWN): selected = (selected + (1 if key == curses.KEY_DOWN else -1)) % len(self.ACTIONS) else: execute = key in (10, 13, curses.KEY_ENTER) if key == curses.KEY_MOUSE: try: _, x, y, _, state = curses.getmouse() action_index = (y - action_top) // action_stride execute = ( 1 <= x < width - 1 and 0 <= action_index < len(self.ACTIONS) and bool( state & ( curses.BUTTON1_CLICKED | curses.BUTTON1_PRESSED ) ) ) if execute: selected = action_index except curses.error: execute = False if execute: action = self.ACTIONS[selected][0] if action in ("update", "update-restart") and not update_attempted and phase == "load" and future is not None: pending_update = action cancel_loading.set() status = "正在准备更新…" elif action == "fork" and source is not None and future is None: phase = "fork" future = executor.submit(self.control.fork, source, messages.put) status = "正在 Fork…" status_offset = 0 elif action == "fork": if future is None: status = source_error or "尚未读取 Fork 会话信息,请按 Ctrl+R 重读。" status_offset = 0 elif phase == "load": status = "正在读取 Fork 会话信息,请稍候再执行。" elif action in ("update", "update-restart") and not update_attempted and future is None: phase = action update_attempted = True method = self.control.update if action == "update" else self.control.update_and_restart future = executor.submit(method, messages.put) status = f"正在更新 {target} 的 Herdr,请等待…" status_offset = 0 @staticmethod def _add( screen, y: int, x: int, value: 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 chunks = wrap_display(value, max(1, width - x - 1)) try: screen.addstr(y, x, chunks[0] if chunks else "", attribute) except curses.error: pass