hermes-agent/tools/terminal_tool_lifecycle.py

"""Sandbox lifecycle for the terminal tool: idle reaping, teardown, manual/atexit
cleanup, and the lazy ensure_task_env bring-up. The env cache dicts and locks
stay in tools.terminal_tool (tests patch them there) and are read through it
at call time.

Split out of ``tools/terminal_tool.py``; every public/patched name is re-imported there,
so ``tools.terminal_tool.<name>`` keeps resolving (and monkeypatching) as before.
"""

import logging
import inspect
import threading
import time
from contextlib import contextmanager
from typing import Any, Dict, Optional
from tools.terminal_tool_backends import (
    _container_config_from_config,
    _local_config_from_config,
    _ssh_config_from_config,
)
from tools.terminal_tool_config import _quiet

# Log-record parity with the origin module.
logger = logging.getLogger("tools.terminal_tool")


def _create_configured_env(
    config: dict[str, Any], env_type: str, *, image: str, cwd: str, timeout: int,
    task_id: str, host_cwd: Optional[str],
):
    """``_create_environment`` with the ssh/local/container kwargs shaped from *config*
    (shared by the terminal tool, the file tools and the lazy :func:`ensure_task_env` bring-up)."""
    from tools.terminal_tool_backends import _create_environment
    from tools.terminal_tool_config import _is_container_backend
    return _create_environment(
        env_type=env_type, image=image, cwd=cwd, timeout=timeout,
        ssh_config=_ssh_config_from_config(config) if env_type == "ssh" else None,
        container_config=(
            _container_config_from_config(config) if _is_container_backend(env_type) else None
        ),
        local_config=_local_config_from_config(config) if env_type == "local" else None,
        task_id=task_id, host_cwd=host_cwd,
    )


def _cleanup_env(env: Any, *, force_remove: Optional[bool] = None) -> None:
    """Tear down one environment via cleanup()/stop()/terminate(), whichever it has.

    ``force_remove`` is forwarded to ``cleanup()`` only when given and the backend's
    signature accepts it (``DockerEnvironment``, issue #20561; other backends don't).
    Shared by ``cleanup_vm``, the idle reaper and the prompt-time backend probe so
    the signature check lives in one place.
    """
    if hasattr(env, 'cleanup'):
        if force_remove is not None and "force_remove" in inspect.signature(env.cleanup).parameters:
            env.cleanup(force_remove=force_remove)
        else:
            env.cleanup()
    elif hasattr(env, 'stop'):
        env.stop()
    elif hasattr(env, 'terminate'):
        env.terminate()


def _teardown_env(env: Any, task_id: str, *, force_remove: Optional[bool] = None, done_msg: str = "Cleaned up inactive environment for task: %s", stop_idle: bool = False) -> None:
    """``_cleanup_env`` (or, with *stop_idle*, the backend's ``stop_idle``) plus outcome logging.
    A 404/"not found" error means the sandbox is already gone — logged at info."""
    try:
        if stop_idle:
            env.stop_idle()
        else:
            _cleanup_env(env, force_remove=force_remove)
        logger.info(done_msg, task_id)
    except Exception as e:
        error_str = str(e)
        if "404" in error_str or "not found" in error_str.lower():
            logger.info("Environment for task %s already cleaned up", task_id)
        else:
            logger.warning("Error cleaning up environment for task %s: %s", task_id, e)


def _clear_file_ops_cache(task_id: str) -> None:
    """Invalidate the file_ops cache entry so ShellFileOperations can't reference a dead sandbox."""
    try:
        from tools.file_tools import clear_file_ops_cache
        clear_file_ops_cache(task_id)
    except ImportError:
        pass


def _unregister_env(task_id: str):
    """Pop *task_id* from the env cache, activity map and creation locks; return
    the env (or None). Callers run the (slow) teardown OUTSIDE the lock —
    Modal/Docker teardown can block 10-15s and would stall every concurrent
    terminal/file tool call."""
    from tools.terminal_tool import (
        _active_environments, _creation_locks, _creation_locks_lock, _env_lock,
        _last_activity,
    )
    with _env_lock:
        env = _active_environments.pop(task_id, None)
        _last_activity.pop(task_id, None)
    with _creation_locks_lock:
        _creation_locks.pop(task_id, None)
    return env


