openenv/echo_env
6
1# SPDX-License-Identifier: BSD-3-Clause2 3"""4Environment Auto-Discovery System5==================================6 7This module provides automatic discovery of OpenEnv environments by:81. Discovering installed openenv-* packages using importlib.metadata92. Loading manifests (openenv.yaml) from package resources103. Caching results for performance114. Supporting HuggingFace Hub downloads12 13This enables AutoEnv to work without coupling to src/envs/ directory.14"""15 16from __future__ import annotations17 18import importlib19import importlib.metadata20import importlib.resources21import json22import logging23import os24import re25import stat26from dataclasses import asdict, dataclass27from pathlib import Path28from typing import Any, Type29 30import yaml31 32logger = logging.getLogger(__name__)33 34 35@dataclass36class EnvironmentInfo:37 """38 Rich information about a discovered environment.39 40 Attributes:41 env_key: Environment key (e.g., "echo", "coding")42 name: Full environment name (e.g., "echo_env")43 package_name: Package name (e.g., "openenv-echo_env")44 version: Version string45 description: Human-readable description46 client_module_path: Full module path to client (e.g., "echo_env.client")47 client_class_name: Client class name (e.g., "EchoEnv")48 action_class_name: Action class name (e.g., "EchoAction")49 observation_class_name: Observation class name (e.g., "EchoObservation")50 default_image: Default Docker image name (e.g., "echo-env:latest")51 spec_version: OpenEnv spec version (from openenv.yaml)52 manifest: Original manifest data53 """54 55 env_key: str56 name: str57 package_name: str58 version: str59 description: str60 client_module_path: str61 client_class_name: str62 action_class_name: str63 observation_class_name: str64 default_image: str65 spec_version: int | None = None66 manifest: dict[str, Any] | None = None67 68 def get_client_class(self) -> Type[Any]:69 """70 Dynamically import and return the client class.71 72 Returns:73 Client class (e.g., EchoEnv)74 75 Raises:76 ImportError: If module or class cannot be imported77 """78 try:79 module = importlib.import_module(self.client_module_path)80 return getattr(module, self.client_class_name)81 except ImportError as e:82 raise ImportError(83 f"Failed to import {self.client_class_name} from {self.client_module_path}: {e}\n"84 f"Make sure the package '{self.package_name}' is installed: "85 f"pip install {self.package_name}"86 ) from e87 except AttributeError as e:88 raise ImportError(89 f"Class {self.client_class_name} not found in {self.client_module_path}: {e}"90 ) from e91 92 def get_action_class(self) -> Type[Any]:93 """94 Dynamically import and return the action class.95 96 Returns:97 Action class (e.g., EchoAction)98 99 Raises:100 ImportError: If module or class cannot be imported101 """102 try:103 module = importlib.import_module(self.client_module_path)104 return getattr(module, self.action_class_name)105 except ImportError as e:106 raise ImportError(107 f"Failed to import {self.action_class_name} from {self.client_module_path}: {e}\n"108 f"Make sure the package '{self.package_name}' is installed: "109 f"pip install {self.package_name}"110 ) from e111 except AttributeError as e:112 raise ImportError(113 f"Class {self.action_class_name} not found in {self.client_module_path}: {e}"114 ) from e115 116 def get_observation_class(self) -> Type[Any]:117 """118 Dynamically import and return the observation class.119 120 Returns:121 Observation class (e.g., EchoObservation)122 123 Raises:124 ImportError: If module or class cannot be imported125 """126 try:127 module = importlib.import_module(self.client_module_path)128 return getattr(module, self.observation_class_name)129 except ImportError as e:130 raise ImportError(131 f"Failed to import {self.observation_class_name} from {self.client_module_path}: {e}\n"132 f"Make sure the package '{self.package_name}' is installed: "133 f"pip install {self.package_name}"134 ) from e135 except AttributeError as e:136 raise ImportError(137 f"Class {self.observation_class_name} not found in {self.client_module_path}: {e}"138 ) from e139 140 141def _normalize_env_name(name: str) -> str:142 """143 Normalize environment name to standard format.144 145 Args:146 name: Input name (e.g., "echo", "echo-env", "echo_env")147 148 Returns:149 Normalized name (e.g., "echo_env")150 151 Examples:152 >>> _normalize_env_name("echo")153 'echo_env'154 >>> _normalize_env_name("echo-env")155 'echo_env'156 >>> _normalize_env_name("echo_env")157 'echo_env'158 """159 # Remove common suffixes160 name = re.sub(r"[-_]env$", "", name)161 # Convert hyphens to underscores162 name = name.replace("-", "_")163 # Add _env suffix if not present164 if not name.endswith("_env"):165 name = f"{name}_env"166 return name167 168 169def _is_hub_url(name: str) -> bool:170 """171 Check if name is a HuggingFace Hub URL or repo ID.172 173 Args:174 name: Input name175 176 Returns:177 True if it looks like a Hub URL178 179 Examples:180 >>> _is_hub_url("openenv/echo_env")181 True182 >>> _is_hub_url("https://huggingface.co/openenv/echo_env")183 True184 >>> _is_hub_url("echo")185 False186 """187 # Contains org/repo pattern or huggingface.co domain188 return "/" in name or "huggingface.co" in name189 190 191def _infer_class_name(env_name: str, class_type: str) -> str:192 """193 Infer class name from environment name using simple conventions.194 195 Args:196 env_name: Environment name (e.g., "echo_env")197 class_type: Type of class ("client", "action", "observation")198 199 Returns:200 Inferred class name201 202 Examples:203 >>> _infer_class_name("echo_env", "client")204 'EchoEnv'205 >>> _infer_class_name("echo_env", "action")206 'EchoAction'207 """208 # Remove _env suffix for base name209 base_name = env_name.replace("_env", "")210 211 # Convert to PascalCase212 pascal_name = "".join(word.capitalize() for word in base_name.split("_"))213 214 # Add suffix based on type215 if class_type == "client":216 return f"{pascal_name}Env"217 elif class_type == "action":218 return f"{pascal_name}Action"219 elif class_type == "observation":220 return f"{pascal_name}Observation"221 else:222 raise ValueError(f"Unknown class type: {class_type}")223 224 225def _load_manifest_from_package(226 package_name: str, module_name: str227) -> dict[str, Any] | None:228 """229 Load openenv.yaml manifest from an installed package.230 231 Args:232 package_name: Package name (e.g., "openenv-echo_env")233 module_name: Module name (e.g., "echo_env")234 235 Returns:236 Parsed manifest dictionary, or None if not found237 238 """239 try:240 # Try to read openenv.yaml from package241 if hasattr(importlib.resources, "files"):242 # Python 3.9+243 package_files = importlib.resources.files(module_name)244 if (package_files / "openenv.yaml").is_file():245 manifest_text = (package_files / "openenv.yaml").read_text()246 return yaml.safe_load(manifest_text)247 else:248 # Python 3.7-3.8 fallback249 with importlib.resources.open_text(module_name, "openenv.yaml") as f:250 return yaml.safe_load(f)251 except (FileNotFoundError, ModuleNotFoundError, AttributeError):252 logger.debug(f"No openenv.yaml found in {module_name}")253 return None254 except Exception as e:255 logger.warning(f"Failed to load openenv.yaml from {module_name}: {e}")256 return None257 258 259def _create_env_info_from_package(260 package_name: str, module_name: str, version: str261) -> EnvironmentInfo | None:262 """263 Create EnvironmentInfo from an installed package.264 265 Args:266 package_name: Package name (e.g., "openenv-echo_env")267 module_name: Module name (e.g., "echo_env")268 version: Package version269 270 Returns:271 EnvironmentInfo instance, or None if invalid272 """273 # Load manifest274 manifest = _load_manifest_from_package(package_name, module_name)275 276 # Get environment name277 if manifest and "name" in manifest:278 env_name = manifest["name"]279 else:280 # Infer from module name281 env_name = module_name282 283 # Normalize to ensure _env suffix284 if not env_name.endswith("_env"):285 env_name = f"{env_name}_env"286 287 # Determine env_key (e.g., "echo_env" → "echo")288 env_key = env_name.replace("_env", "") if env_name.endswith("_env") else env_name289 290 # Get description291 description = (292 manifest.get("description", f"{env_name} environment")293 if manifest294 else f"{env_name} environment"295 )296 297 # Get spec version298 spec_version = manifest.get("spec_version") if manifest else None299 300 # Determine class names301 # Check if manifest has custom class names (custom format)302 if manifest and "action" in manifest and "observation" in manifest:303 # Custom format (like coding_env)304 client_class_name = _infer_class_name(env_name, "client")305 action_class_name = manifest.get(306 "action", _infer_class_name(env_name, "action")307 )308 observation_class_name = manifest.get(309 "observation", _infer_class_name(env_name, "observation")310 )311 else:312 # Use conventions313 client_class_name = _infer_class_name(env_name, "client")314 action_class_name = _infer_class_name(env_name, "action")315 observation_class_name = _infer_class_name(env_name, "observation")316 317 # Module path is just module_name.client318 client_module_path = f"{module_name}.client"319 320 # Determine default Docker image name321 image_name = env_name.replace("_", "-")322 default_image = f"{image_name}:latest"323 324 return EnvironmentInfo(325 env_key=env_key,326 name=env_name,327 package_name=package_name,328 version=version,329 description=description,330 client_module_path=client_module_path,331 client_class_name=client_class_name,332 action_class_name=action_class_name,333 observation_class_name=observation_class_name,334 default_image=default_image,335 spec_version=spec_version,336 manifest=manifest,337 )338 339 340def _default_cache_file() -> Path:341 """342 Return the per-user discovery cache file path.343 344 Uses a per-user cache directory (``$XDG_CACHE_HOME`` or ``~/.cache``) rather than a345 shared, world-writable temporary directory. A fixed path under the shared temp dir346 lets another local user pre-create the cache file and redirect discovery to347 attacker-controlled import paths (`import_module` on a cached `client_module_path`).348 Per the XDG Base Directory specification, relative `XDG_CACHE_HOME` values are349 ignored so an untrusted working tree cannot supply a victim-owned cache file.350 The fallback home must itself be absolute; otherwise discovery fails closed.351 """352 base = os.environ.get("XDG_CACHE_HOME")353 if base and Path(base).is_absolute():354 root = Path(base)355 else:356 home = Path.home()357 if not home.is_absolute():358 raise RuntimeError(359 "discovery cache requires an absolute home directory when "360 "XDG_CACHE_HOME is unset or relative"361 )362 root = home / ".cache"363 return root / "openenv" / "discovery_cache.json"364 365 366def _is_trusted_stat(info: os.stat_result) -> bool:367 """368 Return whether the file *info* describes is safe to load.369 370 On POSIX the file must be owned by the current user and not writable by group or371 others, so a cache file planted by another user is ignored rather than trusted.372 373 Args:374 info (`os.stat_result`):375 Metadata for the file being considered.376 377 Returns:378 `bool`: `True` if the file is owned by the current user and not379 group/other-writable (always `True` on non-POSIX platforms).380 """381 if os.name != "posix":382 return True383 return info.st_uid == os.getuid() and not (384 info.st_mode & (stat.S_IWGRP | stat.S_IWOTH)385 )386 387 388def _is_trusted_cache_file(path: Path) -> bool:389 """390 Return whether *path* is safe to load.391 392 Prefer `_open_trusted_cache`, which checks the descriptor it hands back. This393 variant resolves the path a second time, so on its own it cannot promise that394 the file inspected is the file later read.395 396 Args:397 path (`Path`):398 The cache file to check.399 400 Returns:401 `bool`: `True` if the file is owned by the current user and not402 group/other-writable (always `True` on non-POSIX platforms).403 """404 if os.name != "posix":405 return True406 try:407 return _is_trusted_stat(path.stat())408 except OSError:409 return False410 411 412def _open_trusted_cache(path: Path) -> int | None:413 """414 Open *path* for reading and return the descriptor only if it is trustworthy.415 416 Checking a path and then opening it are two separate resolutions, and an417 attacker who can write in the cache directory can swap the file for a symlink418 in between, so the file that was vetted is not the file that gets read.419 ``O_NOFOLLOW`` refuses a symlink outright and ``fstat`` inspects the descriptor420 itself, which makes the ownership check and the read the same object.421 422 Args:423 path (`Path`):424 The cache file to open.425 426 Returns:427 `int` or `None`: an open read-only descriptor the caller must close, or428 `None` if the file is missing, is a symlink, is not a regular file, or429 is not owned by the current user.430 """431 # O_NONBLOCK matters before the descriptor can be inspected: opening a FIFO432 # read-only otherwise waits for a writer that an attacker simply never433 # provides, so discovery hangs before any check runs. It has no effect on a434 # regular file, which is the only thing accepted here anyway.435 flags = os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0) | getattr(os, "O_NONBLOCK", 0)436 try:437 fd = os.open(path, flags)438 except OSError:439 # Missing, unreadable, or a symlink (ELOOP under O_NOFOLLOW).440 return None441 try:442 info = os.fstat(fd)443 if not stat.S_ISREG(info.st_mode) or not _is_trusted_stat(info):444 os.close(fd)445 return None446 except OSError:447 os.close(fd)448 return None449 return fd450 451 452class EnvironmentDiscovery:453 """454 Auto-discovery system for OpenEnv environments using installed packages.455 456 This class discovers installed openenv-* packages and loads their metadata.457 """458 459 def __init__(self):460 """Initialize discovery system."""461 self._cache: dict[str, EnvironmentInfo] | None = None462 self._cache_file = _default_cache_file()463 464 def _discover_installed_packages(self) -> dict[str, EnvironmentInfo]:465 """466 Discover all installed openenv-* packages.467 468 Returns:469 Dictionary mapping env_key to EnvironmentInfo470 """471 environments: dict[str, EnvironmentInfo] = {}472 473 # Invalidate import caches to ensure we pick up newly installed packages474 importlib.invalidate_caches()475 476 # Get all installed packages477 try:478 distributions = importlib.metadata.distributions()479 except Exception as e:480 logger.warning(f"Failed to get installed packages: {e}")481 return environments482 483 # Filter for openenv-* environment packages; ignore the legacy core wheel.484 for dist in distributions:485 package_name = dist.metadata["Name"]486 487 if not package_name.startswith("openenv-"):488 continue489 490 if package_name == "openenv-core":491 continue492 493 # Get module name (e.g., "openenv-echo_env" → "echo_env")494 module_name = package_name.replace("openenv-", "").replace("-", "_")495 496 # Get version497 version = dist.version498 499 try:500 # Create environment info501 env_info = _create_env_info_from_package(502 package_name, module_name, version503 )504 505 if env_info:506 environments[env_info.env_key] = env_info507 logger.debug(508 f"Discovered environment: {env_info.env_key} ({package_name})"509 )510 511 except Exception as e:512 logger.warning(f"Failed to load environment from {package_name}: {e}")513 continue514 515 return environments516 517 def _load_cache(self) -> dict[str, EnvironmentInfo] | None:518 """519 Load cached discovery results.520 521 Returns:522 Dictionary of env_key -> EnvironmentInfo, or None if cache invalid523 """524 # Only trust a cache file owned by the current user. This prevents another525 # local user from planting a file that would redirect discovery (and the526 # subsequent import_module) to attacker-controlled modules/classes. The527 # descriptor is what gets vetted and then read, so the file cannot be528 # swapped for a symlink after the check.529 fd = _open_trusted_cache(self._cache_file)530 if fd is None:531 if self._cache_file.exists():532 logger.warning(533 f"Ignoring discovery cache {self._cache_file}: not owned by "534 "the current user, writable by group/others, or a symlink."535 )536 return None537 538 try:539 with os.fdopen(fd, "r") as f:540 cache_data = json.load(f)541 542 # Reconstruct EnvironmentInfo objects543 cache: dict[str, EnvironmentInfo] = {}544 for env_key, env_data in cache_data.items():545 cache[env_key] = EnvironmentInfo(**env_data)546 547 return cache548 except Exception as e:549 logger.warning(f"Failed to load discovery cache: {e}")550 return None551 552 def _save_cache(self, environments: dict[str, EnvironmentInfo]) -> None:553 """554 Save discovery results to cache.555 556 Args:557 environments: Dictionary of env_key -> EnvironmentInfo558 """559 try:560 cache_data = {}561 for env_key, env_info in environments.items():562 cache_data[env_key] = asdict(env_info)563 564 self._cache_file.parent.mkdir(parents=True, exist_ok=True)565 # Write a new file and move it into place rather than writing566 # through whatever already sits at the destination. A mode passed to567 # `os.open` only applies when it creates the file, so truncating an568 # existing group/world-writable cache would keep that mode and569 # publish what was just written; and `os.replace` is atomic, so a570 # reader never observes a half-written cache when serialization571 # fails part way. `O_EXCL` requires the temporary name to be new,572 # and 0600 is applied at creation so the umask cannot widen it.573 tmp_path = self._cache_file.with_name(574 f"{self._cache_file.name}.{os.getpid()}.tmp"575 )576 flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0)577 fd = os.open(tmp_path, flags, 0o600)578 try:579 with os.fdopen(fd, "w") as f:580 json.dump(cache_data, f, indent=2)581 os.replace(tmp_path, self._cache_file)582 except BaseException:583 try:584 os.unlink(tmp_path)585 except OSError:586 pass587 raise588 589 except Exception as e:590 logger.warning(f"Failed to save discovery cache: {e}")591 592 def discover(self, use_cache: bool = True) -> dict[str, EnvironmentInfo]:593 """594 Discover all installed OpenEnv environments.595 596 Args:597 use_cache: If True, try to load from cache first598 599 Returns:600 Dictionary mapping env_key to EnvironmentInfo601 602 Examples:603 >>> discovery = EnvironmentDiscovery()604 >>> envs = discovery.discover()605 >>> print(envs.keys())606 dict_keys(['echo', 'coding', ...])607 """608 # Try to load from memory cache first609 if use_cache and self._cache is not None:610 return self._cache611 612 # Try to load from file cache613 if use_cache:614 cached = self._load_cache()615 if cached is not None:616 self._cache = cached617 return self._cache618 619 # Discover from installed packages620 environments = self._discover_installed_packages()621 622 # Save to cache623 self._save_cache(environments)624 self._cache = environments625 626 return environments627 628 def get_environment(self, env_key: str) -> EnvironmentInfo | None:629 """630 Get information about a specific environment.631 632 Args:633 env_key: Environment key (e.g., "echo", "coding")634 635 Returns:636 EnvironmentInfo if found, None otherwise637 638 Examples:639 >>> discovery = EnvironmentDiscovery()640 >>> env = discovery.get_environment("echo")641 >>> print(env.client_class_name)642 'EchoEnv'643 """644 environments = self.discover()645 return environments.get(env_key)646 647 def get_environment_by_name(self, name: str) -> EnvironmentInfo | None:648 """649 Get environment info by flexible name matching.650 651 Args:652 name: Environment name (e.g., "echo", "echo-env", "echo_env")653 654 Returns:655 EnvironmentInfo if found, None otherwise656 """657 # Normalize name to env_key658 normalized = _normalize_env_name(name)659 env_key = normalized.replace("_env", "")660 661 return self.get_environment(env_key)662 663 def list_environments(self) -> None:664 """665 Print a formatted list of all discovered environments.666 667 Examples:668 >>> discovery = EnvironmentDiscovery()669 >>> discovery.list_environments()670 Available OpenEnv Environments:671 ----------------------------------------------------------------------672 echo : Echo Environment (v0.1.0) - openenv-echo_env673 coding : Coding Environment (v0.1.0) - openenv-coding_env674 ...675 """676 environments = self.discover()677 678 print("Available OpenEnv Environments:")679 print("-" * 70)680 681 if not environments:682 print(" No OpenEnv environments found.")683 print(" Install environments with: pip install openenv-<env-name>")684 else:685 for env_key in sorted(environments.keys()):686 env = environments[env_key]687 print(f" {env_key:<15}: {env.description} (v{env.version})")688 print(f" Package: {env.package_name}")689 690 print("-" * 70)691 print(f"Total: {len(environments)} environments")692 693 def clear_cache(self) -> None:694 """Clear the discovery cache."""695 if self._cache_file.exists():696 self._cache_file.unlink()697 self._cache = None698 699 700# Global discovery instance701_global_discovery: EnvironmentDiscovery | None = None702 703 704def get_discovery() -> EnvironmentDiscovery:705 """706 Get or create the global discovery instance.707 708 Returns:709 Global EnvironmentDiscovery instance710 711 Examples:712 >>> discovery = get_discovery()713 >>> envs = discovery.discover()714 """715 global _global_discovery716 717 if _global_discovery is None:718 _global_discovery = EnvironmentDiscovery()719 720 return _global_discovery721 722 723def reset_discovery() -> None:724 """Drop the global discovery singleton (useful for testing).725 726 Does **not** delete the on-disk cache. The cache is now a persistent727 per-user file under `$XDG_CACHE_HOME`; wiping it belongs on an explicit728 [`EnvironmentDiscovery.clear_cache`][] call so test fixtures that only729 need a fresh singleton cannot clobber a developer's real cache.730 """731 global _global_discovery732 _global_discovery = None733