indent 0.0.8__py3-none-any.whl

exponent/core/remote_execution/languages/python_execution.py

@@ -0,0 +1,239 @@
+import asyncio
+import logging
+import queue
+import re
+import sys
+import threading
+import time
+from collections.abc import AsyncGenerator, Callable
+from typing import Any
+
+from jupyter_client.client import KernelClient
+from jupyter_client.manager import KernelManager
+
+from exponent.core.remote_execution.languages.types import (
+    PythonExecutionResult,
+    StreamedOutputPiece,
+)
+from exponent.core.remote_execution.types import PythonEnvInfo
+
+logger = logging.getLogger(__name__)
+
+
+class IOChannelHandler:
+    ESCAPE_SEQUENCE = re.compile(r"\x1B\[[0-?]*[ -/]*[@-~]")
+
+    def __init__(self, user_interrupted: Callable[[], bool] | None = None) -> None:
+        self.output_buffer: queue.Queue[str] = queue.Queue()
+        self.user_interrupted = user_interrupted
+
+    def add_message(self, message: dict[str, Any]) -> None:
+        logger.debug(f"Jupyter kernel message received: {message}")
+        output = None
+        if message["msg_type"] == "stream":
+            output = message["content"]["text"]
+        elif message["msg_type"] == "error":
+            raw_content = "\n".join(message["content"]["traceback"])
+            content = self.ESCAPE_SEQUENCE.sub("", raw_content)
+            output = content
+        if output:
+            self.output_buffer.put(output)
+
+    @staticmethod
+    def is_idle(message: dict[str, Any]) -> bool:
+        return bool(
+            message["header"]["msg_type"] == "status"
+            and message["content"]["execution_state"] == "idle"
+        )
+
+
+class Kernel:
+    def __init__(self, working_directory: str) -> None:
+        self._manager: KernelManager | None = None
+        self._client: KernelClient | None = None
+        self.io_handler: IOChannelHandler = IOChannelHandler()
+        self.working_directory = working_directory
+        self.interrupted_by_user: bool = False
+
+    @property
+    def manager(self) -> KernelManager:
+        if not self._manager:
+            self._manager = KernelManager(kernel_name="python3")
+            self._manager.start_kernel(cwd=self.working_directory)
+        return self._manager
+
+    @property
+    def client(self) -> KernelClient:
+        if not self._client:
+            self._client = self.manager.client()
+
+            while not self._client.is_alive():
+                time.sleep(0.1)
+
+            self._client.start_channels()
+        return self._client
+
+    async def wait_for_ready(self, timeout: int = 5) -> None:
+        manager = self.manager
+        start_time = time.time()
+        while not manager.is_alive():
+            if time.time() - start_time > timeout:
+                raise Exception("Kernel took too long to start")
+            await asyncio.sleep(0.05)
+        await asyncio.sleep(0.5)
+
+    def _clear_channels(self) -> None:
+        """Clear all pending messages from kernel channels."""
+        # First clear shell and control channels
+        channels = [
+            self.client.shell_channel,
+            self.client.control_channel,
+        ]
+        for channel in channels:
+            try:
+                while True:
+                    channel.get_msg(timeout=0.1)
+            except queue.Empty:
+                continue
+
+        # Then process iopub channel until we get an idle state
+        iterations = 0
+        while True:
+            try:
+                msg = self.client.iopub_channel.get_msg(timeout=0.1)
+                self.io_handler.add_message(msg)
+                if self.io_handler.is_idle(msg):
+                    break
+            except queue.Empty:
+                iterations += 1
+                if iterations > 10:
+                    logger.info("Kernel took too long to become idle")
+                    break
+
+    def iopub_listener(self, client: KernelClient) -> None:
+        while True:
+            try:
+                if (
+                    self.io_handler.user_interrupted
+                    and self.io_handler.user_interrupted()
+                ):
+                    logger.info("External halt signal received")
+                    self.manager.interrupt_kernel()
+                    self.interrupted_by_user = True
+
+                    # Wait for kernel to push any final output
+                    time.sleep(0.5)
+
+                    # Clear all channels to reset kernel state
+                    self._clear_channels()
+
+                    break
+
+                try:
+                    msg = client.iopub_channel.get_msg(timeout=1)
+                    logger.debug(f"Received message from kernel: {msg}")
+                    self.io_handler.add_message(msg)
+
+                    if self.io_handler.is_idle(msg):
+                        logger.debug("Kernel is idle.")
+                        break
+                except queue.Empty:
+                    continue
+
+            except Exception as e:  # noqa: BLE001 - TODO: Deep audit potential exceptions
+                logger.info(f"Error getting message from kernel: {e}")
+                break
+
+    # Deprecated, use execute_code_streaming
+    async def execute_code(self, code: str) -> str:
+        async for result in self.execute_code_streaming(code):
+            if isinstance(result, PythonExecutionResult):
+                return result.output
+        # should be unreachable
+        raise Exception("No result from kernel")
+
+    async def execute_code_streaming(
+        self, code: str, user_interrupted: Callable[[], bool] | None = None
+    ) -> AsyncGenerator[StreamedOutputPiece | PythonExecutionResult, None]:
+        await self.wait_for_ready()
+        self.interrupted_by_user = False
+
+        self.io_handler = IOChannelHandler(user_interrupted=user_interrupted)
+
+        client = self.client
+        client.connect_iopub()
+        iopub_thread = threading.Thread(target=self.iopub_listener, args=(client,))
+        logger.info("Starting IO listener thread.")
+        iopub_thread.start()
+
+        logger.info("Executing code in kernel.")
+        client.execute(code)
+
+        stopping = False
+
+        results = []
+        while True:
+            stopping = not iopub_thread.is_alive()
+
+            while not self.io_handler.output_buffer.empty():
+                output = self.io_handler.output_buffer.get()
+                logger.info("Execution output: %s", output)
+                yield StreamedOutputPiece(content=output)
+                results.append(output)
+
+            if stopping:
+                break
+
+            await asyncio.sleep(0.05)
+
+        # Wait for thread to fully exit
+        iopub_thread.join(timeout=1.0)
+        yield PythonExecutionResult(
+            output="".join(results), halted=self.interrupted_by_user
+        )
+
+    def close(self) -> None:
+        if self._client:
+            self._client.stop_channels()
+            self._client = None
+        if self._manager:
+            self._manager.shutdown_kernel()
+            self._manager = None
+
+
+async def execute_python(code: str, kernel: Kernel) -> str:
+    return await kernel.execute_code(code)
+
+
+async def execute_python_streaming(
+    code: str, kernel: Kernel, user_interrupted: Callable[[], bool] | None = None
+) -> AsyncGenerator[StreamedOutputPiece | PythonExecutionResult, None]:
+    async for result in kernel.execute_code_streaming(
+        code, user_interrupted=user_interrupted
+    ):
+        yield result
+
+
+### Environment Helpers
+
+
+def get_python_env_info() -> PythonEnvInfo:
+    return PythonEnvInfo(
+        interpreter_path=get_active_python_interpreter_path(),
+        interpreter_version=get_active_python_interpreter_version(),
+    )
+
+
+def get_active_python_interpreter_path() -> str | None:
+    return sys.executable
+
+
+def get_active_python_interpreter_version() -> str | None:
+    version = sys.version
+
+    match = re.search(r"(\d+\.\d+\.\d+).*", version)
+
+    if match:
+        return match.group(1)
+
+    return None