@contextmanager
def _call_in_flight(*keys: Optional[str]):
    """Hold the env cached under *keys* against the idle reaper for the duration of a call.

    Only background processes on workspace keys refresh ``_last_activity`` mid-run, so a
    foreground command longer than ``lifetime_seconds`` had its sandbox reaped out from under
    it. Refcounted (parallel calls share a key); exit restamps, so the idle countdown starts at
    call end.
    """
    from tools.terminal_tool import _calls_in_flight, _env_lock, _last_activity
    held = [k for k in dict.fromkeys(keys) if k]
    with _env_lock:
        for k in held:
            _calls_in_flight[k] = _calls_in_flight.get(k, 0) + 1
    try:
        yield
    finally:
        with _env_lock:
            now = time.time()
            for k in held:
                left = _calls_in_flight.pop(k) - 1
                if left:
                    _calls_in_flight[k] = left
                if k in _last_activity:
                    _last_activity[k] = now


def _cleanup_inactive_envs(lifetime_seconds: int = 300):
    """Clean up environments that have been inactive for longer than lifetime_seconds."""
    from tools.terminal_tool import (
        _active_environments, _calls_in_flight, _creation_locks, _creation_locks_lock, _env_lock,
        _last_activity,
    )
    current_time = time.time()

    # Sandboxes with active background processes stay alive (refresh activity).
    try:
        from tools.process_registry import process_registry
        for task_id in list(_last_activity.keys()):
            # A workspace container is keyed by its digest, not the spawning task id: also keep it
            # while any process it hosts is still running (stopping it would kill that process).
            if process_registry.has_active_processes(task_id) or (
                    _is_workspace_key(task_id)
                    and process_registry._any_running(lambda s, key=task_id: s.task_id == key)):
                _last_activity[task_id] = current_time
    except ImportError:
        pass

    # Phase 1: unregister stale entries atomically under the lock; phase 2:
    # stop them outside it (see _unregister_env for why).
    with _env_lock:
        # A key with a foreground call in flight is busy, however old its stamp.
        stale = [t for t, last in list(_last_activity.items())
                 if current_time - last > lifetime_seconds and t not in _calls_in_flight]
        envs_to_stop = [(t, _active_environments.pop(t, None)) for t in stale]
        for t in stale:
            _last_activity.pop(t, None)
        with _creation_locks_lock:
            for t in stale:
                _creation_locks.pop(t, None)
    for task_id, env in envs_to_stop:
        if env is not None:
            _clear_file_ops_cache(task_id)
            if _is_workspace_key(task_id) and callable(getattr(env, "stop_idle", None)):
                # Persistent per-workspace container: stop it (state kept, restarted on next use)
                # instead of leaving it running forever like the one profile container.
                _teardown_env(env, task_id, stop_idle=True)
            else:
                _teardown_env(env, task_id)


def _is_workspace_key(key: str) -> bool:
    """A persistent Docker key bound to one session workspace (``<base>:ws:<digest>``)."""
    return ":ws:" in str(key or "")


def get_active_env(task_id: str):
    """Return the active BaseEnvironment for *task_id*, or None."""
    from tools.terminal_tool import _active_environments, _env_lock, _resolve_container_task_id
    lookup = _resolve_container_task_id(task_id)
    with _env_lock:
        return _active_environments.get(lookup) or _active_environments.get(task_id)


