From 46b5314229a0646d84cbc1d0fabad83cfca84608 Mon Sep 17 00:00:00 2001 From: Dominic Bejar Date: Thu, 6 Aug 2026 20:25:50 +0000 Subject: [PATCH] fix(terminal): harden scope fallback and memory override --- tests/tools/test_process_registry.py | 74 ++++++++++++++++++++++++++++ tools/process_registry.py | 27 +++++++--- 2 files changed, 94 insertions(+), 7 deletions(-) diff --git a/tests/tools/test_process_registry.py b/tests/tools/test_process_registry.py index fce8fe500536f..2de785123dae4 100644 --- a/tests/tools/test_process_registry.py +++ b/tests/tools/test_process_registry.py @@ -1793,6 +1793,61 @@ class TestSystemdCgroupIsolation: assert argv[-3:] == ["/bin/bash", "-lic", "set +m; codex"] assert session.systemd_unit == f"hermes-worker-{session.id}.scope" + def test_pty_spawn_failure_reaps_scope_before_distinct_pipe_fallback( + self, registry, monkeypatch + ): + """A failed PTY scope must not collide with the pipe fallback scope.""" + from ptyprocess import PtyProcess + + events = [] + fake_proc = MagicMock() + fake_proc.pid = 4321 + fake_proc.stdout = iter([]) + fake_proc.stdin = MagicMock() + fake_proc.poll.return_value = None + + def fake_popen(argv, **_kwargs): + events.append(("pipe", list(argv))) + return fake_proc + + def fake_stop(unit_name): + events.append(("stop", unit_name)) + return True + + def fail_pty(*_args, **_kwargs): + events.append(("pty", None)) + raise RuntimeError("PTY wrapper failed after scope creation") + + monkeypatch.setattr("tools.process_registry._find_shell", lambda: "/bin/bash") + monkeypatch.setattr( + "tools.process_registry._systemd_run_user_scope_available", + lambda: True, + ) + monkeypatch.setattr( + "gateway.restart.is_gateway_supervisor_process", + lambda: True, + ) + monkeypatch.setattr("shutil.which", lambda name: "/usr/bin/systemd-run") + + with patch.object(PtyProcess, "spawn", side_effect=fail_pty), \ + patch("subprocess.Popen", side_effect=fake_popen), \ + patch("tools.process_registry._stop_systemd_unit", side_effect=fake_stop), \ + patch("threading.Thread", return_value=MagicMock()), \ + patch.object(registry, "_write_checkpoint"): + session = registry.spawn_local("codex", cwd="/tmp", use_pty=True) + + assert [event[0] for event in events] == ["pty", "stop", "pipe"] + stopped_unit = events[1][1] + fallback_argv = events[2][1] + assert stopped_unit == f"hermes-worker-{session.id}.scope" + unit_idx = fallback_argv.index("--unit") + assert fallback_argv[unit_idx + 1] == ( + f"hermes-worker-{session.id}-pipe-fallback" + ) + assert session.systemd_unit == ( + f"hermes-worker-{session.id}-pipe-fallback.scope" + ) + def test_worker_memory_limit_honors_local_guard_mb_override(self, monkeypatch): import tools.process_registry as pr @@ -1806,6 +1861,25 @@ class TestSystemdCgroupIsolation: assert f"MemoryMax={123 * 1024 * 1024}" in argv + def test_worker_memory_limit_caps_oversized_local_guard_override( + self, monkeypatch + ): + import tools.process_registry as pr + + monkeypatch.setenv("TERMINAL_LOCAL_MEMORY_MAX_MB", "999999") + monkeypatch.setattr( + pr.Path, + "read_text", + lambda *_args, **_kwargs: (_ for _ in ()).throw(OSError("no cgroup")), + ) + monkeypatch.setattr( + pr.os, + "sysconf", + lambda *_args: (_ for _ in ()).throw(OSError("no sysconf")), + ) + + assert pr._worker_memory_max_bytes() == pr._DEFAULT_WORKER_MEMORY_MAX_BYTES + def test_kill_recovered_detached_already_exited_stops_persisted_scope( self, registry, monkeypatch ): diff --git a/tools/process_registry.py b/tools/process_registry.py index aec51dd15b86b..bfd6477153166 100644 --- a/tools/process_registry.py +++ b/tools/process_registry.py @@ -108,19 +108,21 @@ _WORKER_MEMORY_MAX_CAP_BYTES = 4 * 1024 * 1024 * 1024 def _worker_memory_max_bytes() -> int: """Return a finite per-worker cgroup limit without widening host risk. - The proposed local-memory-guard environment override is honored so this - isolation composes with PR #57121 instead of inventing a second knob. + The proposed local-memory-guard environment override is honored when it + tightens the safe bound, so this isolation composes with PR #57121 instead + of inventing a second knob. An oversized override cannot widen host risk. Otherwise retain the tighter of the gateway's current cgroup-v2 ``memory.max`` and half of physical RAM, capped at 4 GiB. This keeps the sibling worker outside the gateway cgroup while ensuring the worker cannot consume memory up to the enclosing user slice or host limit. """ + override_bound: Optional[int] = None override = os.getenv("TERMINAL_LOCAL_MEMORY_MAX_MB", "").strip() if override: try: parsed = int(override) * 1024 * 1024 if parsed >= _MIN_WORKER_MEMORY_MAX_BYTES: - return parsed + override_bound = parsed except ValueError: pass logger.warning( @@ -158,7 +160,8 @@ def _worker_memory_max_bytes() -> int: except (OSError, ValueError, TypeError): pass - return min(candidates) if candidates else _DEFAULT_WORKER_MEMORY_MAX_BYTES + safe_bound = min(candidates) if candidates else _DEFAULT_WORKER_MEMORY_MAX_BYTES + return min(override_bound, safe_bound) if override_bound else safe_bound def _systemd_run_user_scope_available() -> bool: @@ -969,6 +972,7 @@ class ProcessRegistry: started_at=time.time(), ) + pty_scope_attempted = False if use_pty: # Try PTY mode for interactive CLI tools try: @@ -1002,6 +1006,7 @@ class ProcessRegistry: pty_argv, unit_suffix=session.id, ) session.systemd_unit = f"hermes-worker-{session.id}.scope" + pty_scope_attempted = True elif not _IS_WINDOWS: try: from gateway.restart import is_gateway_supervisor_process as _sup @@ -1046,6 +1051,9 @@ class ProcessRegistry: logger.warning("ptyprocess not installed, falling back to pipe mode") except Exception as e: logger.warning("PTY spawn failed (%s), falling back to pipe mode", e) + if pty_scope_attempted and session.systemd_unit: + _stop_systemd_unit(session.systemd_unit) + session.systemd_unit = "" # Standard Popen path (non-PTY or PTY fallback) # Use the user's login shell for consistency with LocalEnvironment -- @@ -1080,10 +1088,15 @@ class ProcessRegistry: use_systemd_scope = False if use_systemd_scope: - spawn_argv = _build_systemd_scope_argv( - shell_argv, unit_suffix=session.id, + unit_suffix = ( + f"{session.id}-pipe-fallback" + if pty_scope_attempted + else session.id ) - session.systemd_unit = f"hermes-worker-{session.id}.scope" + spawn_argv = _build_systemd_scope_argv( + shell_argv, unit_suffix=unit_suffix, + ) + session.systemd_unit = f"hermes-worker-{unit_suffix}.scope" # systemd-run creates the new session/cgroup for us; do NOT also # set start_new_session (harmless, but redundant and it can mask # scope-creation failures in some systemd versions).