exponent/core/remote_execution/languages/shell_streaming.py

@@ -0,0 +1,221 @@
+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.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(4096)
+            if not data:
+                break
+            chunk = data.decode(encoding=encoding)
+            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(
+    code: str,
+    working_directory: str,
+    timeout: int,
+    should_halt: Callable[[], bool] | 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,
+        )
+    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,
+            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
+        process.kill()
+
+    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
+
+                    if process.returncode is not None:
+                        break
+
+                    # Schedule next read from the same stream
+                    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/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:  # noqa: BLE001
+        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,54 @@
+import getpass
+import os
+import platform
+
+from anyio import Path as AsyncPath
+
+from exponent.core.remote_execution.git import get_git_info
+from exponent.core.remote_execution.languages import python_execution
+from exponent.core.remote_execution.types import (
+    SystemContextRequest,
+    SystemContextResponse,
+    SystemInfo,
+)
+from exponent.core.remote_execution.utils import safe_read_file
+
+EXPONENT_TXT_FILENAMES = [
+    "exponent.txt",
+]
+
+
+async def get_system_context(
+    request: SystemContextRequest, working_directory: str
+) -> SystemContextResponse:
+    return SystemContextResponse(
+        correlation_id=request.correlation_id,
+        exponent_txt=await _read_exponent_txt(working_directory),
+        system_info=await get_system_info(working_directory),
+    )
+
+
+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(),
+    )
+
+
+async def _read_exponent_txt(working_directory: str) -> str | None:
+    for filename in EXPONENT_TXT_FILENAMES:
+        file_path = AsyncPath(os.path.join(working_directory, filename.lower()))
+        exists = await file_path.exists()
+
+        if exists:
+            return await safe_read_file(file_path)
+
+    return None
+
+
+def _get_user_shell() -> str:
+    return os.environ.get("SHELL", "bash")