refactor(sandboxes): introduce contribution-based provider registry (#3330)

## Related issue

N/A

## Summary

- Add `omnigent/onboarding/sandboxes/types.py` with shared dataclasses
  (`SandboxCapabilities`, `SandboxSpec`, `SandboxInfo`, `HostContext`) and the
  new `SandboxError` exception hierarchy.
- Add `omnigent/onboarding/sandboxes/registry.py` with a contribution-based
  provider registry that mirrors `omnigent/harness_plugins.py`: built-in
  providers are declared as a `SandboxProviderContribution`, community
  packages register via the `omnigent.sandbox_providers` entrypoint group, and
  broken plugins are recorded in `load_errors` without breaking core startup.
- Add `omnigent/community/sandbox/__init__.py` as a namespace package so
  third-party providers can ship code under `omnigent.community.sandbox.*`.
- Validation enforces that community provider code lives under the community
  namespace, rejects name collisions, and checks metadata consistency.
- Add a `capabilities` property to `SandboxLauncher` that derives feature flags
  from existing class variables and overridden transport methods.
- Migrate CLI and managed-host call sites from direct class-var reads
  (`supports_cli_bootstrap`, `can_resume`, `supports_local_port_forward`) to
  the new `capabilities` object.
- Add unit tests for types, registry behavior, validation, and entrypoint
  discovery.

No provider implementations were changed; this is purely a surface-layer
refactor toward a pluggable sandbox provider interface.

## Test Plan

```bash
uv run pytest tests/onboarding/sandboxes tests/sandbox tests/cli/test_cli.py tests/server -k managed -q
pre-commit run --files omnigent/onboarding/sandboxes/types.py omnigent/onboarding/sandboxes/registry.py omnigent/onboarding/sandboxes/base.py omnigent/onboarding/sandboxes/__init__.py omnigent/onboarding/sandboxes/bootstrap.py omnigent/community/sandbox/__init__.py omnigent/cli_sandbox.py omnigent/server/managed_hosts.py tests/onboarding/sandboxes/test_types.py tests/onboarding/sandboxes/test_registry.py
```

All 779 selected tests pass and the targeted pre-commit hooks pass.

## Demo

N/A

## Type of change

- [ ] Bug fix
- [ ] Feature
- [ ] UI / frontend change
- [x] Refactor / chore
- [ ] Docs
- [ ] Test / CI
- [ ] Breaking change

## Test coverage

- [x] Unit tests added / updated
- [ ] Integration tests added / updated
- [ ] E2E tests added / updated
- [ ] Manual verification completed
- [x] Existing tests cover this change
- [ ] Not applicable

## Coverage notes

New unit tests in `tests/onboarding/sandboxes/test_types.py` and
`tests/onboarding/sandboxes/test_registry.py` exercise the registry,
contribution validation, types, and capabilities derivation. Existing
provider and CLI tests pass unchanged, confirming backward compatibility.

Signed-off-by: Zeyi (Rice) Fan <zeyi.f@databricks.com>
This commit is contained in:
Zeyi (Rice) Fan
2026-07-26 23:44:36 -07:00
committed by GitHub
parent c1bedeaabd
commit a7ef194c4f
10 changed files with 980 additions and 39 deletions
+2 -2
View File
@@ -95,7 +95,7 @@ def _require_cli_bootstrap(launcher: SandboxLauncher) -> None:
:raises click.ClickException: When the provider has no CLI
bootstrap flow.
"""
if not launcher.supports_cli_bootstrap:
if not launcher.capabilities.cli_bootstrap:
raise click.ClickException(
f"The '{launcher.provider}' provider supports server-managed "
"sessions only — create one with "
@@ -276,7 +276,7 @@ def sandbox_create(
# The in-sandbox login only exists for providers that can forward
# the browser's callback port — others skip it automatically, no
# --no-auth acknowledgement required.
if not launcher.supports_local_port_forward:
if not launcher.capabilities.local_port_forward:
skip_auth = True
sandbox_id = bootstrap_sandbox_host(
launcher,
+18
View File
@@ -0,0 +1,18 @@
"""Namespace package for optional community sandbox implementations."""
from __future__ import annotations
import sys
from pathlib import Path
from pkgutil import extend_path
__path__ = extend_path(__path__, __name__)
_seen = set(__path__)
for _entry in sys.path:
_candidate = Path(_entry).joinpath(*__name__.split("."))
if _candidate.is_dir():
_candidate_str = str(_candidate.resolve())
if _candidate_str not in _seen:
__path__.append(_candidate_str)
_seen.add(_candidate_str)
+65 -31
View File
@@ -1,17 +1,16 @@
"""
Sandbox launchers: run Omnigent hosts in remote sandboxes.
"""Sandbox launchers: run Omnigent hosts in remote sandboxes.
Public API for the ``omnigent sandbox`` CLI and anything else that
bootstraps a sandbox-backed host. Providers are registered by name in
:data:`_LAUNCHERS`; launcher modules may be absent from a given
distribution (e.g. the Databricks Lakebox launcher), in which case the
provider simply isn't offered.
bootstraps a sandbox-backed host. Core Omnigent contributes built-in
providers directly; third-party packages can contribute providers through
the ``omnigent.sandbox_providers`` entry point group.
"""
from __future__ import annotations
import importlib
import importlib.util
import warnings
import click
@@ -32,28 +31,64 @@ from omnigent.onboarding.sandboxes.bootstrap import (
set_sandbox_host_name,
ship_wheels,
)
from omnigent.onboarding.sandboxes.registry import (
COMMUNITY_MODULE_PREFIX,
SandboxProviderContribution,
SandboxProviderMetadata,
SandboxProviderPluginState,
SandboxRegistryError,
available_providers,
get_provider_metadata,
instantiate,
plugin_state,
reset_plugin_state_for_tests,
)
from omnigent.onboarding.sandboxes.types import (
HostContext,
SandboxCapabilities,
SandboxCommandError,
SandboxConfigError,
SandboxError,
SandboxInfo,
SandboxSpec,
)
__all__ = [
"COMMUNITY_MODULE_PREFIX",
"DEFAULT_SANDBOX_NAME",
"DerivedWorkspace",
"HostContext",
"RemoteCommandResult",
"RemoteProcess",
"SandboxCapabilities",
"SandboxCapabilityError",
"SandboxCommandError",
"SandboxConfigError",
"SandboxError",
"SandboxInfo",
"SandboxLauncher",
"SandboxProviderContribution",
"SandboxProviderMetadata",
"SandboxProviderPluginState",
"SandboxRegistryError",
"SandboxSpec",
"available_providers",
"bootstrap_sandbox_host",
"build_wheels",
"connect_sandbox_host",
"derive_workspace",
"get_launcher",
"get_provider_metadata",
"login_app_oauth_in_sandbox",
"plugin_state",
"reset_plugin_state_for_tests",
"set_sandbox_host_name",
"ship_wheels",
]
# Provider name → "module:ClassName" of its SandboxLauncher. Modules are
# imported lazily (some pull in optional SDKs) and may be absent from a
# distribution entirely (e.g. lakebox).
# Legacy registration surface. It is no longer used by the registry, but is kept
# as a fallback path so any provider not yet loaded through the entrypoint
# mechanism still resolves for one release cycle.
_LAUNCHERS: dict[str, str] = {
"lakebox": "omnigent.onboarding.sandboxes.lakebox:LakeboxLauncher",
"modal": "omnigent.onboarding.sandboxes.modal:ModalSandboxLauncher",
@@ -73,27 +108,6 @@ _LAUNCHERS: dict[str, str] = {
}
def available_providers() -> tuple[str, ...]:
"""
List the sandbox providers whose launcher modules exist in this
build.
Uses ``find_spec`` (no import side effects), so it is cheap enough
to call at CLI startup to decide whether to register the
``omnigent sandbox`` command group.
:returns: Provider names in registration order, e.g.
``("lakebox", "modal")`` internally or ``("modal",)`` in the
OSS build (where the lakebox module is excluded).
"""
available: list[str] = []
for name, target in _LAUNCHERS.items():
module_name = target.partition(":")[0]
if importlib.util.find_spec(module_name) is not None:
available.append(name)
return tuple(available)
def get_launcher(provider: str, *, workspace_host: str | None = None) -> SandboxLauncher:
"""
Resolve a provider name to a launcher instance.
@@ -113,8 +127,28 @@ def get_launcher(provider: str, *, workspace_host: str | None = None) -> Sandbox
:raises click.ClickException: If the provider is unknown or its
launcher module is not present in this build.
"""
if provider in plugin_state():
try:
return instantiate(provider, workspace_host=workspace_host)
except SandboxRegistryError as exc:
raise click.ClickException(str(exc)) from exc
# Legacy fallback for any provider not yet loaded by the registry.
warnings.warn(
f"Sandbox provider '{provider}' was resolved through the legacy "
f"_LAUNCHERS registry. Use the SandboxProviderContribution or "
f"omnigent.sandbox_providers entrypoints instead.",
DeprecationWarning,
stacklevel=2,
)
target = _LAUNCHERS.get(provider)
if target is None or provider not in available_providers():
if target is None:
offered = ", ".join(available_providers()) or "(none in this build)"
raise click.ClickException(
f"Unknown or unavailable sandbox provider '{provider}'. Available: {offered}."
)
module_name = target.partition(":")[0]
if importlib.util.find_spec(module_name) is None:
offered = ", ".join(available_providers()) or "(none in this build)"
raise click.ClickException(
f"Unknown or unavailable sandbox provider '{provider}'. Available: {offered}."
+41 -2
View File
@@ -30,6 +30,7 @@ from typing import TYPE_CHECKING, ClassVar
import click
from omnigent.host.identity import HOST_ID_ENV_VAR, HOST_NAME_ENV_VAR, HOST_TOKEN_ENV_VAR
from omnigent.onboarding.sandboxes import types as _sandbox_types
if TYPE_CHECKING:
from collections.abc import Callable, Iterator
@@ -288,7 +289,7 @@ def render_host_config_write_command(host_config: dict[str, object]) -> str:
return f"python3 -c {shlex.quote(script)}"
class SandboxCapabilityError(click.ClickException):
class SandboxCapabilityError(click.ClickException, _sandbox_types.SandboxError):
"""
Raised when a launcher does not support an optional primitive.
@@ -409,6 +410,44 @@ class SandboxLauncher(ABC):
# of being silently revived onto an empty workspace.
can_resume: ClassVar[bool] = False
# Whether this provider supports the server-managed host flow
# (``host_type="managed"`` sessions). The server checks this before
# launching a sandbox for a managed session. Most providers support both
# CLI bootstrap and managed launch; set this to ``False`` for providers
# that are CLI-only (e.g. Lakebox) or staged-but-not-yet-launched.
supports_managed_launch: ClassVar[bool] = True
@property
def capabilities(self) -> _sandbox_types.SandboxCapabilities:
"""
Feature flags this provider declares.
The returned object is derived from the provider's class vars and
from which optional transport methods it has overridden. It is a
transition shim: providers will set an explicit
:class:`~omnigent.onboarding.sandboxes.types.SandboxCapabilities`
object directly once the refactor is complete.
"""
return _sandbox_types.SandboxCapabilities(
cli_bootstrap=self.supports_cli_bootstrap,
managed_launch=self.supports_managed_launch,
local_port_forward=self.supports_local_port_forward,
resume_stopped=self.can_resume,
programmatic_terminate=self._is_capability_overridden("terminate"),
file_copy=self._is_capability_overridden("put"),
streaming_exec=self._is_capability_overridden("stream_exec"),
foreground_exec=self._is_capability_overridden("exec_foreground"),
)
def _is_capability_overridden(self, name: str) -> bool:
"""
Return whether this provider overrides the named optional method.
Used while the refactor is in transition so the capability object
can reflect overridden methods without provider authors touching it.
"""
return getattr(type(self), name) is not getattr(SandboxLauncher, name)
@abstractmethod
def prepare(self) -> None:
"""
@@ -527,7 +566,7 @@ class SandboxLauncher(ABC):
# host_config the cleanup must run: an operator who removed the block
# expects previously injected entries gone on the next wake. Fresh
# sandboxes can't carry a stale marker — skip the extra exec there.
if host_config is not None or self.can_resume:
if host_config is not None or self.capabilities.resume_stopped:
self.run(sandbox_id, render_host_config_write_command(host_config or {}))
env_prefix = " ".join(
f"{key}={shlex.quote(value)}"
+2 -2
View File
@@ -496,7 +496,7 @@ def login_app_oauth_in_sandbox(
# (e.g. Modal) — BEFORE validating flags or touching the sandbox, so
# the user gets the --no-auth hint instead of a misleading error
# from a doomed in-sandbox login.
if not launcher.supports_local_port_forward:
if not launcher.capabilities.local_port_forward:
raise launcher.forward_capability_error()
if server_url is None:
raise click.ClickException(
@@ -722,7 +722,7 @@ def bootstrap_sandbox_host(
# up front so a misconfigured call fails before the wheel build and
# ship already ran. (The CLI skips auth automatically for providers
# without the capability; this backstops programmatic callers.)
if not skip_auth and not launcher.supports_local_port_forward:
if not skip_auth and not launcher.capabilities.local_port_forward:
raise launcher.forward_capability_error()
launcher.prepare()
if sandbox_id is None:
+308
View File
@@ -0,0 +1,308 @@
"""Sandbox provider registry.
Core Omnigent contributes built-in sandbox providers directly. Optional
community packages contribute additional providers through the
``omnigent.sandbox_providers`` entry point group.
The registry mirrors the design of ``omnigent.harness_plugins``:
contributions are merged into a single plugin state, validation keeps
community provider code in the ``omnigent.community.sandbox`` namespace, and
broken plugins are recorded but never break core startup.
"""
from __future__ import annotations
import importlib
import importlib.metadata
import importlib.util
import logging
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any
from omnigent.onboarding.sandboxes.base import SandboxLauncher
if TYPE_CHECKING:
pass
logger = logging.getLogger(__name__)
COMMUNITY_ENTRY_POINT_GROUP = "omnigent.sandbox_providers"
COMMUNITY_MODULE_PREFIX = "omnigent.community.sandbox."
class SandboxRegistryError(Exception):
"""A provider registry operation failed."""
@dataclass(frozen=True)
class SandboxProviderMetadata:
"""Static metadata for one sandbox provider.
:param name: Provider short name, e.g. ``"modal"``.
:param launcher_class: Fully-qualified launcher class as
``"module.path:ClassName"`` or ``"module.path.ClassName"``.
:param config_model: Optional Pydantic model that validates the
provider-specific ``sandbox.<name>`` config block.
:param managed_token_ttl_s: Optional default managed launch-token
lifetime in seconds.
"""
name: str
launcher_class: str
config_model: type[object] | None = None
managed_token_ttl_s: int | None = None
@dataclass(frozen=True)
class SandboxProviderContribution:
"""One package's contribution to the sandbox provider registry.
:param name: Package / contributor name, e.g. ``"omnigent-acme"``. Used
only for error messages.
:param providers: Mapping of provider short name to metadata.
"""
name: str
providers: dict[str, SandboxProviderMetadata] = field(default_factory=dict)
@dataclass(frozen=True)
class SandboxProviderPluginState:
"""Merged provider registry plus non-fatal plugin load errors."""
providers: dict[str, SandboxProviderMetadata]
load_errors: dict[str, str]
def __contains__(self, name: str) -> bool:
return name in self.providers
def get(self, name: str) -> SandboxProviderMetadata | None:
return self.providers.get(name)
def names(self) -> tuple[str, ...]:
return tuple(self.providers)
_state: SandboxProviderPluginState | None = None
def _entry_points() -> tuple[importlib.metadata.EntryPoint, ...]:
"""Return community sandbox-provider entrypoints, tolerating old stdlib."""
discovered = importlib.metadata.entry_points()
if hasattr(discovered, "select"):
return tuple(discovered.select(group=COMMUNITY_ENTRY_POINT_GROUP))
return tuple(discovered.get(COMMUNITY_ENTRY_POINT_GROUP, ()))
def _load_object(import_path: str) -> Any:
"""Load ``module:attribute`` or ``module.attribute``."""
if ":" in import_path:
module_name, attr = import_path.split(":", 1)
else:
module_name, attr = import_path.rsplit(".", 1)
module = importlib.import_module(module_name)
return getattr(module, attr)
def _launcher_module(import_path: str) -> str:
"""Return the module part of a launcher class import path."""
if ":" in import_path:
return import_path.split(":", 1)[0]
return import_path.rsplit(".", 1)[0]
def _config_model_module(config_model: type[object]) -> str | None:
"""Return the module of a config model class, or ``None``."""
return getattr(config_model, "__module__", None)
def _builtin_contribution() -> SandboxProviderContribution:
"""The built-in sandbox-provider contribution from core Omnigent."""
return SandboxProviderContribution(
name="omnigent",
providers={
"lakebox": SandboxProviderMetadata(
name="lakebox",
launcher_class="omnigent.onboarding.sandboxes.lakebox:LakeboxLauncher",
),
"modal": SandboxProviderMetadata(
name="modal",
launcher_class="omnigent.onboarding.sandboxes.modal:ModalSandboxLauncher",
managed_token_ttl_s=25 * 3600,
),
"daytona": SandboxProviderMetadata(
name="daytona",
launcher_class="omnigent.onboarding.sandboxes.daytona:DaytonaSandboxLauncher",
managed_token_ttl_s=7 * 24 * 3600,
),
"boxlite": SandboxProviderMetadata(
name="boxlite",
launcher_class="omnigent.onboarding.sandboxes.boxlite:BoxliteSandboxLauncher",
managed_token_ttl_s=7 * 24 * 3600,
),
"cwsandbox": SandboxProviderMetadata(
name="cwsandbox",
launcher_class="omnigent.onboarding.sandboxes.cwsandbox:CWSandboxLauncher",
),
"islo": SandboxProviderMetadata(
name="islo",
launcher_class="omnigent.onboarding.sandboxes.islo:IsloSandboxLauncher",
managed_token_ttl_s=7 * 24 * 3600,
),
"e2b": SandboxProviderMetadata(
name="e2b",
launcher_class="omnigent.onboarding.sandboxes.e2b:E2BSandboxLauncher",
),
"openshell": SandboxProviderMetadata(
name="openshell",
launcher_class="omnigent.onboarding.sandboxes.openshell:OpenShellSandboxLauncher",
managed_token_ttl_s=7 * 24 * 3600,
),
"kubernetes": SandboxProviderMetadata(
name="kubernetes",
launcher_class="omnigent.onboarding.sandboxes.kubernetes:KubernetesSandboxLauncher",
managed_token_ttl_s=7 * 24 * 3600,
),
},
)
def _provider_module_available(import_path: str) -> bool:
"""Return whether the module holding a launcher class is importable."""
module_name = _launcher_module(import_path)
return importlib.util.find_spec(module_name) is not None
def _validate_community_contribution(
contribution: SandboxProviderContribution,
*,
entry_point_name: str,
existing: dict[str, SandboxProviderMetadata],
) -> str | None:
"""Validate a community contribution. Returns an error string or ``None``."""
if not contribution.name:
return "sandbox provider plugin must set name"
for name, meta in contribution.providers.items():
if not meta.name:
return f"provider {name!r} must set name"
if name != meta.name:
return f"provider key {name!r} does not match metadata name {meta.name!r}"
if not meta.launcher_class:
return f"provider {name!r} must set launcher_class"
if name in existing:
return (
f"sandbox provider plugin {entry_point_name!r} attempts to "
f"override existing provider {name!r}"
)
module = _launcher_module(meta.launcher_class)
if not module.startswith(COMMUNITY_MODULE_PREFIX):
return (
f"sandbox provider plugin {entry_point_name!r} provider "
f"{name!r} uses module {module!r}; expected "
f"{COMMUNITY_MODULE_PREFIX}*"
)
if meta.config_model is not None:
cfg_module = _config_model_module(meta.config_model)
if cfg_module is None or not cfg_module.startswith(COMMUNITY_MODULE_PREFIX):
return (
f"sandbox provider plugin {entry_point_name!r} config model "
f"for provider {name!r} must live under "
f"{COMMUNITY_MODULE_PREFIX}*"
)
return None
def plugin_state() -> SandboxProviderPluginState:
"""Return the merged built-in + community sandbox provider registry."""
global _state
if _state is not None:
return _state
built_in = _builtin_contribution().providers
providers: dict[str, SandboxProviderMetadata] = {}
load_errors: dict[str, str] = {}
for name, meta in built_in.items():
if _provider_module_available(meta.launcher_class):
providers[name] = meta
for entry_point in _entry_points():
try:
loaded = entry_point.load()
contribution = loaded() if callable(loaded) else loaded
if not isinstance(contribution, SandboxProviderContribution):
raise TypeError(
f"entry point returned {type(contribution).__name__}, "
"expected SandboxProviderContribution"
)
error = _validate_community_contribution(
contribution,
entry_point_name=entry_point.name,
existing=providers,
)
if error is not None:
raise ValueError(error)
providers.update(contribution.providers)
except Exception as exc:
load_errors[entry_point.name] = str(exc)
logger.warning(
"could not load sandbox provider entry point %s (%s)",
entry_point.name,
exc,
exc_info=True,
)
_state = SandboxProviderPluginState(providers=providers, load_errors=load_errors)
return _state
def reset_plugin_state_for_tests() -> None:
"""Clear the cached plugin state."""
global _state
_state = None
def available_providers() -> tuple[str, ...]:
"""All sandbox providers whose launcher modules are available."""
return plugin_state().names()
def get_provider_metadata(name: str) -> SandboxProviderMetadata | None:
"""Return metadata for a provider, or ``None`` if it is not registered."""
return plugin_state().get(name)
def instantiate(
name: str,
*,
workspace_host: str | None = None,
) -> object:
"""Import and instantiate a registered provider's launcher class.
:param name: Registered provider name.
:param workspace_host: Optional Databricks workspace host passed
to the Lakebox launcher constructor.
:returns: A fresh launcher instance.
:raises SandboxRegistryError: If the provider is unknown or its
class cannot be imported/instantiated.
"""
meta = get_provider_metadata(name)
if meta is None:
raise SandboxRegistryError(f"unknown sandbox provider '{name}'")
try:
launcher_cls = _load_object(meta.launcher_class)
except Exception as exc:
raise SandboxRegistryError(
f"could not load sandbox provider '{name}' from {meta.launcher_class!r}: {exc}"
) from exc
if not isinstance(launcher_cls, type) or not issubclass(launcher_cls, SandboxLauncher):
raise SandboxRegistryError(
f"sandbox provider '{name}' resolved to {launcher_cls!r}, "
f"which is not a SandboxLauncher subclass"
)
if name == "lakebox" and workspace_host is not None:
return launcher_cls(workspace_host=workspace_host)
return launcher_cls()
+126
View File
@@ -0,0 +1,126 @@
"""Public types for the sandbox launcher surface.
These dataclasses and exceptions are the vocabulary used by both the
existing :class:`~omnigent.onboarding.sandboxes.base.SandboxLauncher`
interface and the newer pluggable surface in
:mod:`omnigent.onboarding.sandboxes.registry`.
"""
from __future__ import annotations
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
from collections.abc import Callable
class SandboxError(Exception):
"""Base for all sandbox-provider errors."""
class SandboxConfigError(SandboxError):
"""Sandbox provider configuration is malformed or unavailable."""
class SandboxAuthError(SandboxError):
"""Provider credentials or local tooling are missing/invalid."""
class SandboxCommandError(SandboxError):
"""A command executed inside a sandbox failed.
:param message: Human-readable reason.
:param command: The remote command that failed.
:param returncode: Remote exit code.
:param stdout: Captured standard output.
:param stderr: Captured standard error.
"""
def __init__(
self,
message: str,
*,
command: str | None = None,
returncode: int | None = None,
stdout: str | None = None,
stderr: str | None = None,
) -> None:
super().__init__(message)
self.command = command
self.returncode = returncode
self.stdout = stdout
self.stderr = stderr
@dataclass(frozen=True)
class SandboxCapabilities:
"""Feature flags declared by a sandbox provider.
Providers advertise which primitives they support so callers can fail
fast and surface actionable messages.
:param cli_bootstrap: Provider supports ``omnigent sandbox create`` /
``connect`` (``put`` / ``stream_exec`` / ``exec_foreground`` /
``wheel_install_command``).
:param managed_launch: Provider supports server-managed
``host_type="managed"`` sessions (``prepare`` / ``provision`` /
``start_host``).
:param local_port_forward: Provider can bridge a local port into the
sandbox for the App OAuth callback flow.
:param resume_stopped: Provider can resume a stopped sandbox in place
with its persistent volume.
:param programmatic_terminate: Provider can terminate a sandbox
programmatically.
:param file_copy: Provider supports copying files into the sandbox.
:param streaming_exec: Provider supports streaming process execution
inside the sandbox.
:param foreground_exec: Provider supports a foreground exec that
inherits local stdio.
"""
cli_bootstrap: bool = False
managed_launch: bool = False
local_port_forward: bool = False
resume_stopped: bool = False
programmatic_terminate: bool = False
file_copy: bool = False
streaming_exec: bool = False
foreground_exec: bool = False
@dataclass(frozen=True)
class SandboxSpec:
"""Provider-agnostic description of a sandbox to provision."""
name: str
image: str | None = None
cpu: float | None = None
memory_mib: int | None = None
disk_gb: int | None = None
lifetime_s: int | None = None
tags: dict[str, str] = field(default_factory=dict)
@dataclass(frozen=True)
class SandboxInfo:
"""Result of a successful provision or attach."""
sandbox_id: str
workspace_path: str | None = None
metadata: dict[str, Any] = field(default_factory=dict)
@dataclass
class HostContext:
"""Context handed to ``start_host`` when launching a managed host."""
token: str
host_id: str
host_name: str
server_url: str
repo_url: str | None = None
repo_branch: str | None = None
repo_name: str | None = None
host_config: dict[str, object] = field(default_factory=dict)
on_stage: Callable[[str], None] | None = None
+6 -2
View File
@@ -2307,7 +2307,11 @@ def host_resume_supported(
host with no recorded ``sandbox_id``.
"""
launcher = _launcher_for_teardown(host, config)
return launcher is not None and launcher.can_resume and host.sandbox_id is not None
return (
launcher is not None
and launcher.capabilities.resume_stopped
and host.sandbox_id is not None
)
def host_sandbox_is_running(
@@ -2386,7 +2390,7 @@ async def resume_managed_host(
# Resume needs a reattachable volume; others (e.g. Modal) fall through to
# the caller's host-offline path (the user starts a new session).
launcher = _launcher_for_teardown(host, config)
if launcher is None or not launcher.can_resume or host.sandbox_id is None:
if launcher is None or not launcher.capabilities.resume_stopped or host.sandbox_id is None:
return
sandbox_id = host.sandbox_id
# Single-flight per host (see _resume_locks).
+313
View File
@@ -0,0 +1,313 @@
"""Tests for :mod:`omnigent.onboarding.sandboxes.registry`."""
from __future__ import annotations
import contextlib
from dataclasses import dataclass
from typing import Any
import click
import pytest
from omnigent.onboarding.sandboxes import get_launcher
from omnigent.onboarding.sandboxes.registry import (
COMMUNITY_MODULE_PREFIX,
SandboxProviderContribution,
SandboxProviderMetadata,
SandboxProviderPluginState,
SandboxRegistryError,
available_providers,
get_provider_metadata,
instantiate,
plugin_state,
reset_plugin_state_for_tests,
)
@dataclass
class _FakeEntryPoint:
"""Minimal stub for importlib.metadata.EntryPoint."""
name: str
value: Any
def load(self) -> Any:
return self.value
class _AcmeLauncher:
"""Dummy launcher used by the entrypoint test."""
provider = "acme"
class _NoOpLauncher:
"""Dummy launcher for built-in registration tests."""
provider = "noop"
NoOpSandboxLauncher = _NoOpLauncher
AcmeSandboxLauncher = _AcmeLauncher
def _set_entry_points(
monkeypatch: pytest.MonkeyPatch,
*eps: _FakeEntryPoint,
) -> None:
"""Patch the registry's internal entrypoint discovery."""
monkeypatch.setattr(
"omnigent.onboarding.sandboxes.registry._entry_points",
lambda: eps,
)
def test_plugin_state_loads_builtins() -> None:
"""The cached plugin state includes core Omnigent built-in providers."""
reset_plugin_state_for_tests()
state = plugin_state()
assert isinstance(state, SandboxProviderPluginState)
assert "modal" in state
assert "kubernetes" in state
def test_plugin_state_is_cached() -> None:
"""Successive calls return the same state object."""
reset_plugin_state_for_tests()
first = plugin_state()
second = plugin_state()
assert first is second
def test_reset_clears_state() -> None:
"""reset_plugin_state_for_tests lets the next call rebuild the state."""
reset_plugin_state_for_tests()
first = plugin_state()
reset_plugin_state_for_tests()
second = plugin_state()
assert first is not second
def test_available_providers_returns_builtins() -> None:
"""available_providers exposes built-in provider names."""
reset_plugin_state_for_tests()
names = available_providers()
assert "modal" in names
assert "kubernetes" in names
def test_get_provider_metadata_known_provider() -> None:
"""Metadata for a registered provider is available."""
reset_plugin_state_for_tests()
meta = get_provider_metadata("modal")
assert meta is not None
assert meta.name == "modal"
assert "omnigent.onboarding.sandboxes.modal:ModalSandboxLauncher" in meta.launcher_class
def test_get_provider_metadata_unknown_returns_none() -> None:
"""Metadata for an unknown provider is ``None``."""
reset_plugin_state_for_tests()
assert get_provider_metadata("not-a-provider") is None
def test_instantiate_loads_built_in_provider() -> None:
"""instantiate imports and constructs a built-in launcher."""
reset_plugin_state_for_tests()
launcher = instantiate("modal")
assert launcher.provider == "modal"
def test_instantiate_unknown_raises() -> None:
"""Instantiating an unknown provider raises SandboxRegistryError."""
reset_plugin_state_for_tests()
with pytest.raises(SandboxRegistryError, match="unknown sandbox provider"):
instantiate("not-a-provider")
def _acme_contribution() -> SandboxProviderContribution:
return SandboxProviderContribution(
name="omnigent-acme",
providers={
"acme": SandboxProviderMetadata(
name="acme",
launcher_class=f"{COMMUNITY_MODULE_PREFIX}acme:AcmeSandboxLauncher",
)
},
)
def test_community_entrypoints_are_loaded(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Community entrypoints in the right namespace are registered."""
reset_plugin_state_for_tests()
_set_entry_points(
monkeypatch,
_FakeEntryPoint(name="acme", value=_acme_contribution),
)
state = plugin_state()
assert "acme" in state
meta = state.get("acme")
assert meta is not None
assert meta.launcher_class == f"{COMMUNITY_MODULE_PREFIX}acme:AcmeSandboxLauncher"
def test_broken_entrypoint_is_recorded_not_raised(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""A broken community plugin is captured in load_errors."""
reset_plugin_state_for_tests()
def _broken() -> SandboxProviderContribution:
raise TypeError("bad")
_set_entry_points(
monkeypatch,
_FakeEntryPoint(name="bad", value=_broken),
)
state = plugin_state()
assert "bad" not in state.providers
assert "bad" in state.load_errors
def test_entrypoint_must_return_contribution(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""An entrypoint returning the wrong type is recorded as an error."""
reset_plugin_state_for_tests()
_set_entry_points(
monkeypatch,
_FakeEntryPoint(name="wrong", value=lambda: "string"),
)
state = plugin_state()
assert "wrong" not in state.providers
assert "wrong" in state.load_errors
assert "expected SandboxProviderContribution" in state.load_errors["wrong"]
def test_validation_rejects_community_module_outside_namespace(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Community providers must live under the community namespace."""
reset_plugin_state_for_tests()
contribution = SandboxProviderContribution(
name="omnigent-external",
providers={
"external": SandboxProviderMetadata(
name="external",
launcher_class="external_provider.sandbox:ExternalLauncher",
)
},
)
_set_entry_points(
monkeypatch,
_FakeEntryPoint(name="external", value=lambda: contribution),
)
state = plugin_state()
assert "external" not in state.providers
assert "external" in state.load_errors
assert COMMUNITY_MODULE_PREFIX in state.load_errors["external"]
def test_validation_rejects_name_collision(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""A community plugin cannot override an existing provider name."""
reset_plugin_state_for_tests()
contribution = SandboxProviderContribution(
name="modal-duplicate",
providers={
"modal": SandboxProviderMetadata(
name="modal",
launcher_class=f"{COMMUNITY_MODULE_PREFIX}modal:ModalSandboxLauncher",
)
},
)
_set_entry_points(
monkeypatch,
_FakeEntryPoint(name="modal-duplicate", value=lambda: contribution),
)
state = plugin_state()
assert "modal" in state.providers # original built-in still present
assert "modal-duplicate" in state.load_errors
assert "override existing provider" in state.load_errors["modal-duplicate"]
def test_validation_rejects_name_mismatch(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""A provider's metadata name must match its provider-key."""
reset_plugin_state_for_tests()
contribution = SandboxProviderContribution(
name="mismatch",
providers={
"foo": SandboxProviderMetadata(
name="bar",
launcher_class=f"{COMMUNITY_MODULE_PREFIX}foo:FooLauncher",
)
},
)
_set_entry_points(
monkeypatch,
_FakeEntryPoint(name="mismatch", value=lambda: contribution),
)
state = plugin_state()
assert "foo" not in state.providers
assert "mismatch" in state.load_errors
assert "does not match metadata name" in state.load_errors["mismatch"]
def test_get_launcher_uses_registry() -> None:
"""The public get_launcher resolves providers through the merged state."""
reset_plugin_state_for_tests()
launcher = get_launcher("modal")
assert launcher.provider == "modal"
def test_get_launcher_unknown_raises_click_exception() -> None:
"""An unknown provider still surfaces as a click.ClickException."""
reset_plugin_state_for_tests()
with contextlib.suppress(DeprecationWarning):
with pytest.raises(
click.ClickException,
match="Unknown or unavailable sandbox provider",
):
get_launcher("definitely-not-real")
def test_instantiate_rejects_non_launcher_class(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""A contribution whose launcher class is not a SandboxLauncher is rejected."""
reset_plugin_state_for_tests()
contribution = SandboxProviderContribution(
name="not-a-launcher",
providers={
"not-a-launcher": SandboxProviderMetadata(
name="not-a-launcher",
launcher_class=f"{COMMUNITY_MODULE_PREFIX}fake:FakeLauncher",
)
},
)
_set_entry_points(
monkeypatch,
_FakeEntryPoint(name="not-a-launcher", value=lambda: contribution),
)
class _NotALauncher:
pass
monkeypatch.setattr(
"omnigent.onboarding.sandboxes.registry._load_object",
lambda _: _NotALauncher,
)
with pytest.raises(SandboxRegistryError, match="not a SandboxLauncher subclass"):
instantiate("not-a-launcher")
+99
View File
@@ -0,0 +1,99 @@
"""Tests for :mod:`omnigent.onboarding.sandboxes.types`."""
from __future__ import annotations
import click
import pytest
from omnigent.onboarding.sandboxes.types import (
HostContext,
SandboxCapabilities,
SandboxCommandError,
SandboxConfigError,
SandboxError,
SandboxInfo,
SandboxSpec,
)
def test_capabilities_defaults() -> None:
"""The default capability set has every feature disabled."""
caps = SandboxCapabilities()
assert caps.cli_bootstrap is False
assert caps.managed_launch is False
assert caps.local_port_forward is False
assert caps.resume_stopped is False
assert caps.programmatic_terminate is False
assert caps.file_copy is False
assert caps.streaming_exec is False
assert caps.foreground_exec is False
def test_capabilities_custom() -> None:
"""Capabilities can be enabled field-by-field."""
caps = SandboxCapabilities(cli_bootstrap=True, foreground_exec=True)
assert caps.cli_bootstrap is True
assert caps.foreground_exec is True
assert caps.managed_launch is False
def test_sandbox_spec_defaults() -> None:
"""SandboxSpec has sensible defaults for optional fields."""
spec = SandboxSpec(name="test")
assert spec.name == "test"
assert spec.image is None
assert spec.cpu is None
assert spec.memory_mib is None
assert spec.disk_gb is None
assert spec.lifetime_s is None
assert spec.tags == {}
def test_sandbox_info_defaults() -> None:
"""SandboxInfo carries an id and optional workspace/metadata."""
info = SandboxInfo(sandbox_id="sb_123")
assert info.sandbox_id == "sb_123"
assert info.workspace_path is None
assert info.metadata == {}
def test_host_context_defaults() -> None:
"""HostContext has defaults for optional repo/config/stage args."""
ctx = HostContext(token="tok", host_id="hid", host_name="hname", server_url="https://srv")
assert ctx.token == "tok"
assert ctx.host_id == "hid"
assert ctx.host_name == "hname"
assert ctx.server_url == "https://srv"
assert ctx.repo_url is None
assert ctx.on_stage is None
assert ctx.host_config == {}
def test_errors_inherit_from_sandbox_error() -> None:
"""All sandbox error types are catchable as SandboxError."""
with pytest.raises(SandboxError):
raise SandboxConfigError("bad config")
def test_sandbox_command_error_carries_fields() -> None:
"""SandboxCommandError exposes command, returncode, and streams."""
exc = SandboxCommandError(
"failed",
command="echo hi",
returncode=1,
stdout="out",
stderr="err",
)
assert exc.command == "echo hi"
assert exc.returncode == 1
assert exc.stdout == "out"
assert exc.stderr == "err"
assert str(exc) == "failed"
def test_sandbox_capability_error_is_click_exception() -> None:
"""SandboxCapabilityError is catchable as click.ClickException (transition)."""
from omnigent.onboarding.sandboxes import SandboxCapabilityError
with pytest.raises(click.ClickException):
raise SandboxCapabilityError("not supported")