Spaces:
Running on Zero
Running on Zero
| """ | |
| Universal 24/7 Command Engine & Multi-Project Sandboxing | |
| Production-grade execution of arbitrary bash commands with process tracking, | |
| timeout protection, and multi-project workspace isolation. | |
| """ | |
| import os | |
| import sys | |
| import time | |
| import shutil | |
| import asyncio | |
| import logging | |
| import psutil | |
| from typing import Dict, List, Optional, Any | |
| from pydantic import BaseModel, Field | |
| logger = logging.getLogger("antigravity.command_engine") | |
| # Base directory configuration | |
| BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) | |
| PROJECTS_ROOT = os.environ.get("PROJECTS_DIR", os.path.join(BASE_DIR, "projects")) | |
| os.makedirs(PROJECTS_ROOT, exist_ok=True) | |
| class CommandResult(BaseModel): | |
| command: str | |
| project_id: str | |
| exit_code: int | |
| stdout: str | |
| stderr: str | |
| duration_seconds: float | |
| timestamp: float | |
| is_success: bool | |
| class ProjectMeta(BaseModel): | |
| project_id: str | |
| name: str | |
| path: str | |
| size_bytes: int | |
| file_count: int | |
| created_at: float | |
| updated_at: float | |
| class FileItem(BaseModel): | |
| name: str | |
| relative_path: str | |
| is_directory: bool | |
| size_bytes: int | |
| modified_at: float | |
| class RunningProcessTracker: | |
| """Tracks active subprocesses to enable cancellation and monitoring.""" | |
| def __init__(self): | |
| self._active_processes: Dict[str, asyncio.subprocess.Process] = {} | |
| def register(self, task_id: str, process: asyncio.subprocess.Process): | |
| self._active_processes[task_id] = process | |
| def unregister(self, task_id: str): | |
| self._active_processes.pop(task_id, None) | |
| async def terminate(self, task_id: str, timeout: float = 3.0) -> bool: | |
| process = self._active_processes.get(task_id) | |
| if not process: | |
| return False | |
| try: | |
| if process.returncode is None: | |
| process.terminate() | |
| try: | |
| await asyncio.wait_for(process.wait(), timeout=timeout) | |
| except asyncio.TimeoutError: | |
| process.kill() | |
| await process.wait() | |
| return True | |
| except Exception as e: | |
| logger.error(f"Error terminating process {task_id}: {e}") | |
| return False | |
| finally: | |
| self._active_processes.pop(task_id, None) | |
| # Global process tracker instance | |
| process_tracker = RunningProcessTracker() | |
| def get_project_dir(project_id: str) -> str: | |
| """Sanitize and resolve project sandbox directory.""" | |
| # Sanitize project_id to avoid path traversal | |
| safe_id = "".join(c for c in project_id if c.isalnum() or c in ("-", "_", ".")).strip(".") | |
| if not safe_id: | |
| safe_id = "default" | |
| project_path = os.path.abspath(os.path.join(PROJECTS_ROOT, safe_id)) | |
| # Ensure resolved path is within PROJECTS_ROOT | |
| if not project_path.startswith(os.path.abspath(PROJECTS_ROOT)): | |
| project_path = os.path.join(PROJECTS_ROOT, "default") | |
| os.makedirs(project_path, exist_ok=True) | |
| return project_path | |
| async def run_bash_command( | |
| command: str, | |
| project_id: str = "default", | |
| task_id: Optional[str] = None, | |
| timeout_seconds: int = 600, | |
| custom_env: Optional[Dict[str, str]] = None, | |
| log_callback: Optional[Any] = None | |
| ) -> CommandResult: | |
| """ | |
| Executes an arbitrary bash command inside the sandboxed project directory. | |
| Captures stdout/stderr in real-time and provides timeout & cancellation protection. | |
| """ | |
| start_time = time.time() | |
| cwd = get_project_dir(project_id) | |
| env = os.environ.copy() | |
| env["HOME"] = os.path.expanduser("~") | |
| env["PROJECT_DIR"] = cwd | |
| env["PYTHONUNBUFFERED"] = "1" | |
| if custom_env: | |
| env.update(custom_env) | |
| logger.info(f"[{project_id}] Executing command: {command} (timeout={timeout_seconds}s)") | |
| if log_callback: | |
| await log_callback(f"[$] {command}\n") | |
| try: | |
| process = await asyncio.create_subprocess_shell( | |
| command, | |
| stdout=asyncio.subprocess.PIPE, | |
| stderr=asyncio.subprocess.PIPE, | |
| cwd=cwd, | |
| env=env, | |
| shell=True | |
| ) | |
| if task_id: | |
| process_tracker.register(task_id, process) | |
| stdout_chunks: List[str] = [] | |
| stderr_chunks: List[str] = [] | |
| async def read_stream(stream, chunks: List[str], prefix: str = ""): | |
| while True: | |
| line = await stream.readline() | |
| if not line: | |
| break | |
| decoded = line.decode("utf-8", errors="replace") | |
| chunks.append(decoded) | |
| if log_callback: | |
| await log_callback(f"{prefix}{decoded}") | |
| # Stream stdout and stderr concurrently with timeout | |
| try: | |
| await asyncio.wait_for( | |
| asyncio.gather( | |
| read_stream(process.stdout, stdout_chunks), | |
| read_stream(process.stderr, stderr_chunks, prefix="[STDERR] "), | |
| process.wait() | |
| ), | |
| timeout=float(timeout_seconds) | |
| ) | |
| exit_code = process.returncode if process.returncode is not None else 0 | |
| except asyncio.TimeoutError: | |
| logger.warning(f"Command timed out after {timeout_seconds}s: {command}") | |
| process.kill() | |
| await process.wait() | |
| exit_code = -1 | |
| stderr_chunks.append(f"\n[ERROR] Command timed out after {timeout_seconds} seconds.\n") | |
| duration = round(time.time() - start_time, 3) | |
| full_stdout = "".join(stdout_chunks) | |
| full_stderr = "".join(stderr_chunks) | |
| return CommandResult( | |
| command=command, | |
| project_id=project_id, | |
| exit_code=exit_code, | |
| stdout=full_stdout, | |
| stderr=full_stderr, | |
| duration_seconds=duration, | |
| timestamp=start_time, | |
| is_success=(exit_code == 0) | |
| ) | |
| except Exception as e: | |
| duration = round(time.time() - start_time, 3) | |
| err_msg = f"Failed to execute command: {str(e)}" | |
| logger.error(err_msg, exc_info=True) | |
| return CommandResult( | |
| command=command, | |
| project_id=project_id, | |
| exit_code=1, | |
| stdout="", | |
| stderr=err_msg, | |
| duration_seconds=duration, | |
| timestamp=start_time, | |
| is_success=False | |
| ) | |
| finally: | |
| if task_id: | |
| process_tracker.unregister(task_id) | |
| def list_projects() -> List[ProjectMeta]: | |
| """Lists all sandboxed project workspaces with metadata and disk footprints.""" | |
| projects: List[ProjectMeta] = [] | |
| if not os.path.exists(PROJECTS_ROOT): | |
| return projects | |
| for entry in os.scandir(PROJECTS_ROOT): | |
| if entry.is_dir(): | |
| project_path = entry.path | |
| total_size = 0 | |
| file_count = 0 | |
| try: | |
| for root, _, files in os.walk(project_path): | |
| for f in files: | |
| fp = os.path.join(root, f) | |
| if not os.path.islink(fp): | |
| try: | |
| total_size += os.path.getsize(fp) | |
| file_count += 1 | |
| except OSError: | |
| pass | |
| stat = entry.stat() | |
| projects.append(ProjectMeta( | |
| project_id=entry.name, | |
| name=entry.name, | |
| path=project_path, | |
| size_bytes=total_size, | |
| file_count=file_count, | |
| created_at=stat.st_ctime, | |
| updated_at=stat.st_mtime | |
| )) | |
| except Exception as e: | |
| logger.warning(f"Error reading project {entry.name}: {e}") | |
| # Sort most recently updated first | |
| projects.sort(key=lambda p: p.updated_at, reverse=True) | |
| return projects | |
| def create_project(project_id: str, name: Optional[str] = None) -> ProjectMeta: | |
| """Creates a new project sandbox directory.""" | |
| path = get_project_dir(project_id) | |
| stat = os.stat(path) | |
| return ProjectMeta( | |
| project_id=os.path.basename(path), | |
| name=name or os.path.basename(path), | |
| path=path, | |
| size_bytes=0, | |
| file_count=0, | |
| created_at=stat.st_ctime, | |
| updated_at=stat.st_mtime | |
| ) | |
| def list_project_files(project_id: str, relative_subpath: str = "") -> List[FileItem]: | |
| """Lists files and folders inside a specific project workspace.""" | |
| base_proj_dir = get_project_dir(project_id) | |
| target_dir = os.path.abspath(os.path.join(base_proj_dir, relative_subpath.lstrip("/"))) | |
| if not target_dir.startswith(base_proj_dir): | |
| target_dir = base_proj_dir | |
| if not os.path.exists(target_dir): | |
| return [] | |
| items: List[FileItem] = [] | |
| try: | |
| for entry in os.scandir(target_dir): | |
| try: | |
| stat = entry.stat() | |
| rel_path = os.path.relpath(entry.path, base_proj_dir) | |
| items.append(FileItem( | |
| name=entry.name, | |
| relative_path=rel_path, | |
| is_directory=entry.is_dir(), | |
| size_bytes=stat.st_size if not entry.is_dir() else 0, | |
| modified_at=stat.st_mtime | |
| )) | |
| except OSError: | |
| continue | |
| except Exception as e: | |
| logger.error(f"Failed to list files for project {project_id}: {e}") | |
| items.sort(key=lambda x: (not x.is_directory, x.name.lower())) | |
| return items | |
| def delete_project(project_id: str) -> bool: | |
| """Safely removes a project sandbox directory.""" | |
| if project_id.lower() in ("default", "system"): | |
| return False | |
| path = get_project_dir(project_id) | |
| if os.path.exists(path) and path.startswith(os.path.abspath(PROJECTS_ROOT)): | |
| shutil.rmtree(path, ignore_errors=True) | |
| return True | |
| return False | |
| def get_system_telemetry() -> Dict[str, Any]: | |
| """Retrieves hardware, CPU, RAM, 50GB disk capacity, and OS telemetry.""" | |
| cpu_percent = psutil.cpu_percent(interval=None) | |
| cpu_count = psutil.cpu_count(logical=True) | |
| virtual_mem = psutil.virtual_memory() | |
| # Hugging Face Spaces free tier gives 50GB disk | |
| disk_usage = shutil.disk_usage(BASE_DIR) | |
| total_gb = round(disk_usage.total / (1024 ** 3), 2) | |
| used_gb = round(disk_usage.used / (1024 ** 3), 2) | |
| free_gb = round(disk_usage.free / (1024 ** 3), 2) | |
| disk_percent = round((disk_usage.used / disk_usage.total) * 100, 1) | |
| boot_time = psutil.boot_time() | |
| uptime_seconds = int(time.time() - boot_time) | |
| return { | |
| "cpu_usage_percent": cpu_percent, | |
| "cpu_cores": cpu_count, | |
| "memory_total_mb": round(virtual_mem.total / (1024 ** 2), 1), | |
| "memory_used_mb": round(virtual_mem.used / (1024 ** 2), 1), | |
| "memory_free_mb": round(virtual_mem.available / (1024 ** 2), 1), | |
| "memory_percent": virtual_mem.percent, | |
| "disk_total_gb": total_gb, | |
| "disk_used_gb": used_gb, | |
| "disk_free_gb": free_gb, | |
| "disk_percent": disk_percent, | |
| "uptime_seconds": uptime_seconds, | |
| "active_processes_count": len(process_tracker._active_processes) | |
| } | |