indent 0.1.26__py3-none-any.whl

exponent/core/remote_execution/languages/shell_streaming.py

@@ -0,0 +1,226 @@
+import asyncio
+import os
+import platform
+import shutil
+import signal
+from collections.abc import AsyncGenerator, Callable
+from typing import Any
+
+from exponent.core.remote_execution.default_env import get_process_env
+from exponent.core.remote_execution.languages.types import (
+    ShellExecutionResult,
+    StreamedOutputPiece,
+)
+from exponent.core.remote_execution.utils import smart_decode
+
+STDOUT_FD = 1
+STDERR_FD = 2
+MAX_TIMEOUT = 300
+
+
+def get_rc_file_source_command(shell_path: str) -> str:
+    """
+    Returns a command to source the user's shell rc file
+    Login profiles are already sourced via the -l flag
+    """
+    # On Windows, shell behavior is different
+    if platform.system() == "Windows":
+        return ""  # Windows shells don't typically use rc files in the same way
+
+    shell_name = os.path.basename(shell_path)
+    home_dir = os.path.expanduser("~")
+
+    if shell_name == "zsh":
+        zshrc = os.path.join(home_dir, ".zshrc")
+        if os.path.exists(zshrc):
+            return f"source {zshrc} 2>/dev/null || true; "
+    elif shell_name == "bash":
+        bashrc = os.path.join(home_dir, ".bashrc")
+        if os.path.exists(bashrc):
+            return f"source {bashrc} 2>/dev/null || true; "
+
+    return ""  # No rc file found or unsupported shell
+
+
+async def read_stream(
+    stream: asyncio.StreamReader, fd: int, output: list[tuple[int, str]]
+) -> AsyncGenerator[StreamedOutputPiece, None]:
+    data: bytes = b""
+    encoding = "utf-8"
+
+    while True:
+        try:
+            data = await stream.read(50_000)
+            if not data:
+                break
+            chunk = data.decode(encoding=encoding, errors="replace")
+            output.append((fd, chunk))
+            yield StreamedOutputPiece(content=chunk)
+        except UnicodeDecodeError:
+            decode_result = smart_decode(data)
+
+            if decode_result is None:
+                break
+
+            # Store the detected encoding as a hint for how
+            # to decode the remaining chunks of data
+            chunk, encoding = decode_result
+
+            output.append((fd, chunk))
+            yield StreamedOutputPiece(content=chunk)
+
+            continue
+        except asyncio.CancelledError:
+            raise
+        except Exception:
+            break
+
+
+async def execute_shell_streaming(  # noqa: PLR0915
+    code: str,
+    working_directory: str,
+    timeout: int,
+    should_halt: Callable[[], bool] | None = None,
+    env: dict[str, str] | None = None,
+) -> AsyncGenerator[StreamedOutputPiece | ShellExecutionResult, None]:
+    timeout_seconds = min(timeout, MAX_TIMEOUT)
+
+    shell_path = os.environ.get("SHELL") or shutil.which("bash") or shutil.which("sh")
+
+    if not shell_path:
+        process = await asyncio.create_subprocess_shell(
+            code,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+            cwd=working_directory,
+            env=get_process_env(env),
+        )
+    else:
+        # Add rc file sourcing to the command
+        rc_source_cmd = get_rc_file_source_command(shell_path)
+        full_command = f"{rc_source_cmd}{code}"
+
+        process = await asyncio.create_subprocess_exec(
+            shell_path,
+            "-l",
+            "-c",
+            full_command,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+            cwd=working_directory,
+            env=get_process_env(env),
+            start_new_session=True if platform.system() != "Windows" else False,
+        )
+
+    exit_code = None
+    output: list[tuple[int, str]] = []
+    halted = False
+    timed_out = False
+    assert process.stdout
+    assert process.stderr
+
+    async def monitor_halt() -> None:
+        nonlocal halted
+
+        while True:
+            if should_halt and should_halt():
+                # Send signal to process group for proper interrupt propagation
+                try:
+                    if platform.system() != "Windows":
+                        # Send SIGTERM to the process group
+                        os.killpg(process.pid, signal.SIGTERM)
+                        try:
+                            await asyncio.wait_for(process.wait(), timeout=2.0)
+                        except TimeoutError:
+                            # Fall back to SIGKILL to process group
+                            os.killpg(process.pid, signal.SIGKILL)
+                    else:
+                        # Windows fallback
+                        process.terminate()
+                        try:
+                            await asyncio.wait_for(process.wait(), timeout=2.0)
+                        except TimeoutError:
+                            process.kill()
+                except (ProcessLookupError, OSError):
+                    # Process or process group already terminated
+                    pass
+                halted = True
+                break
+            if process.returncode is not None:
+                break
+            await asyncio.sleep(0.1)
+
+    def on_timeout() -> None:
+        nonlocal timed_out
+        timed_out = True
+        try:
+            process.kill()
+        except ProcessLookupError:
+            pass
+
+    try:
+        halt_task = asyncio.create_task(monitor_halt()) if should_halt else None
+        timeout_handle = asyncio.get_running_loop().call_later(
+            timeout_seconds, on_timeout
+        )
+
+        # Stream stdout and stderr concurrently using wait
+        stdout_gen = read_stream(process.stdout, STDOUT_FD, output)
+        stderr_gen = read_stream(process.stderr, STDERR_FD, output)
+
+        stdout_task = asyncio.create_task(stdout_gen.__anext__())
+        stderr_task = asyncio.create_task(stderr_gen.__anext__())
+        pending = {stdout_task, stderr_task}
+
+        while pending:
+            done, pending = await asyncio.wait(
+                pending, return_when=asyncio.FIRST_COMPLETED
+            )
+
+            for task in done:
+                try:
+                    piece = await task
+                    yield piece
+
+                    # Schedule next read from the same stream
+                    # Don't check process.returncode here - we need to drain all buffered output
+                    if task is stdout_task and not process.stdout.at_eof():
+                        stdout_task = asyncio.create_task(stdout_gen.__anext__())
+                        pending.add(stdout_task)
+                    elif task is stderr_task and not process.stderr.at_eof():
+                        stderr_task = asyncio.create_task(stderr_gen.__anext__())
+                        pending.add(stderr_task)
+                except StopAsyncIteration:
+                    continue
+
+        exit_code = await process.wait()
+        timeout_handle.cancel()
+
+    except asyncio.CancelledError:
+        process.kill()
+        raise
+    finally:
+        # Explicitly kill the process if it's still running
+        if process and process.returncode is None:
+            try:
+                if platform.system() != "Windows":
+                    os.killpg(process.pid, signal.SIGTERM)
+                else:
+                    process.terminate()
+            except (ProcessLookupError, OSError):
+                pass
+
+        tasks_to_cancel: list[asyncio.Task[Any]] = [stdout_task, stderr_task]
+        if halt_task:
+            tasks_to_cancel.append(halt_task)
+
+        await asyncio.gather(*tasks_to_cancel, return_exceptions=True)
+
+    formatted_output = "".join([chunk for (_, chunk) in output]).strip() + "\n\n"
+
+    yield ShellExecutionResult(
+        output=formatted_output,
+        cancelled_for_timeout=timed_out,
+        exit_code=None if timed_out else exit_code,
+        halted=halted,
+    )

