agentworks-cli 0.2.1__py3-none-any.whl

agentworks/output.py

@@ -0,0 +1,296 @@
+"""Output contract between business logic and the presentation layer.
+
+Business logic reports data through the handler (info, detail, warn, progress)
+and signals errors by raising exceptions from the hierarchy below. The
+presentation layer (CLI, web, test) sets the handler implementation and
+catches exceptions.
+
+Business logic must never import typer, call sys.exit, or format output.
+"""
+
+from __future__ import annotations
+
+import sys
+import time
+from typing import Protocol
+
+# ---------------------------------------------------------------------------
+# Progress handle
+# ---------------------------------------------------------------------------
+
+
+class Progress(Protocol):
+    """Handle returned by OutputHandler.progress() for tracking a long operation."""
+
+    def update(self, current: int | None = None, message: str | None = None) -> None:
+        """Report progress. current is meaningful when total was provided."""
+        ...
+
+    def done(self, message: str | None = None) -> None:
+        """Mark the operation as complete."""
+        ...
+
+
+# ---------------------------------------------------------------------------
+# Handler protocol
+# ---------------------------------------------------------------------------
+
+
+class OutputHandler(Protocol):
+    """Contract for all user-facing output from business logic.
+
+    Implementations decide rendering: terminal, web, test capture, etc.
+    """
+
+    def info(self, message: str) -> None:
+        """One-shot status message (top-level)."""
+        ...
+
+    def detail(self, message: str, indent: int = 1) -> None:
+        """Sub-step or detail message. indent controls nesting depth (1 = 2 spaces, 2 = 4, etc.)."""
+        ...
+
+    def warn(self, message: str) -> None:
+        """Non-fatal warning."""
+        ...
+
+    def confirm(self, message: str, default: bool = False) -> bool:
+        """Present a yes/no question. Returns True for yes, False for no."""
+        ...
+
+    def choose(self, message: str, options: list[str]) -> int:
+        """Present a list of options. Returns the index of the selected option."""
+        ...
+
+    def pause(self, message: str) -> None:
+        """Wait for user acknowledgment (press Enter)."""
+        ...
+
+    def prompt(self, label: str, default: str | None = None) -> str:
+        """Collect a string value. If default is provided and user enters nothing, returns default."""
+        ...
+
+    def prompt_secret(self, label: str, hint: str | None = None) -> str:
+        """Collect a secret value with masked input. Rejects empty values."""
+        ...
+
+    def progress(self, label: str, total: int | None = None) -> Progress:
+        """Start a tracked operation. Returns a Progress handle.
+
+        If total is provided, the operation is determinate (percentage-based).
+        Otherwise it is indeterminate (elapsed time only).
+        """
+        ...
+
+
+# ---------------------------------------------------------------------------
+# Default handler (plain print, no terminal magic)
+# ---------------------------------------------------------------------------
+
+
+class _DefaultProgress:
+    def __init__(self, label: str, total: int | None = None) -> None:
+        self._label = label
+        self._total = total
+        self._start = time.monotonic()
+
+    def update(self, current: int | None = None, message: str | None = None) -> None:
+        parts = [f"  {self._label}..."]
+        if current is not None and self._total is not None and self._total > 0:
+            pct = current / self._total * 100
+            parts.append(f" {pct:.0f}% ({current}/{self._total})")
+        if message:
+            parts.append(f" {message}")
+        print("".join(parts))
+
+    def done(self, message: str | None = None) -> None:
+        elapsed = time.monotonic() - self._start
+        suffix = f" {message}" if message else ""
+        print(f"  {self._label} done ({elapsed:.0f}s){suffix}")
+
+
+class _DefaultHandler:
+    def info(self, message: str) -> None:
+        print(message)
+
+    def detail(self, message: str, indent: int = 1) -> None:
+        print(f"{'  ' * indent}{message}")
+
+    def warn(self, message: str) -> None:
+        print(f"Warning: {message}", file=sys.stderr)
+
+    def confirm(self, message: str, default: bool = False) -> bool:
+        try:
+            suffix = " [Y/n]" if default else " [y/N]"
+            response = input(message + suffix + " ").strip().lower()
+            if not response:
+                return default
+            return response in ("y", "yes")
+        except (EOFError, KeyboardInterrupt):
+            raise UserAbort("interrupted") from None
+
+    def choose(self, message: str, options: list[str]) -> int:
+        try:
+            print(message)
+            for i, option in enumerate(options, 1):
+                print(f"  {i}) {option}")
+            while True:
+                try:
+                    choice = int(input("Choice: "))
+                    if 1 <= choice <= len(options):
+                        return choice - 1
+                except ValueError:
+                    pass
+                print(f"Invalid choice. Enter 1-{len(options)}.")
+        except (EOFError, KeyboardInterrupt):
+            raise UserAbort("interrupted") from None
+
+    def pause(self, message: str) -> None:
+        try:
+            input(message)
+        except (EOFError, KeyboardInterrupt):
+            raise UserAbort("interrupted") from None
+
+    def prompt(self, label: str, default: str | None = None) -> str:
+        try:
+            suffix = f" [{default}]" if default else ""
+            value = input(f"{label}{suffix}: ").strip()
+            return value if value else (default or "")
+        except (EOFError, KeyboardInterrupt):
+            raise UserAbort("interrupted") from None
+
+    def prompt_secret(self, label: str, hint: str | None = None) -> str:
+        import getpass
+
+        try:
+            if hint:
+                print(f"  {hint}", file=sys.stderr)
+            while True:
+                value = getpass.getpass(f"{label}: ")
+                if value.strip():
+                    return value
+                print("(empty, try again)", file=sys.stderr)
+        except (EOFError, KeyboardInterrupt):
+            raise UserAbort("interrupted") from None
+
+    def progress(self, label: str, total: int | None = None) -> Progress:
+        print(f"  {label}...")
+        return _DefaultProgress(label, total)
+
+
+# ---------------------------------------------------------------------------
+# Module API
+# ---------------------------------------------------------------------------
+
+_handler: OutputHandler = _DefaultHandler()
+
+
+def info(message: str) -> None:
+    """Emit a top-level status message."""
+    _handler.info(message)
+
+
+def detail(message: str, indent: int = 1) -> None:
+    """Emit an indented detail/sub-step message. indent controls nesting depth."""
+    _handler.detail(message, indent)
+
+
+def warn(message: str) -> None:
+    """Emit a non-fatal warning."""
+    _handler.warn(message)
+
+
+def confirm(message: str, default: bool = False) -> bool:
+    """Present a yes/no question. Returns True for yes, False for no."""
+    return _handler.confirm(message, default)
+
+
+def choose(message: str, options: list[str]) -> int:
+    """Present a list of options. Returns the index of the selected option."""
+    return _handler.choose(message, options)
+
+
+def pause(message: str) -> None:
+    """Wait for user acknowledgment (press Enter)."""
+    _handler.pause(message)
+
+
+def prompt(label: str, default: str | None = None) -> str:
+    """Collect a string value. Returns default if user enters nothing."""
+    return _handler.prompt(label, default)
+
+
+def prompt_secret(label: str, hint: str | None = None) -> str:
+    """Collect a secret value with masked input. Rejects empty values."""
+    return _handler.prompt_secret(label, hint)
+
+
+def progress(label: str, total: int | None = None) -> Progress:
+    """Start a tracked operation. Returns a Progress handle."""
+    return _handler.progress(label, total)
+
+
+def set_handler(handler: OutputHandler) -> None:
+    """Replace the global output handler.
+
+    Call from the application entrypoint to route output through the appropriate
+    mechanism (typer.echo for CLI, websocket for web, list collector for tests).
+    """
+    global _handler
+    _handler = handler
+
+
+def get_handler() -> OutputHandler:
+    """Return the current output handler."""
+    return _handler
+
+
+# ---------------------------------------------------------------------------
+# Exception hierarchy
+# ---------------------------------------------------------------------------
+
+
+class AgentworksError(Exception):
+    """Base exception for all agentworks business logic errors.
+
+    The presentation layer catches this (and subclasses) at the entrypoint
+    and decides how to render the error.
+    """
+
+
+class VMError(AgentworksError):
+    """Error related to VM operations."""
+
+
+class WorkspaceError(AgentworksError):
+    """Error related to workspace operations."""
+
+
+class AgentError(AgentworksError):
+    """Error related to agent operations."""
+
+
+class SessionError(AgentworksError):
+    """Error related to session operations."""
+
+
+class BrokenSessionError(SessionError):
+    """Session is BROKEN (PID alive but tmux unreachable). Requires --force."""
+
+
+class ConnectivityError(AgentworksError):
+    """Error related to network, SSH, or Tailscale connectivity."""
+
+
+class BackupError(AgentworksError):
+    """Error related to backup-specific failures."""
+
+
+class ValidationError(AgentworksError):
+    """Invalid user input (name, argument, etc.)."""
+
+
+class UserAbort(AgentworksError):
+    """Raised when the user declines a confirmation prompt."""
+
+

