from __future__ import annotations import argparse import base64 import curses import json import os import subprocess 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, ) 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_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_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 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() class CliControlTests(unittest.TestCase): 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"))