exponent/core/remote_execution/languages/types.py

@@ -0,0 +1,20 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class StreamedOutputPiece:
+    content: str
+
+
+@dataclass
+class ShellExecutionResult:
+    output: str
+    cancelled_for_timeout: bool
+    exit_code: int | None
+    halted: bool = False
+
+
+@dataclass
+class PythonExecutionResult:
+    output: str
+    halted: bool = False

exponent/core/remote_execution/port_utils.py

@@ -0,0 +1,73 @@
+import time
+
+import psutil
+
+from exponent.core.remote_execution.types import PortInfo
+
+
+def get_port_usage() -> list[PortInfo] | None:
+    """
+    Get information about all listening ports on the system.
+
+    Returns:
+        List of PortInfo objects containing process name, port, protocol, pid, and uptime.
+        Returns None if there's a permission error.
+        Returns empty list if no listening ports are found.
+    """
+    try:
+        connections = psutil.net_connections(kind="tcp")
+    except (psutil.AccessDenied, PermissionError):
+        # If we don't have permission to see connections, return None
+        return None
+    except Exception:
+        # For any other unexpected errors, return None
+        return None
+
+    port_info_list: list[PortInfo] = []
+    current_time = time.time()
+
+    for conn in connections:
+        # Only include TCP ports in LISTEN state
+        if conn.status != "LISTEN":
+            continue
+
+        # Skip if no local address (shouldn't happen for LISTEN, but be safe)
+        if not conn.laddr:
+            continue
+
+        port = conn.laddr.port
+        pid = conn.pid
+
+        # Try to get process information
+        process_name = "unknown"
+        uptime_seconds = None
+
+        if pid:
+            try:
+                process = psutil.Process(pid)
+                process_name = process.name()
+
+                # Calculate uptime
+                create_time = process.create_time()
+                uptime_seconds = current_time - create_time
+            except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
+                # Process disappeared or we don't have permission
+                pass
+            except Exception:
+                # Any other unexpected error, just skip process info
+                pass
+
+        port_info = PortInfo(
+            process_name=process_name,
+            port=port,
+            protocol="TCP",
+            pid=pid,
+            uptime_seconds=uptime_seconds,
+        )
+        port_info_list.append(port_info)
+
+        # Limit to 50 ports to avoid bloating the heartbeat payload
+        if len(port_info_list) >= 50:
+            break
+
+    return port_info_list