agentworks/remote_exec.py

@@ -0,0 +1,286 @@
+"""Detached remote command execution with nohup + poll.
+
+Runs long-running commands on remote hosts in a way that survives SSH
+disconnects. The command runs under nohup with output redirected to a file.
+The workstation polls for completion by checking the process status and
+tailing new output.
+
+If the workstation reconnects after a drop, it detects the still-running
+process and resumes polling instead of starting a new one.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import time
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from agentworks import output
+from agentworks.ssh import SSHError
+
+if TYPE_CHECKING:
+    from agentworks.ssh import ExecTarget
+
+
+@dataclass
+class DetachedResult:
+    """Result of a detached remote command."""
+
+    exit_code: int
+    output: str
+
+
+# Shell wrapper that writes PID, runs the command, then writes exit status.
+_WRAPPER_TEMPLATE = """\
+#!/bin/bash
+echo $$ > {pid_file}
+{command} > {output_file} 2>&1
+echo $? > {status_file}
+"""
+
+
+def run_detached(
+    target: ExecTarget,
+    command: str,
+    *,
+    label: str = "Remote command",
+    base_path: str = "/tmp/agentworks-detached",
+    poll_interval: int = 3,
+    quiet_timeout: int = 300,
+    timeout: int | None = None,
+    as_root: bool = False,
+    quiet: bool = False,
+) -> DetachedResult:
+    """Run a command detached on a remote host, polling for completion.
+
+    If a previous run is still in progress (PID file exists, process alive),
+    resumes polling instead of starting a new one.
+
+    Running as root: prefer ``as_root=True`` to embedding ``sudo -n`` in the
+    command. With ``as_root=True``, the wrapper script itself runs as root so
+    the command, its output, and cleanup all happen uniformly with root
+    privileges. Inline ``sudo -n`` is only appropriate when parts of a
+    multi-step command need different privilege levels (e.g., a pipeline that
+    mixes root and non-root stages).
+
+    Args:
+        target: Remote execution target.
+        command: Shell command to run.
+        label: Human-readable label for progress messages.
+        base_path: Base path for output/pid/status files (unique per operation).
+        poll_interval: Seconds between polls.
+        quiet_timeout: Warn if no new output for this many seconds.
+        timeout: Hard timeout in seconds. The remote process is killed and
+            exit code 1 is returned. Partial output is still captured.
+        as_root: Run the wrapper script as root. Prefer this over embedding
+            ``sudo -n`` in the command.
+        quiet: Suppress progress output (still captured in the result).
+
+    Returns:
+        DetachedResult with exit code and full output.
+    """
+    output_file = f"{base_path}.out"
+    pid_file = f"{base_path}.pid"
+    status_file = f"{base_path}.status"
+    wrapper_file = f"{base_path}.sh"
+
+    # Check for a completed previous run (reconnect after process finished)
+    if _status_file_exists(target, status_file):
+        if not quiet:
+            output.detail(f"{label}: found completed result from previous run")
+    # Check for an existing running process (resume scenario)
+    elif _is_running(target, pid_file):
+        if not quiet:
+            output.detail(f"{label}: resuming in-progress operation...")
+    else:
+        # Write and start the wrapper script
+        wrapper = _WRAPPER_TEMPLATE.format(
+            command=command,
+            output_file=output_file,
+            pid_file=pid_file,
+            status_file=status_file,
+        )
+        target.write_file(wrapper_file, wrapper)
+
+        # Clear any stale files from a previous run
+        target.run(f"rm -f {output_file} {pid_file} {status_file}", sudo=as_root, check=False)
+
+        # Launch detached. nohup must be OUTSIDE sudo so that SIGHUP (from
+        # SSH PTY teardown) hits nohup first, not sudo. tty=False is the
+        # primary protection (no PTY = no SIGHUP), but the nohup ordering
+        # provides defense-in-depth. We don't use sudo=True here because
+        # that wraps in bash -c, putting nohup inside the sudo'd shell.
+        if as_root:
+            nohup_cmd = f"nohup sudo -n /bin/bash {wrapper_file} </dev/null >/dev/null 2>&1 &"
+        else:
+            nohup_cmd = f"nohup /bin/bash {wrapper_file} </dev/null >/dev/null 2>&1 &"
+        target.run(nohup_cmd, tty=False, check=False)
+
+        # Brief pause for PID file to be written
+        time.sleep(0.5)
+
+        if not quiet:
+            output.detail(f"{label}: started (detached)")
+
+    # Poll for completion
+    captured = _poll_until_done(
+        target,
+        output_file,
+        pid_file,
+        status_file,
+        label=label,
+        poll_interval=poll_interval,
+        quiet_timeout=quiet_timeout,
+        timeout=timeout,
+        quiet=quiet,
+    )
+
+    # Read exit code (retry on SSH failure like the output read)
+    exit_code = 1
+    for _ec_attempt in range(6):
+        try:
+            exit_code = _read_exit_code(target, status_file)
+            break
+        except SSHError:
+            time.sleep(5)
+
+    # Cleanup remote files (best-effort, may fail if SSH is still recovering)
+    with contextlib.suppress(SSHError):
+        target.run(f"rm -f {wrapper_file} {pid_file} {status_file} {output_file}", sudo=as_root, check=False)
+
+    return DetachedResult(exit_code=exit_code, output=captured)
+
+
+def _is_running(target: ExecTarget, pid_file: str) -> bool:
+    """Check if a detached process is still running."""
+    # Check PID file exists
+    result = target.run(f"test -f {pid_file}", check=False)
+    if result.returncode != 0:
+        return False
+    # Read PID and check if process is alive (ps -p works regardless of user)
+    result = target.run(f"ps -p $(cat {pid_file}) > /dev/null 2>&1", check=False)
+    return result.returncode == 0
+
+
+def _status_file_exists(target: ExecTarget, status_file: str) -> bool:
+    """Check if a status file exists (process completed)."""
+    result = target.run(f"test -f {status_file}", check=False)
+    return result.returncode == 0
+
+
+def _poll_until_done(
+    target: ExecTarget,
+    output_file: str,
+    pid_file: str,
+    status_file: str,
+    *,
+    label: str,
+    poll_interval: int,
+    quiet_timeout: int,
+    timeout: int | None = None,
+    quiet: bool = False,
+) -> str:
+    """Poll the remote process until it completes, streaming new output."""
+    last_size = 0
+    last_output_time = time.monotonic()
+    start_time = time.monotonic()
+    warned_quiet = False
+
+    ssh_failures = 0
+
+    while True:
+        time.sleep(poll_interval)
+
+        # Hard timeout -- kill the remote process to avoid orphans
+        if timeout is not None and (time.monotonic() - start_time) > timeout:
+            output.warn(
+                f"{label}: timed out after {timeout}s, killing remote process"
+            )
+            with contextlib.suppress(SSHError):
+                target.run(
+                    f"test -f {pid_file} && kill $(cat {pid_file}) 2>/dev/null",
+                    check=False,
+                )
+            break
+
+        # All polling commands go through SSH which may be temporarily
+        # down (e.g., tailscale logout disrupts Azure networking). Catch
+        # SSHError and retry -- the wrapper script on the VM keeps running.
+        try:
+            # Read new output since last poll
+            new_output = _read_new_output(target, output_file, last_size)
+            if new_output:
+                last_output_time = time.monotonic()
+                warned_quiet = False
+                if not quiet:
+                    for line in new_output.splitlines():
+                        output.detail(line)
+                last_size += len(new_output.encode("utf-8"))
+
+            # Check if process finished (status file exists)
+            status_check = target.run(f"test -f {status_file}", check=False)
+            if status_check.returncode == 0:
+                # Process done -- read any remaining output
+                final_output = _read_new_output(target, output_file, last_size)
+                if final_output and not quiet:
+                    for line in final_output.splitlines():
+                        output.detail(line)
+                break
+
+            # Check if process is still alive (PID check)
+            if not _is_running(target, pid_file):
+                # Process gone but no status file -- unexpected termination
+                break
+
+            # Reset SSH failure counter on success
+            if ssh_failures > 0 and not quiet:
+                output.detail(f"{label}: connection restored")
+            ssh_failures = 0
+
+        except SSHError:
+            ssh_failures += 1
+            if not quiet:
+                if ssh_failures == 1:
+                    output.detail(f"{label}: connection lost, waiting for recovery...")
+                elif ssh_failures % 6 == 0:
+                    output.detail(f"{label}: still waiting... ({ssh_failures * poll_interval}s)")
+            # Don't break -- the wrapper script is still running on the VM
+            continue
+
+        # Warn if no output for a while
+        quiet_secs = time.monotonic() - last_output_time
+        if quiet_secs > quiet_timeout and not warned_quiet:
+            output.warn(
+                f"{label}: no output for {int(quiet_secs)}s (still running)..."
+            )
+            warned_quiet = True
+
+    # Read the full output for the caller. Retry on SSH failure since the
+    # connection may still be recovering after a transient disruption.
+    for _read_attempt in range(6):
+        try:
+            result = target.run(f"cat {output_file} 2>/dev/null", check=False)
+            return result.stdout
+        except SSHError:
+            time.sleep(5)
+    output.warn(f"{label}: unable to retrieve remote output after repeated SSH failures")
+    return ""
+
+
+def _read_new_output(target: ExecTarget, output_file: str, offset: int) -> str:
+    """Read new bytes from the output file since the given offset."""
+    result = target.run(
+        f"tail -c +{offset + 1} {output_file} 2>/dev/null",
+        check=False,
+    )
+    return result.stdout if result.returncode == 0 else ""
+
+
+def _read_exit_code(target: ExecTarget, status_file: str) -> int:
+    """Read the exit code from the status file."""
+    result = target.run(f"cat {status_file} 2>/dev/null", check=False)
+    try:
+        return int(result.stdout.strip())
+    except (ValueError, AttributeError):
+        return 1  # assume failure if we can't read it