procler 0.2.0__py3-none-any.whl

procler/core/context_docker.py

@@ -0,0 +1,384 @@
+"""Docker container execution context using docker-py SDK."""
+
+import asyncio
+import re
+import time
+from collections.abc import AsyncIterator, Callable
+
+try:
+    import docker
+    from docker.errors import APIError, NotFound
+
+    DOCKER_AVAILABLE = True
+except ImportError:
+    DOCKER_AVAILABLE = False
+
+from .context_base import ExecResult, ExecutionContext, ProcessHandle
+
+# Docker container name pattern: alphanumeric, underscore, dash, dot
+# Must start with alphanumeric
+CONTAINER_NAME_PATTERN = re.compile(r"^[a-zA-Z0-9][a-zA-Z0-9_.-]*$")
+
+
+def validate_container_name(name: str) -> bool:
+    """Validate container name matches Docker naming rules."""
+    if not name or len(name) > 255:
+        return False
+    return bool(CONTAINER_NAME_PATTERN.match(name))
+
+
+class DockerContext(ExecutionContext):
+    """Execute processes inside Docker containers using docker-py SDK."""
+
+    def __init__(self):
+        if not DOCKER_AVAILABLE:
+            raise RuntimeError("Docker SDK not available. Install with: pip install docker")
+        self._client = docker.from_env()
+        self._exec_instances: dict[int, tuple] = {}  # pid -> (container, exec_id)
+        self._stream_tasks: dict[int, asyncio.Task] = {}  # pid -> stream output task
+
+    @property
+    def context_type(self) -> str:
+        return "docker"
+
+    def _get_container(self, container_name: str):
+        """Get a container by name or ID."""
+        if not validate_container_name(container_name):
+            raise ValueError(f"Invalid container name: '{container_name}'")
+        try:
+            return self._client.containers.get(container_name)
+        except NotFound:
+            raise ValueError(f"Container '{container_name}' not found")
+
+    def _is_container_running(self, container_name: str) -> bool:
+        """Check if a container is running."""
+        try:
+            container = self._get_container(container_name)
+            return container.status == "running"
+        except ValueError:
+            return False
+
+    def check_container_available(self, container_name: str) -> tuple[bool, str | None]:
+        """Pre-flight check if container is available and running.
+
+        Returns (is_available, error_message).
+        """
+        if not validate_container_name(container_name):
+            return False, f"Invalid container name: '{container_name}'"
+        try:
+            container = self._client.containers.get(container_name)
+            if container.status != "running":
+                return False, f"Container '{container_name}' is not running (status: {container.status})"
+            return True, None
+        except NotFound:
+            return False, f"Container '{container_name}' not found"
+        except Exception as e:
+            return False, f"Error checking container '{container_name}': {e}"
+
+    def list_containers(self, running_only: bool = True) -> list[dict]:
+        """List available containers."""
+        containers = self._client.containers.list(all=not running_only)
+        return [
+            {
+                "id": c.short_id,
+                "name": c.name,
+                "status": c.status,
+                "image": c.image.tags[0] if c.image.tags else c.image.short_id,
+            }
+            for c in containers
+        ]
+
+    async def start_process(
+        self,
+        command: str,
+        cwd: str | None = None,
+        env: dict[str, str] | None = None,
+        on_stdout: Callable[[str], None] | None = None,
+        on_stderr: Callable[[str], None] | None = None,
+        on_exit: Callable[[int], None] | None = None,
+        container_name: str | None = None,
+    ) -> ProcessHandle:
+        """
+        Start a long-running process inside a Docker container.
+
+        Note: Docker exec doesn't support background processes the same way
+        as local subprocesses. This implementation runs the command and
+        streams output, but the "PID" is actually a unique identifier
+        for the exec instance, not a real container PID.
+        """
+        if not container_name:
+            raise ValueError("container_name is required for Docker context")
+
+        container = self._get_container(container_name)
+
+        if container.status != "running":
+            raise RuntimeError(f"Container '{container_name}' is not running (status: {container.status})")
+
+        # Build environment variables list
+        env_list = [f"{k}={v}" for k, v in (env or {}).items()]
+
+        # Create exec instance
+        exec_result = container.client.api.exec_create(
+            container.id,
+            command,
+            workdir=cwd,
+            environment=env_list if env_list else None,
+            stdout=True,
+            stderr=True,
+            tty=False,
+        )
+        exec_id = exec_result["Id"]
+
+        # Start exec and get stream handle
+        try:
+            output = container.client.api.exec_start(exec_id, stream=True, demux=True)
+        except Exception as e:
+            raise RuntimeError(f"Failed to start exec in container '{container_name}': {e}") from e
+
+        def resolve_exec_pid() -> int:
+            """Resolve the real PID for this exec session inside the container."""
+            try:
+                inspect = container.client.api.exec_inspect(exec_id)
+                pid = inspect.get("Pid")
+                if isinstance(pid, int) and pid > 0:
+                    return pid
+            except Exception:
+                return 0
+            return 0
+
+        exec_pid = resolve_exec_pid()
+        if exec_pid <= 0:
+            # Wait briefly for PID to appear
+            for _ in range(10):
+                await asyncio.sleep(0.1)
+                exec_pid = resolve_exec_pid()
+                if exec_pid > 0:
+                    break
+
+        # Fall back to a pseudo pid if Docker doesn't report one
+        if exec_pid <= 0:
+            import random
+
+            exec_pid = random.randint(100000, 999999)
+
+        self._exec_instances[exec_pid] = (container, exec_id)
+
+        # Start streaming output in background
+        async def stream_output():
+            try:
+                for stdout_chunk, stderr_chunk in output:
+                    if stdout_chunk:
+                        for line in stdout_chunk.decode("utf-8", errors="replace").splitlines():
+                            if on_stdout:
+                                on_stdout(line)
+                    if stderr_chunk:
+                        for line in stderr_chunk.decode("utf-8", errors="replace").splitlines():
+                            if on_stderr:
+                                on_stderr(line)
+
+                # Get exit code
+                inspect = container.client.api.exec_inspect(exec_id)
+                exit_code = inspect.get("ExitCode", 0)
+
+                if on_exit:
+                    on_exit(exit_code)
+
+            except Exception as e:
+                if on_stderr:
+                    on_stderr(f"Error streaming output: {e}")
+                if on_exit:
+                    on_exit(-1)
+            finally:
+                # Cleanup exec instance and task tracking
+                self._exec_instances.pop(exec_pid, None)
+                self._stream_tasks.pop(exec_pid, None)
+
+        # Run in background and track the task
+        task = asyncio.create_task(stream_output())
+        self._stream_tasks[exec_pid] = task
+
+        return ProcessHandle(pid=exec_pid, context_type=self.context_type)
+
+    async def stop_process(self, handle: ProcessHandle, timeout: float = 10.0) -> int:
+        """
+        Stop a running process in a Docker container.
+
+        Note: Docker exec instances cannot be stopped directly. We can only
+        wait for them to complete or kill the container (which is destructive).
+        """
+        if handle.pid not in self._exec_instances:
+            # Cleanup any orphaned task
+            task = self._stream_tasks.pop(handle.pid, None)
+            if task and not task.done():
+                task.cancel()
+                try:
+                    await task
+                except asyncio.CancelledError:
+                    pass
+            return 0  # Already done
+
+        container, exec_id = self._exec_instances[handle.pid]
+
+        # Check if still running
+        try:
+            inspect = container.client.api.exec_inspect(exec_id)
+            if not inspect.get("Running", False):
+                exit_code = inspect.get("ExitCode", 0)
+                self._exec_instances.pop(handle.pid, None)
+                # Cancel stream task
+                task = self._stream_tasks.pop(handle.pid, None)
+                if task and not task.done():
+                    task.cancel()
+                    try:
+                        await task
+                    except asyncio.CancelledError:
+                        pass
+                return exit_code
+        except Exception:
+            pass
+
+        # Wait for completion (can't really kill exec)
+        start_time = time.monotonic()
+        while time.monotonic() - start_time < timeout:
+            await asyncio.sleep(0.1)
+            try:
+                inspect = container.client.api.exec_inspect(exec_id)
+                if not inspect.get("Running", False):
+                    exit_code = inspect.get("ExitCode", 0)
+                    self._exec_instances.pop(handle.pid, None)
+                    task = self._stream_tasks.pop(handle.pid, None)
+                    if task and not task.done():
+                        task.cancel()
+                        try:
+                            await task
+                        except asyncio.CancelledError:
+                            pass
+                    return exit_code
+            except Exception:
+                break
+
+        # Timeout - process may still be running, cleanup anyway
+        self._exec_instances.pop(handle.pid, None)
+        task = self._stream_tasks.pop(handle.pid, None)
+        if task and not task.done():
+            task.cancel()
+            try:
+                await task
+            except asyncio.CancelledError:
+                pass
+        return -1
+
+    async def is_running(self, handle: ProcessHandle) -> bool:
+        """Check if a process is still running."""
+        if handle.pid not in self._exec_instances:
+            return False
+
+        container, exec_id = self._exec_instances[handle.pid]
+
+        try:
+            inspect = container.client.api.exec_inspect(exec_id)
+            return inspect.get("Running", False)
+        except Exception:
+            return False
+
+    async def exec_command(
+        self,
+        command: str,
+        cwd: str | None = None,
+        env: dict[str, str] | None = None,
+        timeout: float | None = None,
+        container_name: str | None = None,
+    ) -> ExecResult:
+        """Execute a one-shot command inside a Docker container."""
+        if not container_name:
+            raise ValueError("container_name is required for Docker context")
+
+        container = self._get_container(container_name)
+
+        if container.status != "running":
+            return ExecResult(
+                exit_code=-1,
+                stdout="",
+                stderr=f"Container '{container_name}' is not running (status: {container.status})",
+            )
+
+        # Build environment variables list
+        env_list = [f"{k}={v}" for k, v in (env or {}).items()]
+
+        try:
+            # Run command with exec
+            exit_code, output = container.exec_run(
+                command,
+                workdir=cwd,
+                environment=env_list if env_list else None,
+                demux=True,
+            )
+
+            stdout = ""
+            stderr = ""
+
+            if output:
+                stdout_bytes, stderr_bytes = output
+                if stdout_bytes:
+                    stdout = stdout_bytes.decode("utf-8", errors="replace")
+                if stderr_bytes:
+                    stderr = stderr_bytes.decode("utf-8", errors="replace")
+
+            return ExecResult(
+                exit_code=exit_code,
+                stdout=stdout,
+                stderr=stderr,
+            )
+
+        except APIError as e:
+            return ExecResult(
+                exit_code=-1,
+                stdout="",
+                stderr=f"Docker API error: {e}",
+            )
+        except Exception as e:
+            return ExecResult(
+                exit_code=-1,
+                stdout="",
+                stderr=f"Error executing command: {e}",
+            )
+
+    async def stream_logs(
+        self,
+        handle: ProcessHandle,
+        follow: bool = True,
+    ) -> AsyncIterator[tuple[str, str]]:
+        """
+        Stream logs from a running process in a Docker container.
+
+        Note: For Docker exec, logs are captured during exec_start.
+        This method is a placeholder for compatibility.
+        """
+        # Docker exec logs are streamed during execution, not after
+        # For container logs, use container.logs()
+        return
+        yield  # Make this a generator
+
+
+# Global singleton for Docker context
+_docker_context: DockerContext | None = None
+
+
+def get_docker_context() -> DockerContext:
+    """Get the global DockerContext instance."""
+    global _docker_context
+    if _docker_context is None:
+        _docker_context = DockerContext()
+    return _docker_context
+
+
+def is_docker_available() -> bool:
+    """Check if Docker is available."""
+    if not DOCKER_AVAILABLE:
+        return False
+    try:
+        client = docker.from_env()
+        client.ping()
+        return True
+    except Exception:
+        return False