exponent/core/remote_execution/session.py

@@ -0,0 +1,128 @@
+import logging
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
+from typing import TYPE_CHECKING, Optional
+
+import httpx
+from httpx import AsyncClient, Request, Response
+
+from exponent.core.remote_execution.exceptions import ExponentError
+from exponent.core.remote_execution.languages.python_execution import Kernel
+from exponent.core.remote_execution.utils import format_error_log
+
+if TYPE_CHECKING:
+    from exponent.core.config import Settings
+
+logger = logging.getLogger(__name__)
+
+
+class SessionLog:
+    def __init__(self) -> None:
+        self.log_buffer: list[str] = []
+        self.max_size = 5
+
+    def append_log(self, log: str) -> None:
+        self.log_buffer.append(log)
+        self.log_buffer = self.log_buffer[-self.max_size :]
+
+    def get_logs(self) -> list[str]:
+        return self.log_buffer
+
+    async def log_request(self, request: Request) -> None:
+        self.append_log(f"Request: {request.method} {request.url}")
+
+    async def log_response(self, response: Response) -> None:
+        request = response.request
+        await response.aread()
+        self.append_log(
+            f"Response for request: {request.method} {request.url}\n"
+            f"Response: {response.status_code}, {response.text}"
+        )
+
+
+class RemoteExecutionClientSession:
+    def __init__(
+        self, working_directory: str, base_url: str, base_ws_url: str, api_key: str
+    ):
+        self.chat_uuid: str | None = None
+
+        self.working_directory = working_directory
+        self.kernel = Kernel(working_directory=working_directory)
+        self.api_log = SessionLog()
+
+        self.api_client = AsyncClient(
+            base_url=base_url,
+            headers={"API-KEY": api_key},
+            event_hooks={
+                "request": [self.api_log.log_request],
+                "response": [self.api_log.log_response],
+            },
+        )
+
+        self.ws_client = AsyncClient(
+            base_url=base_ws_url,
+            headers={"API-KEY": api_key},
+            event_hooks={
+                "request": [self.api_log.log_request],
+                "response": [self.api_log.log_response],
+            },
+        )
+
+    def set_chat_uuid(self, chat_uuid: str) -> None:
+        self.chat_uuid = chat_uuid
+
+
+async def send_exception_log(
+    exc: Exception,
+    session: RemoteExecutionClientSession | None = None,
+    settings: Optional["Settings"] = None,
+) -> None:
+    error_log = format_error_log(
+        exc=exc,
+        chat_uuid=session.chat_uuid if session else None,
+        attachment_lines=session.api_log.get_logs() if session else None,
+    )
+
+    if session:
+        api_client = session.api_client
+    elif settings:
+        api_key = settings.api_key
+        if not api_key:
+            raise ValueError("No API key provided")
+        api_client = AsyncClient(
+            base_url=settings.get_base_api_url(),
+            headers={"API-KEY": api_key},
+        )
+    else:
+        raise ValueError("No session or settings provided")
+
+    if not error_log:
+        return
+    try:
+        await api_client.post(
+            "/api/remote_execution/log_error",
+            content=error_log.model_dump_json(),
+            timeout=60,
+        )
+    except httpx.ConnectError:
+        logger.info("Failed to send error log")
+
+
+@asynccontextmanager
+async def get_session(
+    working_directory: str,
+    base_url: str,
+    base_ws_url: str,
+    api_key: str,
+) -> AsyncGenerator[RemoteExecutionClientSession, None]:
+    session = RemoteExecutionClientSession(
+        working_directory, base_url, base_ws_url, api_key
+    )
+    try:
+        yield session
+    except Exception as exc:
+        await send_exception_log(exc, session=session, settings=None)
+        raise ExponentError(str(exc))
+    finally:
+        session.kernel.close()
+        await session.api_client.aclose()

exponent/core/remote_execution/system_context.py

@@ -0,0 +1,26 @@
+import getpass
+import os
+import platform
+
+from exponent.core.remote_execution.git import get_git_info
+from exponent.core.remote_execution.languages import python_execution
+from exponent.core.remote_execution.port_utils import get_port_usage
+from exponent.core.remote_execution.types import (
+    SystemInfo,
+)
+
+
+async def get_system_info(working_directory: str) -> SystemInfo:
+    return SystemInfo(
+        name=getpass.getuser(),
+        cwd=working_directory,
+        os=platform.system(),
+        shell=_get_user_shell(),
+        git=await get_git_info(working_directory),
+        python_env=python_execution.get_python_env_info(),
+        port_usage=get_port_usage(),
+    )
+
+
+def _get_user_shell() -> str:
+    return os.environ.get("SHELL", "bash")