458 lines
19 KiB
Python
458 lines
19 KiB
Python
"""Herdr command menu and non-destructive conversation forks."""
|
||
from __future__ import annotations
|
||
|
||
from concurrent.futures import ThreadPoolExecutor
|
||
import curses
|
||
from dataclasses import dataclass
|
||
import json
|
||
import os
|
||
import subprocess
|
||
import threading
|
||
import time
|
||
from typing import TYPE_CHECKING, Callable, Sequence
|
||
import uuid
|
||
import unicodedata
|
||
|
||
from .remotes import (
|
||
RemoteEntry, decode_remote_output, parse_remote_session_result,
|
||
remote_control_command, remote_session_command,
|
||
)
|
||
|
||
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_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 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 = object_field(pane, "agent_session")
|
||
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,
|
||
) -> subprocess.CompletedProcess:
|
||
environment = {k: v for k, v in os.environ.items() if k not in ROUTING_ENV}
|
||
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]
|
||
if self.endpoint.kind == "remote":
|
||
entry = RemoteEntry(self.endpoint.target or "", self.endpoint.remote_keybindings,
|
||
self.endpoint.remote_transport)
|
||
command = remote_control_command(self.ssh_bin, entry, argv,
|
||
platform=self._remote_platform(entry, cancel_event),
|
||
unset_environment=ROUTING_ENV)
|
||
else:
|
||
command = [self.herdr_bin, *argv]
|
||
result = self._execute(command, 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 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
|
||
|
||
|
||
def display_chunks(text: str, width: int) -> list[str]:
|
||
chunks: list[str] = []
|
||
line = ""
|
||
cells = 0
|
||
for character in text:
|
||
if character == "\n":
|
||
chunks.append(line)
|
||
line, cells = "", 0
|
||
continue
|
||
if ord(character) < 32 or ord(character) == 127:
|
||
character = " "
|
||
size = 0 if unicodedata.combining(character) else (2 if unicodedata.east_asian_width(character) in ("W", "F") else 1)
|
||
if cells + size > width and line:
|
||
chunks.append(line)
|
||
line, cells = "", 0
|
||
line += character
|
||
cells += size
|
||
chunks.append(line)
|
||
return chunks
|
||
|
||
|
||
class ControlMenu:
|
||
"""Small action menu; all subprocess work stays off the curses thread."""
|
||
ACTIONS = (
|
||
(
|
||
"fork",
|
||
"Fork 当前会话到新 Tab",
|
||
"保留当前 Tab，在同一 Workspace 和工作目录继续",
|
||
),
|
||
)
|
||
|
||
def __init__(self, control: HerdrControl) -> None:
|
||
self.control = control
|
||
|
||
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
|
||
status = ""
|
||
selected = 0
|
||
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():
|
||
try:
|
||
value = future.result()
|
||
if phase == "load":
|
||
source = value
|
||
status = ""
|
||
else:
|
||
return 0
|
||
except ControlCancelled:
|
||
return 0
|
||
except Exception as error:
|
||
status = str(error)
|
||
if phase == "fork":
|
||
source = None # Avoid accidental duplicate forks after ambiguous failure.
|
||
future = None
|
||
screen.erase()
|
||
height, width = screen.getmaxyx()
|
||
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 = "无法读取当前页面"
|
||
self._add(screen, 1, 2, context, curses.A_BOLD)
|
||
action_top = 3
|
||
for index, (_, label, detail) in enumerate(self.ACTIONS):
|
||
y = action_top + index * 2
|
||
enabled = source is not None and future is None
|
||
attribute = (
|
||
curses.A_REVERSE
|
||
if enabled and index == selected
|
||
else curses.A_NORMAL
|
||
if enabled
|
||
else curses.A_DIM
|
||
)
|
||
self._add(screen, y, 2, f" {label} ", attribute)
|
||
self._add(screen, y + 1, 4, detail, curses.A_DIM)
|
||
|
||
status_top = action_top + len(self.ACTIONS) * 2 + 1
|
||
for offset, line in enumerate(
|
||
display_chunks(status, max(1, width - 5)) if status else []
|
||
):
|
||
self._add(screen, status_top + offset, 2, line, curses.A_DIM)
|
||
|
||
footer = (
|
||
"操作执行中，请等待结果"
|
||
if future is not None and phase == "fork"
|
||
else "Enter 执行 · Esc 关闭"
|
||
)
|
||
self._add(screen, height - 2, 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 = "操作执行中；完成后可关闭（避免误以为已取消 Fork）。"
|
||
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) // 2
|
||
execute = (
|
||
1 <= x < width - 1
|
||
and 0 <= action_index < len(self.ACTIONS)
|
||
and y in {
|
||
action_top + action_index * 2,
|
||
action_top + action_index * 2 + 1,
|
||
}
|
||
and bool(
|
||
state
|
||
& (
|
||
curses.BUTTON1_CLICKED
|
||
| curses.BUTTON1_PRESSED
|
||
)
|
||
)
|
||
)
|
||
if execute:
|
||
selected = action_index
|
||
except curses.error:
|
||
execute = False
|
||
if execute and future is None and source is not None:
|
||
phase = "fork"
|
||
future = executor.submit(self.control.fork, source, messages.put)
|
||
status = "正在 Fork…"
|
||
|
||
@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 = display_chunks(value, max(1, width - x - 1))
|
||
try:
|
||
screen.addstr(y, x, chunks[0] if chunks else "", attribute)
|
||
except curses.error:
|
||
pass
|