procler/core/context_local.py

@@ -0,0 +1,287 @@
+"""Local subprocess execution context."""
+
+import asyncio
+import os
+from collections.abc import AsyncIterator, Callable
+
+from .context_base import ExecResult, ExecutionContext, ProcessHandle
+
+
+class ManagedProcess:
+    """Internal wrapper for a managed asyncio subprocess."""
+
+    def __init__(
+        self,
+        process: asyncio.subprocess.Process,
+        stdout_task: asyncio.Task | None = None,
+        stderr_task: asyncio.Task | None = None,
+    ):
+        self.process = process
+        self.stdout_task = stdout_task
+        self.stderr_task = stderr_task
+        self._stdout_lines: list[str] = []
+        self._stderr_lines: list[str] = []
+
+    @property
+    def pid(self) -> int:
+        return self.process.pid
+
+    def is_running(self) -> bool:
+        return self.process.returncode is None
+
+    async def wait(self) -> int:
+        """Wait for process to complete and return exit code."""
+        return await self.process.wait()
+
+    def terminate(self) -> None:
+        """Send SIGTERM to the process."""
+        if self.is_running():
+            self.process.terminate()
+
+    def kill(self) -> None:
+        """Send SIGKILL to the process."""
+        if self.is_running():
+            self.process.kill()
+
+    async def cancel_io_tasks(self) -> None:
+        """Cancel the stdout/stderr reading tasks."""
+        for task in [self.stdout_task, self.stderr_task]:
+            if task and not task.done():
+                task.cancel()
+                try:
+                    await task
+                except asyncio.CancelledError:
+                    pass
+
+
+class LocalContext(ExecutionContext):
+    """Execute processes as local subprocesses using asyncio."""
+
+    def __init__(self):
+        self._processes: dict[int, ManagedProcess] = {}
+
+    @property
+    def context_type(self) -> str:
+        return "local"
+
+    async def start_process(
+        self,
+        command: str,
+        cwd: str | None = None,
+        env: dict[str, str] | None = None,
+        on_stdout: Callable[[str], None] | None = None,
+        on_stderr: Callable[[str], None] | None = None,
+        on_exit: Callable[[int], None] | None = None,
+    ) -> ProcessHandle:
+        """Start a long-running local subprocess."""
+        # Merge environment
+        process_env = os.environ.copy()
+        if env:
+            process_env.update(env)
+
+        # Start the process
+        process = await asyncio.create_subprocess_shell(
+            command,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+            cwd=cwd,
+            env=process_env,
+            start_new_session=True,  # Allows killing the whole process group
+        )
+
+        managed = ManagedProcess(process)
+
+        # Start tasks to read stdout/stderr
+        if process.stdout:
+            managed.stdout_task = asyncio.create_task(self._read_stream(process.stdout, "stdout", managed, on_stdout))
+
+        if process.stderr:
+            managed.stderr_task = asyncio.create_task(self._read_stream(process.stderr, "stderr", managed, on_stderr))
+
+        # Start a task to monitor process exit
+        if on_exit:
+            asyncio.create_task(self._monitor_exit(managed, on_exit))
+
+        self._processes[process.pid] = managed
+
+        return ProcessHandle(pid=process.pid, context_type=self.context_type)
+
+    async def _read_stream(
+        self,
+        stream: asyncio.StreamReader,
+        stream_type: str,
+        managed: ManagedProcess,
+        callback: Callable[[str], None] | None,
+        read_timeout: float = 60.0,
+    ) -> None:
+        """Read lines from a stream and invoke callback.
+
+        Args:
+            stream: The stream to read from
+            stream_type: "stdout" or "stderr"
+            managed: The managed process instance
+            callback: Optional callback for each line
+            read_timeout: Timeout for each readline() call (default 60s)
+        """
+        try:
+            while True:
+                try:
+                    # Use timeout to prevent hanging forever on blocked streams
+                    line = await asyncio.wait_for(stream.readline(), timeout=read_timeout)
+                except TimeoutError:
+                    # Check if process is still running
+                    if not managed.is_running():
+                        break
+                    # Process still running but no output - continue waiting
+                    continue
+
+                if not line:
+                    break
+                decoded = line.decode("utf-8", errors="replace").rstrip("\n\r")
+                if stream_type == "stdout":
+                    managed._stdout_lines.append(decoded)
+                else:
+                    managed._stderr_lines.append(decoded)
+                if callback:
+                    callback(decoded)
+        except asyncio.CancelledError:
+            pass
+
+    async def _monitor_exit(
+        self,
+        managed: ManagedProcess,
+        callback: Callable[[int], None],
+    ) -> None:
+        """Monitor process exit and invoke callback."""
+        exit_code = await managed.wait()
+        callback(exit_code)
+
+    async def stop_process(self, handle: ProcessHandle, timeout: float = 10.0) -> int:
+        """Stop a running process gracefully, then forcefully if needed."""
+        managed = self._processes.get(handle.pid)
+        if not managed:
+            return -1
+
+        try:
+            if not managed.is_running():
+                exit_code = managed.process.returncode or 0
+                return exit_code
+
+            # Try graceful termination first
+            managed.terminate()
+
+            try:
+                exit_code = await asyncio.wait_for(managed.wait(), timeout=timeout)
+            except TimeoutError:
+                # Force kill if graceful shutdown times out
+                managed.kill()
+                exit_code = await managed.wait()
+
+            return exit_code
+        finally:
+            # Always cleanup IO tasks and remove from tracking dict
+            try:
+                await managed.cancel_io_tasks()
+            except Exception:
+                pass  # Best effort cleanup
+            self._processes.pop(handle.pid, None)
+
+    async def is_running(self, handle: ProcessHandle) -> bool:
+        """Check if a process is still running."""
+        managed = self._processes.get(handle.pid)
+        if not managed:
+            return False
+        return managed.is_running()
+
+    async def exec_command(
+        self,
+        command: str,
+        cwd: str | None = None,
+        env: dict[str, str] | None = None,
+        timeout: float | None = None,
+    ) -> ExecResult:
+        """Execute a one-shot command and wait for completion."""
+        process_env = os.environ.copy()
+        if env:
+            process_env.update(env)
+
+        process = await asyncio.create_subprocess_shell(
+            command,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+            cwd=cwd,
+            env=process_env,
+        )
+
+        try:
+            stdout, stderr = await asyncio.wait_for(
+                process.communicate(),
+                timeout=timeout,
+            )
+        except TimeoutError:
+            process.kill()
+            await process.wait()
+            return ExecResult(
+                exit_code=-1,
+                stdout="",
+                stderr=f"Command timed out after {timeout} seconds",
+            )
+
+        return ExecResult(
+            exit_code=process.returncode or 0,
+            stdout=stdout.decode("utf-8", errors="replace"),
+            stderr=stderr.decode("utf-8", errors="replace"),
+        )
+
+    async def stream_logs(
+        self,
+        handle: ProcessHandle,
+        follow: bool = True,
+    ) -> AsyncIterator[tuple[str, str]]:
+        """Stream logs from a running process."""
+        managed = self._processes.get(handle.pid)
+        if not managed:
+            return
+
+        # First yield any buffered lines
+        for line in managed._stdout_lines:
+            yield ("stdout", line)
+        for line in managed._stderr_lines:
+            yield ("stderr", line)
+
+        if not follow:
+            return
+
+        # For follow mode, we'd need a more sophisticated approach
+        # with queues. For now, just poll the buffer.
+        last_stdout_idx = len(managed._stdout_lines)
+        last_stderr_idx = len(managed._stderr_lines)
+
+        while managed.is_running():
+            await asyncio.sleep(0.1)
+
+            # Check for new stdout lines
+            while last_stdout_idx < len(managed._stdout_lines):
+                yield ("stdout", managed._stdout_lines[last_stdout_idx])
+                last_stdout_idx += 1
+
+            # Check for new stderr lines
+            while last_stderr_idx < len(managed._stderr_lines):
+                yield ("stderr", managed._stderr_lines[last_stderr_idx])
+                last_stderr_idx += 1
+
+    def get_managed_process(self, pid: int) -> ManagedProcess | None:
+        """Get a managed process by PID (for internal use)."""
+        return self._processes.get(pid)
+
+
+# Global singleton for the local context
+_local_context: LocalContext | None = None
+
+
+def get_local_context() -> LocalContext:
+    """Get the global LocalContext instance."""
+    global _local_context
+    if _local_context is None:
+        _local_context = LocalContext()
+    return _local_context