def ensure_task_env(task_id: Optional[str] = None):
    """Lazily create and cache the sandbox env for *task_id* if none is active.

    Lets non-terminal callers (``tools.image_source`` reading container-only
    paths) bring the sandbox up on demand with the same machinery as the
    terminal tool. No-op on local. Returns the env, or ``None`` when local or
    when creation fails (best-effort; the caller's fail-closed path stays intact).

    :func:`terminal_tool` creates the environment on the first terminal command, but nothing else did — so
    under a non-local backend (ssh, docker, …) a session whose first action is ``vision_analyze`` on a
    container-only path hit "no active sandbox session" because the SSH/Docker handshake never ran (issue
    #62825). vision reads such paths inside the sandbox (see ``tools.image_source``), so it calls this to
    bring the env up on demand, reusing the same creation machinery as the terminal tool.
    """
    from tools.terminal_tool import (
        _active_environments, _creation_locks, _creation_locks_lock, _env_lock,
        _get_env_config, _last_activity, _resolve_container_task_id,
        _resolve_task_host_cwd, _select_image, _start_cleanup_thread, resolve_task_overrides,
    )
    config = _get_env_config()
    env_type = config["env_type"]
    if env_type == "local":
        return None

    effective_task_id = _resolve_container_task_id(task_id)

    existing = get_active_env(effective_task_id)
    if existing is not None:
        with _env_lock:
            _last_activity[effective_task_id] = time.time()
        return existing

    image = _select_image(env_type, resolve_task_overrides(task_id), config)

    _start_cleanup_thread()

    with _creation_locks_lock:
        task_lock = _creation_locks.setdefault(effective_task_id, threading.Lock())

    with task_lock:
        existing = get_active_env(effective_task_id)
        if existing is not None:
            return existing
        try:
            new_env = _create_configured_env(
                config, env_type, image=image, cwd=config["cwd"],
                timeout=config["timeout"], task_id=effective_task_id,
                host_cwd=_resolve_task_host_cwd(config, task_id),
            )
        except Exception as exc:
            logger.warning(
                "Lazy %s environment init failed for task %s: %s",
                env_type, effective_task_id[:8], exc,
            )
            return None

        with _env_lock:
            _active_environments[effective_task_id] = new_env
            _last_activity[effective_task_id] = time.time()
        logger.info(
            "%s environment lazily initialized for task %s",
            env_type, effective_task_id[:8],
        )
        return new_env


def is_persistent_env(task_id: str) -> bool:
    """True if *task_id*'s active env persists across turns.

    The agent loop skips per-turn teardown for these (persistent docker,
    daytona, modal, …); non-persistent backends are torn down at end of turn
    to prevent leakage, and the idle reaper handles the rest. Session-scoped
    docker containers count as persistent HERE: their lifetime is the session
    (removed by ``AIAgent.close()`` → ``cleanup_vm`` and the idle reaper).
    """
    env = get_active_env(task_id)
    if env is None:
        return False
    return bool(getattr(env, "_session_scoped", False) or getattr(env, "_persistent", False))


def cleanup_all_environments():
    """Clean up ALL active environments (process exit). Use with caution."""
    from tools.environments.base import kill_live_foreground_processes
    from tools.terminal_tool import _active_environments
    # A command still running when the host exits would outlive it in its own process group.
    kill_live_foreground_processes()
    cleaned = 0
    for task_id in list(_active_environments.keys()):
        try:
            cleanup_vm(task_id)
            cleaned += 1
        except Exception as e:
            logger.error("Error cleaning %s: %s", task_id, e, exc_info=True)

    if cleaned > 0:
        logger.info("Cleaned %d environments", cleaned)
    return cleaned


def cleanup_vm(task_id: str, *, force_remove: bool = False):
    """Manually clean up a specific environment by task_id.

    *force_remove* is forwarded to backends that accept it (currently only
    ``DockerEnvironment``). Default False matches session-lifecycle semantics:
    callers (``AIAgent.close()`` on TUI/gateway session teardown, the per-turn
    cleanup of non-persistent envs) must honor the user's persist-mode
    preference — stopping the container here would break the "ONE long-lived
    container shared across sessions" contract. Pass ``force_remove=True``
    only for user-initiated teardown. The idle reaper calls ``env.cleanup()``
    directly, so persist-mode idle envs are likewise no-op'd; only the orphan
    reaper at next startup reclaims them.
    """
    env = _unregister_env(task_id)
    _clear_file_ops_cache(task_id)
    if env is None:
        return
    _teardown_env(
        env, task_id, force_remove=force_remove,
        done_msg="Manually cleaned up environment for task: %s",
    )


def _evict_environment_for_task(task_id: Optional[str]) -> None:
    """Drop any cached env for *task_id* (and its collapsed key) after an
    infrastructure failure, so later calls don't reuse a dead connection."""
    from tools.terminal_tool import (
        _active_environments, _env_lock, _last_activity, _resolve_container_task_id,
    )
    keys = {_resolve_container_task_id(task_id)}
    if task_id:
        keys.add(task_id)
    evicted = []
    with _env_lock:
        for key in keys:
            env = _active_environments.pop(key, None)
            _last_activity.pop(key, None)
            if env is not None:
                evicted.append(env)
    for env in evicted:
        with _quiet("cleanup of degraded environment failed"):
            env.cleanup()