from __future__ import annotations import argparse import base64 import curses import json import os import subprocess import sys import tempfile from pathlib import Path import threading import time import unittest from unittest.mock import Mock, patch from herdr_hub.cli import Endpoint, HubError, ManagedWindow, build_parser, command_control from herdr_hub.control import ( ControlCancelled, ControlError, ControlMenu, HerdrControl, parse_snapshot, update_failure, ) from herdr_hub.remotes import RemoteEntry, remote_control_command SESSION_ID = "0198f891-d531-745e-96f4-ef7d9092298b" def snapshot(kind="codex"): return {"id": "request", "result": {"snapshot": { "focused_workspace_id": "w1", "focused_pane_id": "w1:p1", "panes": [{"pane_id": "w1:p1", "workspace_id": "w1", "tab_id": "t1", "agent": kind, "cwd": "/tmp/project with spaces", "agent_session": { "agent": kind, "kind": "id", "value": SESSION_ID}}], "agents": [{"name": "existing-agent"}], "workspaces": [{"workspace_id": "w1", "label": "Project"}], }}} def completed(payload, code=0): return subprocess.CompletedProcess([], code, json.dumps(payload).encode(), b"") class SnapshotTests(unittest.TestCase): def test_missing_session_reports_identity_problem_not_version(self): payload = snapshot() del payload["result"]["snapshot"]["panes"][0]["agent_session"] with self.assertRaisesRegex(ControlError, "尚未向 Herdr 上报会话 ID") as caught: parse_snapshot(payload) self.assertNotIn("需要 Herdr 0.9", str(caught.exception)) def test_focused_pane_and_exact_native_arguments(self): for kind, args in (("codex", ["fork", SESSION_ID]), ("claude", ["--resume", SESSION_ID, "--fork-session"])): with self.subTest(kind=kind): source = parse_snapshot(snapshot(kind)) self.assertEqual(source.arguments, args) self.assertEqual(source.workspace_id, "w1") self.assertEqual(source.workspace_label, "Project") self.assertEqual(source.existing_names, {"existing-agent"}) def test_invalid_snapshots_fail_before_mutation(self): for field, value in (("agent", None), ("agent", "gemini"), ("cwd", None), ("workspace_id", "other"), ("agent_session", None), ("agent_session", {"agent": "codex", "kind": "id", "value": "wrong"}), ("agent_session", {"agent": "codex", "kind": "path", "value": "x"})): payload = snapshot() payload["result"]["snapshot"]["panes"][0][field] = value with self.subTest(field=field, value=value), self.assertRaises(ControlError): parse_snapshot(payload) for field in ("focused_workspace_id", "focused_pane_id", "panes"): payload = snapshot() del payload["result"]["snapshot"][field] with self.subTest(field=field), self.assertRaises(ControlError): parse_snapshot(payload) class ControlTests(unittest.TestCase): def test_windows_hot_restart_rejected_before_update(self): control = HerdrControl(Endpoint.remote("windows-host", remote_transport="ssh"), "unused") control.platform = "windows" with patch.object(control, "_execute") as execute: with self.assertRaisesRegex(ControlError, "Windows 不支持热切换"): control.update_and_restart(lambda _: None) execute.assert_not_called() def test_remote_hot_restart_runs_helper_on_exact_host(self): control = HerdrControl(Endpoint.remote("ssh://me@host:2222", "named"), "unused", "fake-ssh") control.platform = "posix" with patch.object(control, "_execute", return_value=subprocess.CompletedProcess([], 0, b"ok", b"")) as execute: control.update_and_restart(lambda _: None) command = execute.call_args.args[0] self.assertEqual(command[0], "fake-ssh") self.assertIn("host", command) self.assertIn("2222", command) self.assertIn("exec python3 -u -c", command[-1]) self.assertIn(' "$hub_herdr" named', command[-1]) self.assertIn("unset HERDR_ENV;", command[-1]) def test_streaming_shows_unterminated_output_before_process_exits(self): with tempfile.TemporaryDirectory() as directory: marker = Path(directory) / "seen" script = ( "import os, sys, time\n" "from pathlib import Path\n" "os.write(2, b'downloading 10%')\n" "deadline = time.monotonic() + 3\n" "while not Path(sys.argv[1]).exists():\n" " if time.monotonic() > deadline: sys.exit(9)\n" " time.sleep(.01)\n" "print(' done')\n" ) updates = [] def progress(value): updates.append(value) marker.touch() result = HerdrControl(Endpoint.local(), "unused")._execute( [sys.executable, "-c", script, str(marker)], 5, progress=progress ) self.assertEqual(result.returncode, 0) self.assertIn("downloading 10%", updates[0]) self.assertIn(b"done", result.stdout) def test_streaming_timeout_preserves_diagnostic(self): with self.assertRaisesRegex(ControlError, "waiting for network"): HerdrControl(Endpoint.local(), "unused")._execute( [sys.executable, "-c", "import os,time; os.write(2,b'waiting for network'); time.sleep(10)"], .3, progress=lambda _: None, ) def test_update_local_is_host_scoped_and_accepts_plain_text(self): result = subprocess.CompletedProcess([], 0, b"Already up to date", b"notice") with patch.dict(os.environ, {"HERDR_SESSION": "wrong"}), patch( "herdr_hub.control.subprocess.run", return_value=result ) as run: output = HerdrControl(Endpoint.local("named"), "/fake/herdr").update() self.assertIn("Already up to date", output) self.assertIn("notice", output) self.assertEqual(run.call_args.args[0], ["/fake/herdr", "update"]) self.assertEqual(run.call_args.kwargs["timeout"], 300) self.assertEqual(run.call_args.kwargs["stdin"], subprocess.DEVNULL) self.assertNotIn("HERDR_SESSION", run.call_args.kwargs["env"]) def test_update_remote_routes_once_on_both_platforms(self): for transport, path in (("native", "/home/me/sessions"), ("ssh", "C:\\Users\\me\\sessions")): listing = {"sessions": [{"name": "named", "running": True, "session_dir": path}]} failure = subprocess.CompletedProcess([], 1, b"download started", b"network error") with self.subTest(transport=transport), patch( "herdr_hub.control.subprocess.run", side_effect=[completed(listing), failure] ) as run: control = HerdrControl(Endpoint.remote("host", "named", remote_transport=transport), "unused", "fake-ssh") with self.assertRaisesRegex(ControlError, "network error"): control.update() self.assertEqual(run.call_count, 2) command = run.call_args.args[0] self.assertEqual(command[0], "fake-ssh") self.assertIn("host", command) script = command[-1] if transport == "ssh": script = base64.b64decode(script.split()[-1]).decode("utf-16le") self.assertIn("update", script) self.assertNotIn("--session", script) self.assertNotIn("--handoff", script) def test_update_timeout_is_not_replayed(self): with patch("herdr_hub.control.subprocess.run", side_effect=subprocess.TimeoutExpired("herdr", 300)) as run: with self.assertRaisesRegex(ControlError, "操作可能已完成"): HerdrControl(Endpoint.local(), "herdr").update() run.assert_called_once() def test_local_explicit_session_and_clean_environment(self): for session in ("default", "named"): with self.subTest(session=session), patch.dict(os.environ, { "HERDR_SOCKET_PATH": "/wrong/socket", "HERDR_SESSION": "wrong", "HERDR_PANE_ID": "wrong", "HERDR_CLIENT_SOCKET_PATH": "/wrong/client", }), patch("herdr_hub.control.subprocess.run", return_value=completed(snapshot())) as run: source = HerdrControl(Endpoint.local(session), "/fake/herdr").source() self.assertEqual(source.kind, "codex") self.assertEqual(run.call_args.args[0], ["/fake/herdr", "--session", session, "api", "snapshot"]) self.assertNotIn("HERDR_SOCKET_PATH", run.call_args.kwargs["env"]) self.assertNotIn("HERDR_SESSION", run.call_args.kwargs["env"]) def test_remote_native_and_ssh_route_to_target_named_server(self): for transport, path in (("native", "/home/user/.herdr/default"), ("ssh", "C:\\Users\\me\\.herdr\\default")): listing = {"sessions": [{"name": "named", "running": True, "session_dir": path}]} with self.subTest(transport=transport), patch("herdr_hub.control.subprocess.run", side_effect=[completed(listing), completed(snapshot())]) as run: control = HerdrControl(Endpoint.remote("ssh://user@host:2222", "named", "server", transport), "/unused", "fake-ssh") control.source() command = run.call_args.args[0] self.assertEqual(command[0], "fake-ssh") self.assertIn("2222", command) self.assertIn("host", command) script = command[-1] if transport == "ssh": script = base64.b64decode(script.split()[-1]).decode("utf-16le") self.assertIn("--session named api snapshot", script) self.assertIn("HERDR_SOCKET_PATH", script) def test_windows_missing_path_uses_read_only_fallback_before_mutation(self): missing = subprocess.CompletedProcess([], 1, b"", b"'herdr' is not recognized") listing = {"sessions": [{"name": "named", "running": True, "session_dir": "C:\\Users\\me\\sessions"}]} with patch("herdr_hub.control.subprocess.run", side_effect=[missing, completed(listing), completed(snapshot())]) as run: HerdrControl(Endpoint.remote("host", "named", remote_transport="ssh"), "unused").source() script = base64.b64decode(run.call_args_list[1].args[0][-1].split()[-1]).decode("utf-16le") self.assertIn("session list --json", script) self.assertIn(".local", script) self.assertEqual(len(run.call_args_list), 3) def test_success_creates_unfocused_tab_starts_then_focuses(self): for kind, readiness in (("codex", {"interactive_ready": True}), ("claude", {})): control = HerdrControl(Endpoint.local(), "fake") control.call = Mock(side_effect=[ {"tab": {"tab_id": "t2"}, "root_pane": {"pane_id": "w1:p2", "workspace_id": "w1", "tab_id": "t2"}}, {"agent": {"pane_id": "w1:p2", "agent": kind, **readiness}}, {}, ]) source = parse_snapshot(snapshot(kind)) self.assertEqual(control.fork(source), "t2") calls = [call.args[0] for call in control.call.call_args_list] self.assertEqual(calls[0], ["tab", "create", "--workspace", "w1", "--cwd", source.cwd, "--no-focus"]) self.assertEqual(calls[1][3:], ["--kind", kind, "--pane", "w1:p2", "--", *source.arguments]) self.assertRegex(calls[1][2], rf"^{kind}-fork-[a-f0-9]+$") self.assertEqual(calls[2], ["tab", "focus", "t2"]) def test_failure_never_closes_original_or_created_tab(self): for failure in (ControlError("agent_not_ready: blocked"), ControlError("timeout"), {"agent": {"interactive_ready": False}}, {"bad": "response"}): with self.subTest(failure=failure): control = HerdrControl(Endpoint.local(), "fake") control.call = Mock(side_effect=[ {"tab": {"tab_id": "t2"}, "root_pane": {"pane_id": "w1:p2", "workspace_id": "w1", "tab_id": "t2"}}, failure]) with self.assertRaisesRegex(ControlError, "新增 Tab t2 未自动关闭"): control.fork(parse_snapshot(snapshot())) self.assertEqual(len(control.call.call_args_list), 2) self.assertFalse(any("close" in c.args[0] or "focus" in c.args[0] for c in control.call.call_args_list)) def test_wrong_created_pane_never_starts_in_original_tab(self): for tab_id, pane_id, workspace in (("t1", "w1:p2", "w1"), ("t2", "w1:p1", "w1"), ("t2", "w2:p1", "w2")): control = HerdrControl(Endpoint.local(), "fake") control.call = Mock(return_value={"tab": {"tab_id": tab_id}, "root_pane": {"pane_id": pane_id, "tab_id": tab_id, "workspace_id": workspace}}) with self.assertRaises(ControlError): control.fork(parse_snapshot(snapshot())) control.call.assert_called_once() def test_remote_mutation_failure_is_never_replayed(self): control = HerdrControl(Endpoint.remote("host", "named"), "unused") control.platform = "posix" with patch("herdr_hub.control.subprocess.run", return_value=completed({"error": "failed"}, 127)) as run: with self.assertRaises(ControlError): control.call(["tab", "create", "--no-focus"]) run.assert_called_once() def test_ambiguous_create_failure_is_not_replayed_or_cleaned_up(self): control = HerdrControl(Endpoint.local(), "fake") control.call = Mock(side_effect=ControlError("timeout")) with self.assertRaisesRegex(ControlError, "原 Tab 保留"): control.fork(parse_snapshot(snapshot())) control.call.assert_called_once() def test_timeout_and_invalid_json_are_diagnostics(self): for result in (subprocess.TimeoutExpired(["fake"], 10), subprocess.CompletedProcess([], 0, b"not json", b"")): with patch("herdr_hub.control.subprocess.run", side_effect=[result]): with self.assertRaises(ControlError): HerdrControl(Endpoint.local(), "fake").source() def test_read_only_request_can_be_cancelled_quickly(self): control = HerdrControl(Endpoint.local(), "fake") cancelled = threading.Event() timer = threading.Timer(0.05, cancelled.set) timer.start() started = time.monotonic() try: with self.assertRaises(ControlCancelled): control._execute(["sleep", "5"], 5, cancelled) # type: ignore[attr-defined] finally: timer.cancel() self.assertLess(time.monotonic() - started, 1) class RemoteCommandTests(unittest.TestCase): def test_posix_quoting_executes_exact_argv_with_path_fallback(self): # Shell injection is checked by running only a temporary fake executable. import tempfile from pathlib import Path with tempfile.TemporaryDirectory() as directory: root = Path(directory) executable = root / ".local/bin/herdr" executable.parent.mkdir(parents=True) executable.write_text("#!/usr/bin/env python3\nimport json,sys\nprint(json.dumps(sys.argv[1:]))\n") executable.chmod(0o755) arguments = ["--session", "named", "tab", "create", "--cwd", "a b'\";$HOME`id`$(id)\nnext"] command = remote_control_command("ssh", RemoteEntry("host"), arguments, platform="posix") result = subprocess.run(["sh", "-c", command[-1]], env={**os.environ, "HOME": directory}, capture_output=True, text=True, check=True) self.assertEqual(json.loads(result.stdout), arguments) def test_windows_encoded_script_keeps_untrusted_text_out_of_shell(self): value = "C:\\a b\\quote'\"; $(whoami) & %PATH%" command = remote_control_command("ssh", RemoteEntry("host", remote_transport="ssh"), ["tab", "create", "--cwd", value], platform="windows") self.assertNotIn(value, command[-1]) script = base64.b64decode(command[-1].split()[-1]).decode("utf-16le") self.assertIn("UseShellExecute=$false", script) self.assertIn(subprocess.list2cmdline(["tab", "create", "--cwd", value]).replace("'", "''"), script) class ImmediateFuture: def __init__(self, fn, *args): try: self.value, self.error = fn(*args), None except Exception as error: self.value, self.error = None, error def done(self): return True def result(self): if self.error: raise self.error return self.value class ImmediateExecutor: def __init__(self, **kwargs): pass def __enter__(self): return self def __exit__(self, *args): pass def submit(self, fn, *args): return ImmediateFuture(fn, *args) class MenuTests(unittest.TestCase): def test_small_popup_shows_download_failure_before_proxy_notice(self): self.screen.getmaxyx.return_value = (8, 71) self.control.update_and_restart.side_effect = ControlError(update_failure( "更新失败", "未检测到代理\n正在下载\ncurl: (28) Connection timed out\n没有自动执行普通重启；请检查当前 Session。" )) self.run_menu([curses.KEY_DOWN, curses.KEY_DOWN, 10, 27]) rendered = str(self.screen.addstr.call_args_list) self.assertIn("curl: (28) Connection timed out", rendered) self.assertNotIn("未检测到代理", rendered) def test_hot_restart_success_requests_reconnect(self): self.control.update_and_restart.return_value = "verified" self.screen.getch.side_effect = [curses.KEY_DOWN, curses.KEY_DOWN, 10] menu = ControlMenu(self.control) with patch("herdr_hub.control.ThreadPoolExecutor", ImmediateExecutor), patch("curses.set_escdelay"), patch("curses.curs_set"), patch("curses.mousemask"): self.assertEqual(menu.run(self.screen), 0) self.assertTrue(menu.reconnect_requested) self.control.update_and_restart.assert_called_once() def test_hot_restart_failure_stays_visible_without_reconnect(self): self.control.update_and_restart.side_effect = ControlError("handoff failed") self.screen.getch.side_effect = [curses.KEY_DOWN, curses.KEY_DOWN, 10, 10, 27] menu = ControlMenu(self.control) with patch("herdr_hub.control.ThreadPoolExecutor", ImmediateExecutor), patch("curses.set_escdelay"), patch("curses.curs_set"), patch("curses.mousemask"): self.assertEqual(menu.run(self.screen), 0) self.assertFalse(menu.reconnect_requested) self.control.update_and_restart.assert_called_once() self.assertIn("handoff failed", str(self.screen.addstr.call_args_list)) def setUp(self): self.control = Mock(endpoint=Endpoint.local()) self.control.source.return_value = parse_snapshot(snapshot()) self.control.fork.return_value = "t2" self.screen = Mock() self.screen.getmaxyx.return_value = (24, 90) def run_menu(self, keys): self.screen.getch.side_effect = keys with patch("herdr_hub.control.ThreadPoolExecutor", ImmediateExecutor), patch("curses.set_escdelay"), patch("curses.curs_set"), patch("curses.mousemask"): return ControlMenu(self.control).run(self.screen) def test_enter_executes_and_esc_does_not(self): self.assertEqual(self.run_menu([27]), 0) self.control.fork.assert_not_called() self.assertEqual(self.run_menu([10]), 0) self.control.fork.assert_called_once() def test_click_executes_action(self): with patch("curses.getmouse", return_value=(0, 3, 3, 0, curses.BUTTON1_PRESSED)): self.run_menu([curses.KEY_MOUSE]) self.control.fork.assert_called_once() def test_error_is_visible_and_enter_cannot_repeat_failed_fork(self): self.control.fork.side_effect = ControlError("blocked: inspect t2") self.run_menu([10, 10, 27]) self.control.fork.assert_called_once() self.assertTrue(any("blocked: inspect t2" in str(call) for call in self.screen.addstr.call_args_list)) def test_snapshot_error_does_not_execute(self): self.control.source.side_effect = ControlError("没有 Agent") self.run_menu([10, 27]) self.control.fork.assert_not_called() def test_refresh_restores_fork_after_session_identity_becomes_available(self): self.screen.getmaxyx.return_value = (8, 71) self.control.source.side_effect = [ControlError("尚未上报会话 ID"), parse_snapshot(snapshot())] self.run_menu([10, 18, 10]) self.assertEqual(self.control.source.call_count, 2) self.control.fork.assert_called_once() self.assertIn("不可用", str(self.screen.addstr.call_args_list)) def test_compact_mouse_fork_still_executes(self): self.screen.getmaxyx.return_value = (8, 71) with patch("curses.getmouse", return_value=(0, 3, 1, 0, curses.BUTTON1_PRESSED)): self.run_menu([curses.KEY_MOUSE]) self.control.fork.assert_called_once() self.control.update.assert_not_called() def test_update_without_agent_keeps_result_and_prevents_repeat(self): self.control.source.side_effect = ControlError("没有 Agent") self.control.update.return_value = "更新成功" self.run_menu([curses.KEY_DOWN, 10, 10, 27]) self.control.update.assert_called_once() self.control.fork.assert_not_called() rendered = str(self.screen.addstr.call_args_list) self.assertIn("更新成功", rendered) self.assertIn("更新目标：本机", rendered) def test_click_update_reports_remote_target_and_failure(self): self.control.endpoint = Endpoint.remote("my-host") self.control.update.side_effect = ControlError("更新失败：连接断开") with patch("curses.getmouse", return_value=(0, 3, 5, 0, curses.BUTTON1_PRESSED)): self.run_menu([curses.KEY_MOUSE, 10, 27]) self.control.update.assert_called_once() rendered = str(self.screen.addstr.call_args_list) self.assertIn("更新目标：my-host", rendered) self.assertIn("更新失败：连接断开", rendered) def test_update_output_can_be_scrolled(self): self.control.update.return_value = "\n".join(f"line-{i}" for i in range(60)) self.run_menu([curses.KEY_DOWN, 10, curses.KEY_NPAGE, curses.KEY_PPAGE, 27]) self.assertIn("line-20", str(self.screen.addstr.call_args_list)) def test_update_result_visible_in_actual_small_popup(self): self.screen.getmaxyx.return_value = (8, 71) self.control.update.return_value = "更新命令已完成\nAlready up to date" self.run_menu([curses.KEY_DOWN, 10, 27]) calls = self.screen.addstr.call_args_list results = [call.args for call in calls if "Already up to date" in call.args[2]] self.assertTrue(results, "结果必须在真实的 8 行弹窗内绘制") self.assertTrue(all(3 <= args[0] < 7 for args in results)) self.assertTrue(all(0 <= call.args[0] < 8 for call in calls)) def test_compact_mouse_update_row(self): self.screen.getmaxyx.return_value = (8, 71) self.control.update.return_value = "完成" with patch("curses.getmouse", return_value=(0, 3, 2, 0, curses.BUTTON1_PRESSED)): self.run_menu([curses.KEY_MOUSE, 27]) self.control.update.assert_called_once() self.control.fork.assert_not_called() def test_update_during_snapshot_load_is_not_ignored(self): cancel = None class LoadingFuture: def done(self): return cancel.is_set() def result(self): raise ControlCancelled() class LoadingExecutor(ImmediateExecutor): def submit(inner, fn, *args): nonlocal cancel if fn == self.control.source: cancel = args[0] return LoadingFuture() return ImmediateFuture(fn, *args) self.control.update.return_value = "完成" self.screen.getch.side_effect = [curses.KEY_DOWN, 10, -1, 27] with patch("herdr_hub.control.ThreadPoolExecutor", LoadingExecutor), patch("curses.set_escdelay"), patch("curses.curs_set"), patch("curses.mousemask"): ControlMenu(self.control).run(self.screen) self.control.update.assert_called_once() class CliControlTests(unittest.TestCase): def test_hot_restart_reconnects_original_window_after_popup_closes(self): args = build_parser().parse_args(["control"]) endpoint = Endpoint.local() original = ManagedWindow("@17", endpoint.key, "Local", False, True, endpoint) other = ManagedWindow("@18", "local:other", "Other", False, True, Endpoint.local("other")) hub = Mock(ssh_bin="fake-ssh") hub.managed_windows.side_effect = [[original], [original, other]] def close_popup(run): hub.reconnect.assert_not_called() run.__self__.reconnect_requested = True return 0 with patch("herdr_hub.cli.make_hub", return_value=hub), patch("herdr_hub.cli.curses.wrapper", side_effect=close_popup): self.assertEqual(command_control(args), 0) hub.reconnect.assert_called_once_with("@17") hub.wait_for_window_startup.assert_called_once_with("@17") def test_active_window_metadata_selects_control_endpoint(self): args = build_parser().parse_args(["control"]) hub = Mock(ssh_bin="fake-ssh") endpoint = Endpoint.remote("host", "named", "server", "ssh") hub.managed_windows.return_value = [ ManagedWindow("@17", endpoint.key, "Remote", False, True, endpoint), ManagedWindow( "@18", "local:default", "Other", False, False, Endpoint.local() ), ] with patch("herdr_hub.cli.make_hub", return_value=hub), patch("herdr_hub.cli.HerdrControl") as controller, patch("herdr_hub.cli.curses.wrapper", return_value=0): self.assertEqual(command_control(args), 0) controller.assert_called_once_with(endpoint, "herdr", "fake-ssh") def test_explicit_window_remains_available_for_cli_use(self): args = build_parser().parse_args(["control", "--window", "@17"]) hub = Mock(ssh_bin="fake-ssh") endpoint = Endpoint.local() hub.managed_windows.return_value = [ ManagedWindow("@17", endpoint.key, "Local", False, False, endpoint) ] with patch("herdr_hub.cli.make_hub", return_value=hub), patch( "herdr_hub.cli.HerdrControl" ) as controller, patch("herdr_hub.cli.curses.wrapper", return_value=0): self.assertEqual(command_control(args), 0) controller.assert_called_once() def test_unmanaged_window_rejected(self): hub = Mock() hub.managed_windows.return_value = [] with patch("herdr_hub.cli.make_hub", return_value=hub), self.assertRaises(HubError): command_control(argparse.Namespace(window="@404"))