From 8ccaea6d2db6aa0d89b0a05b0072e48ba1f1c50a Mon Sep 17 00:00:00 2001 From: Sean Knowles Date: Wed, 22 Jul 2026 19:20:58 +0100 Subject: [PATCH 1/4] Add remote sandbox runtime provider Signed-off-by: Sean Knowles --- omnigent/onboarding/sandboxes/__init__.py | 1 + omnigent/onboarding/sandboxes/base.py | 8 + omnigent/onboarding/sandboxes/daytona.py | 114 +----------- omnigent/onboarding/sandboxes/remote.py | 201 +++++++++++++++++++++ omnigent/server/managed_hosts.py | 56 ++++-- tests/onboarding/sandboxes/test_daytona.py | 63 +------ tests/onboarding/sandboxes/test_remote.py | 105 +++++++++++ tests/server/test_managed_hosts.py | 23 +++ 8 files changed, 383 insertions(+), 188 deletions(-) create mode 100644 omnigent/onboarding/sandboxes/remote.py create mode 100644 tests/onboarding/sandboxes/test_remote.py diff --git a/omnigent/onboarding/sandboxes/__init__.py b/omnigent/onboarding/sandboxes/__init__.py index af5514335d..f1f9f1965c 100644 --- a/omnigent/onboarding/sandboxes/__init__.py +++ b/omnigent/onboarding/sandboxes/__init__.py @@ -58,6 +58,7 @@ "lakebox": "omnigent.onboarding.sandboxes.lakebox:LakeboxLauncher", "modal": "omnigent.onboarding.sandboxes.modal:ModalSandboxLauncher", "daytona": "omnigent.onboarding.sandboxes.daytona:DaytonaSandboxLauncher", + "remote": "omnigent.onboarding.sandboxes.remote:RemoteSandboxLauncher", "boxlite": "omnigent.onboarding.sandboxes.boxlite:BoxliteSandboxLauncher", # CoreWeave Sandbox via the official cwsandbox SDK (the # `omnigent[cwsandbox]` extra), imported lazily like modal/daytona. diff --git a/omnigent/onboarding/sandboxes/base.py b/omnigent/onboarding/sandboxes/base.py index 137174bcda..ced4c18a04 100644 --- a/omnigent/onboarding/sandboxes/base.py +++ b/omnigent/onboarding/sandboxes/base.py @@ -409,6 +409,14 @@ class SandboxLauncher(ABC): # of being silently revived onto an empty workspace. can_resume: ClassVar[bool] = False + def set_launch_context(self, *, owner: str, session_id: str | None) -> None: + """Bind authenticated launch metadata to this launcher instance. + + Providers that delegate provisioning to an external control plane can + override this hook. Direct providers ignore it by default. + """ + del owner, session_id + @abstractmethod def prepare(self) -> None: """ diff --git a/omnigent/onboarding/sandboxes/daytona.py b/omnigent/onboarding/sandboxes/daytona.py index 5fffd750bc..131b185629 100644 --- a/omnigent/onboarding/sandboxes/daytona.py +++ b/omnigent/onboarding/sandboxes/daytona.py @@ -38,14 +38,11 @@ from __future__ import annotations -import json import os import time import uuid from collections.abc import Sequence from typing import TYPE_CHECKING, ClassVar -from urllib.parse import urlencode -from urllib.request import Request, urlopen import click @@ -218,8 +215,6 @@ def __init__(self, *, image: str | None = None, env: Sequence[str] | None = None self._env_names = tuple(env) if env is not None else None self._client: daytona_sdk.Daytona | None = None self._sandboxes: dict[str, DaytonaSandbox] = {} - self._platform_owner: str | None = None - self._platform_session_id: str | None = None def _daytona(self) -> daytona_sdk.Daytona: """ @@ -303,56 +298,6 @@ def _resolve_sandbox_env(self) -> dict[str, str]: resolved[name] = value return resolved - def set_platform_owner(self, owner: str) -> None: - """Bind this per-launch instance to the authenticated session owner.""" - self._platform_owner = owner - - def set_platform_session(self, session_id: str) -> None: - """Bind provider metadata to the Omnigent session being launched.""" - self._platform_session_id = session_id - - def _resolve_owner_profile(self) -> dict[str, object]: - """Resolve non-secret Daytona mount references for this owner.""" - base_url = os.environ.get("PLATFORM_MODEL_CREDENTIAL_BROKER_URL") - token = os.environ.get("PLATFORM_MODEL_CREDENTIAL_BROKER_TOKEN") - if not base_url and not token: - return {} - if not base_url or not token or not self._platform_owner: - raise click.ClickException( - "owner-scoped model credential broker is incompletely configured" - ) - separator = "&" if "?" in base_url else "?" - url = f"{base_url}{separator}{urlencode({'owner': self._platform_owner})}" - request = Request( - url, - headers={ - "Accept": "application/json", - "Authorization": f"Bearer {token}", - }, - ) - try: - with urlopen(request, timeout=10) as response: - profile = json.loads(response.read().decode("utf-8")) - except Exception as exc: - raise click.ClickException( - "could not resolve the execution owner's model connection profile" - ) from exc - if not isinstance(profile, dict): - raise click.ClickException("model connection broker returned an invalid profile") - return profile - - @staticmethod - def _append_runner_passthrough(env_vars: dict[str, str]) -> None: - names = [ - item.strip() - for item in env_vars.get("OMNIGENT_RUNNER_ENV_PASSTHROUGH", "").split(",") - if item.strip() - ] - for name in ("CODEX_HOME", "CLAUDE_CONFIG_DIR"): - if name not in names: - names.append(name) - env_vars["OMNIGENT_RUNNER_ENV_PASSTHROUGH"] = ",".join(names) - def prepare(self) -> None: """ Local preflight: the Daytona SDK must be installed and an API @@ -388,52 +333,17 @@ def provision(self, name: str) -> str: resolved_ref = self._image_ref or os.environ.get(HOST_IMAGE_ENV_VAR) or DEFAULT_HOST_IMAGE env_vars = self._resolve_sandbox_env() - profile = self._resolve_owner_profile() - env_vars.setdefault("CODEX_HOME", "/root/.codex") - env_vars.setdefault("CLAUDE_CONFIG_DIR", "/root/.claude") - self._append_runner_passthrough(env_vars) - volumes = None - raw_volume = profile.get("volume") - if raw_volume is not None: - if not isinstance(raw_volume, dict): - raise click.ClickException("model connection profile volume is invalid") - volume_id = raw_volume.get("id") - subpath = raw_volume.get("subpath") - mount_path = raw_volume.get("mountPath", "/root/.tellimer-auth") - if not all( - isinstance(value, str) and value for value in (volume_id, subpath, mount_path) - ): - raise click.ClickException("model connection profile volume is incomplete") - volumes = [ - daytona.VolumeMount( - volume_id=volume_id, - mount_path=mount_path, - subpath=subpath, - ) - ] - raw_secrets = profile.get("secrets", {}) - if not isinstance(raw_secrets, dict) or not all( - isinstance(key, str) and isinstance(value, str) for key, value in raw_secrets.items() - ): - raise click.ClickException("model connection profile secrets are invalid") - labels = {"omnigent-name": name} - if self._platform_session_id is not None: - # Provider-dashboard audit trail: the canonical session id makes - # the one-session/one-sandbox ownership visible outside Omnigent. - labels["omnigent-session-id"] = self._platform_session_id click.echo(f"▸ Creating Daytona sandbox '{name}' from {resolved_ref}") try: handle = self._daytona().create( daytona.CreateSandboxFromImageParams( image=resolved_ref, env_vars=env_vars or None, - labels=labels, + labels={"omnigent-name": name}, # Release compute after one hour of inactivity. The # persistent sandbox is resumed in place on next message. auto_stop_interval=_MANAGED_AUTO_STOP_MINUTES, resources=daytona.Resources(cpu=_SANDBOX_CPU, memory=_SANDBOX_MEMORY_GIB), - volumes=volumes, - secrets=raw_secrets or None, ), timeout=_CREATE_TIMEOUT_S, # First-use image pulls stream build logs; echo them so a @@ -448,28 +358,6 @@ def provision(self, name: str) -> str: # instead of a generic "internal error". raise click.ClickException(f"Daytona sandbox creation failed: {exc}") from exc self._sandboxes[handle.id] = handle - if raw_volume is not None: - self.run( - handle.id, - "set -eu; mkdir -p /root/.codex; chmod 700 /root/.codex; " - "if [ -s /root/.tellimer-auth/codex/auth.json ]; then " - "cp -f /root/.tellimer-auth/codex/auth.json /root/.codex/auth.json; " - "chmod 600 /root/.codex/auth.json; fi; " - "if [ -f /root/.tellimer-auth/codex/config.toml ]; then " - "cp -f /root/.tellimer-auth/codex/config.toml /root/.codex/config.toml; fi; " - "rm -f /root/.tellimer-auth/codex/state_*.sqlite*", - ) - self.run_background( - handle.id, - "while sleep 15; do " - "if [ -s /root/.codex/auth.json ]; then " - "mkdir -p /root/.tellimer-auth/codex; " - "tmp=/root/.tellimer-auth/codex/.auth.json.$$.tmp; " - "cp -f /root/.codex/auth.json $tmp && chmod 600 $tmp " - "&& mv -f $tmp /root/.tellimer-auth/codex/auth.json; " - "fi; done", - log_path="/tmp/codex-auth-sync.log", - ) click.echo(f" → created {handle.id}") return handle.id diff --git a/omnigent/onboarding/sandboxes/remote.py b/omnigent/onboarding/sandboxes/remote.py new file mode 100644 index 0000000000..7a44e96a0d --- /dev/null +++ b/omnigent/onboarding/sandboxes/remote.py @@ -0,0 +1,201 @@ +"""Managed sandbox launcher backed by an external runtime controller.""" + +from __future__ import annotations + +import json +import os +from collections.abc import Mapping, Sequence +from typing import ClassVar +from urllib.error import HTTPError, URLError +from urllib.request import Request, urlopen + +import click + +from omnigent.onboarding.sandboxes.base import RemoteCommandResult, SandboxLauncher + +DEFAULT_TOKEN_ENV = "OMNIGENT_REMOTE_SANDBOX_TOKEN" + + +class RemoteSandboxLauncher(SandboxLauncher): + """Delegate managed sandbox primitives to a versioned HTTP control plane.""" + + provider: ClassVar[str] = "remote" + supports_cli_bootstrap: ClassVar[bool] = False + can_resume: ClassVar[bool] = True + + def __init__( + self, + *, + url: str, + token_env: str | None = None, + env: Sequence[str] | None = None, + ) -> None: + self._url = url.rstrip("/") + self._token_env = token_env or DEFAULT_TOKEN_ENV + self._env_names = tuple(env or ()) + self._owner: str | None = None + self._session_id: str | None = None + + def set_launch_context(self, *, owner: str, session_id: str | None) -> None: + self._owner = owner + self._session_id = session_id + + def prepare(self) -> None: + if not self._url.startswith(("https://", "http://localhost", "http://127.0.0.1")): + raise click.ClickException( + "remote sandbox controller URL must use HTTPS (or localhost for development)" + ) + if not os.environ.get(self._token_env): + raise click.ClickException( + f"remote sandbox controller token is not set in {self._token_env}" + ) + if self._owner is None or self._session_id is None: + raise click.ClickException("remote sandbox launch context is incomplete") + + def provision(self, name: str) -> str: + self.prepare() + env: dict[str, str] = {} + for env_name in self._env_names: + value = os.environ.get(env_name) + if value is None: + raise click.ClickException( + f"sandbox.remote.env names '{env_name}' but it is not set" + ) + env[env_name] = value + body = self._request( + "POST", + "/api/v1/sandbox-runtimes", + { + "name": name, + "owner": self._owner, + "sessionId": self._session_id, + "env": env, + }, + timeout=15 * 60, + ) + runtime = self._mapping(body.get("runtime"), "runtime") + runtime_id = runtime.get("id") + if not isinstance(runtime_id, str) or not runtime_id: + raise click.ClickException("remote sandbox controller returned no runtime id") + return runtime_id + + def run(self, sandbox_id: str, command: str, *, check: bool = True) -> RemoteCommandResult: + body = self._request( + "POST", + f"/api/v1/sandbox-runtimes/{sandbox_id}/commands", + {"command": command, "timeoutSeconds": 15 * 60}, + timeout=16 * 60, + ) + result = self._mapping(body.get("result"), "command result") + exit_code = result.get("exitCode") + if exit_code is not None and not isinstance(exit_code, int): + raise click.ClickException("remote sandbox controller returned an invalid exit code") + stdout = result.get("stdout", "") + stderr = result.get("stderr", "") + if not isinstance(stdout, str) or not isinstance(stderr, str): + raise click.ClickException("remote sandbox controller returned invalid command output") + completed = RemoteCommandResult(returncode=exit_code or 0, stdout=stdout, stderr=stderr) + if check and completed.returncode != 0: + detail = stderr.strip() or stdout.strip() or "no output" + raise click.ClickException( + f"remote command failed in sandbox '{sandbox_id}' " + f"(exit {completed.returncode}): {detail}" + ) + return completed + + def run_background( + self, + sandbox_id: str, + command: str, + *, + log_path: str = "/tmp/omnigent-host.log", + ) -> RemoteCommandResult: + del log_path + body = self._request( + "POST", + f"/api/v1/sandbox-runtimes/{sandbox_id}/commands", + {"command": command, "detached": True}, + timeout=30, + ) + result = self._mapping(body.get("result"), "command result") + return RemoteCommandResult( + returncode=int(result.get("exitCode") or 0), + stdout=str(result.get("stdout") or "launched\n"), + stderr=str(result.get("stderr") or ""), + ) + + def terminate(self, sandbox_id: str) -> None: + self._request("DELETE", f"/api/v1/sandbox-runtimes/{sandbox_id}") + + def resume(self, sandbox_id: str) -> None: + self._request("POST", f"/api/v1/sandbox-runtimes/{sandbox_id}/resume") + + def is_running(self, sandbox_id: str) -> bool | None: + runtime = self._runtime(sandbox_id) + return None if runtime is None else runtime.get("state") == "running" + + def exists(self, sandbox_id: str) -> bool | None: + runtime = self._runtime(sandbox_id) + return runtime is not None and runtime.get("state") != "deleted" + + def _runtime(self, sandbox_id: str) -> Mapping[str, object] | None: + try: + body = self._request("GET", f"/api/v1/sandbox-runtimes/{sandbox_id}") + except click.ClickException as exc: + if "(404)" in exc.message: + return None + raise + return self._mapping(body.get("runtime"), "runtime") + + def _request( + self, + method: str, + path: str, + body: Mapping[str, object] | None = None, + *, + timeout: int = 90, + ) -> Mapping[str, object]: + token = os.environ.get(self._token_env) + if not token: + raise click.ClickException( + f"remote sandbox controller token is not set in {self._token_env}" + ) + data = json.dumps(body).encode() if body is not None else None + request = Request( + f"{self._url}{path}", + method=method, + data=data, + headers={ + "Accept": "application/json", + "Authorization": f"Bearer {token}", + "Content-Type": "application/json", + "X-Sandbox-Runtime-API-Version": "1", + }, + ) + try: + with urlopen(request, timeout=timeout) as response: + raw = response.read() + except HTTPError as exc: + detail = exc.read().decode("utf-8", errors="replace")[:500] + raise click.ClickException( + f"remote sandbox controller request failed ({exc.code}): {detail}" + ) from exc + except (URLError, TimeoutError, OSError) as exc: + raise click.ClickException( + f"remote sandbox controller is unavailable: {exc}" + ) from exc + if not raw: + return {} + try: + value = json.loads(raw) + except ValueError as exc: + raise click.ClickException( + "remote sandbox controller returned invalid JSON" + ) from exc + return self._mapping(value, "response") + + @staticmethod + def _mapping(value: object, name: str) -> Mapping[str, object]: + if not isinstance(value, dict): + raise click.ClickException(f"remote sandbox controller returned an invalid {name}") + return value diff --git a/omnigent/server/managed_hosts.py b/omnigent/server/managed_hosts.py index 8c8714b760..b07c2eb61a 100644 --- a/omnigent/server/managed_hosts.py +++ b/omnigent/server/managed_hosts.py @@ -160,10 +160,21 @@ "e2b", "openshell", "kubernetes", + "remote", } ) PROVIDERS_WITH_MANAGED_LAUNCH: frozenset[str] = frozenset( - {"modal", "daytona", "boxlite", "cwsandbox", "islo", "e2b", "openshell", "kubernetes"} + { + "modal", + "daytona", + "boxlite", + "cwsandbox", + "islo", + "e2b", + "openshell", + "kubernetes", + "remote", + } ) # How long a managed launch waits for the sandboxed host to register @@ -218,6 +229,10 @@ # mints a fresh token (and the per-Pod token Secret is replaced). KUBERNETES_MANAGED_TOKEN_TTL_S = 7 * 24 * 3600 +# Remote controllers own the backing provider lifecycle. Keep Omnigent's host +# token bounded while allowing a stopped runtime to resume in place. +REMOTE_MANAGED_TOKEN_TTL_S = 7 * 24 * 3600 + # The cwsandbox launch-token TTL is NOT a constant: CW Sandbox's lifetime is # operator-overridable (OMNIGENT_CWSANDBOX_MAX_LIFETIME_S), so the TTL is # derived from the resolved lifetime at parse time via @@ -833,6 +848,13 @@ def parse_sandbox_config(raw: object) -> ManagedSandboxConfig | None: resources=_parse_kubernetes_resources(raw), ) token_ttl_s = KUBERNETES_MANAGED_TOKEN_TTL_S + elif provider == "remote": + launcher_factory = _remote_launcher_factory( + url=_parse_provider_string(raw, "remote", "url"), + token_env=_parse_provider_string(raw, "remote", "token_env"), + env=_parse_provider_env(raw, "remote"), + ) + token_ttl_s = REMOTE_MANAGED_TOKEN_TTL_S else: launcher_factory = _unsupported_launcher_factory(provider) # Never consulted (the factory rejects before any token is @@ -948,6 +970,24 @@ def _build() -> SandboxLauncher: return _build +def _remote_launcher_factory( + *, + url: str | None, + token_env: str | None, + env: list[str] | None, +) -> Callable[[], SandboxLauncher]: + """Build a launcher that delegates lifecycle operations over HTTP.""" + if url is None: + raise ValueError("server config 'sandbox.remote.url' is required") + + def _build() -> SandboxLauncher: + from omnigent.onboarding.sandboxes.remote import RemoteSandboxLauncher + + return RemoteSandboxLauncher(url=url, token_env=token_env, env=env) + + return _build + + def _parse_daytona_image(raw: dict[str, object]) -> str | None: """ Extract and validate the daytona image from the ``sandbox`` dict. @@ -1891,12 +1931,7 @@ async def launch_managed_host( startup, or registration fails. """ launcher = config.launcher_factory() - owner_setter = getattr(launcher, "set_platform_owner", None) - if callable(owner_setter): - owner_setter(owner) - session_setter = getattr(launcher, "set_platform_session", None) - if session_id is not None and callable(session_setter): - session_setter(session_id) + launcher.set_launch_context(owner=owner, session_id=session_id) host_id = uuid.uuid4().hex # Visible label in the host picker; (owner, name) is the hosts # table PK, so embed the host_id's leading hex for uniqueness @@ -1977,12 +2012,7 @@ async def relaunch_managed_host( "was launched with is no longer configured on this server" ), ) - owner_setter = getattr(launcher, "set_platform_owner", None) - if callable(owner_setter): - owner_setter(host.user_id) - session_setter = getattr(launcher, "set_platform_session", None) - if session_id is not None and callable(session_setter): - session_setter(session_id) + launcher.set_launch_context(owner=host.user_id, session_id=session_id) # The old generation is normally already dead (that is why we are # here), but terminate defensively so a transient tunnel outage # can never leave two live sandboxes claiming one host identity. diff --git a/tests/onboarding/sandboxes/test_daytona.py b/tests/onboarding/sandboxes/test_daytona.py index dce1e1c157..aa4fe4f791 100644 --- a/tests/onboarding/sandboxes/test_daytona.py +++ b/tests/onboarding/sandboxes/test_daytona.py @@ -448,11 +448,7 @@ def test_provision_defaults_official_image_and_uses_1h_autostop( [create] = fake_daytona.create_calls assert create.params.image == DEFAULT_HOST_IMAGE assert create.params.auto_stop_interval == 60 - assert create.params.env_vars == { - "CODEX_HOME": "/root/.codex", - "CLAUDE_CONFIG_DIR": "/root/.claude", - "OMNIGENT_RUNNER_ENV_PASSTHROUGH": "CODEX_HOME,CLAUDE_CONFIG_DIR", - } + assert create.params.env_vars is None assert create.params.labels == {"omnigent-name": "managed-abc"} assert create.params.resources == _FakeResources(cpu=2, memory=4) # Cold creates pull + snapshot the image (minutes); the SDK's 60s @@ -461,22 +457,6 @@ def test_provision_defaults_official_image_and_uses_1h_autostop( assert create.has_log_callback is True -def test_provision_labels_sandbox_with_canonical_session_id( - fake_daytona: _FakeDaytonaState, -) -> None: - """Provider dashboard metadata identifies the owning Omnigent session.""" - launcher = DaytonaSandboxLauncher() - launcher.set_platform_session("conv_owner_123") - - launcher.provision("managed-owner") - - [create] = fake_daytona.create_calls - assert create.params.labels == { - "omnigent-name": "managed-owner", - "omnigent-session-id": "conv_owner_123", - } - - def test_provision_image_resolution_order( fake_daytona: _FakeDaytonaState, monkeypatch: pytest.MonkeyPatch ) -> None: @@ -512,9 +492,6 @@ def test_provision_env_passthrough_resolves_from_server_env( assert create.params.env_vars == { "OPENAI_API_KEY": "sk-test-123", "GIT_TOKEN": "ghp-test-456", - "CODEX_HOME": "/root/.codex", - "CLAUDE_CONFIG_DIR": "/root/.claude", - "OMNIGENT_RUNNER_ENV_PASSTHROUGH": "CODEX_HOME,CLAUDE_CONFIG_DIR", } @@ -536,47 +513,9 @@ def test_provision_env_passthrough_env_var_fallback( assert create.params.env_vars == { "OPENAI_API_KEY": "sk-test-123", "GIT_TOKEN": "ghp-test-456", - "CODEX_HOME": "/root/.codex", - "CLAUDE_CONFIG_DIR": "/root/.claude", - "OMNIGENT_RUNNER_ENV_PASSTHROUGH": "CODEX_HOME,CLAUDE_CONFIG_DIR", } -def test_provision_attaches_owner_profile_credentials( - fake_daytona: _FakeDaytonaState, monkeypatch: pytest.MonkeyPatch -) -> None: - """The execution owner's volume and secrets are attached to only their sandbox.""" - launcher = DaytonaSandboxLauncher() - launcher.set_platform_owner("alice@example.com") - monkeypatch.setattr( - launcher, - "_resolve_owner_profile", - lambda: { - "volume": { - "id": "vol-personal", - "subpath": "users/alice", - "mountPath": "/root/.tellimer-auth", - }, - "secrets": {"ANTHROPIC_SETUP_TOKEN": "daytona-secret-ref"}, - }, - ) - - sandbox_id = launcher.provision("managed-alice") - - [create] = fake_daytona.create_calls - assert create.params.volumes == [ - _FakeVolumeMount( - volume_id="vol-personal", - mount_path="/root/.tellimer-auth", - subpath="users/alice", - ) - ] - assert create.params.secrets == {"ANTHROPIC_SETUP_TOKEN": "daytona-secret-ref"} - commands = [call.command for call in fake_daytona.sandboxes[sandbox_id].process.exec_calls] - assert any("/root/.tellimer-auth/codex/auth.json" in command for command in commands) - assert any("codex-auth-sync.log" in command for command in commands) - - def test_provision_env_passthrough_missing_var_fails_loud( fake_daytona: _FakeDaytonaState, monkeypatch: pytest.MonkeyPatch ) -> None: diff --git a/tests/onboarding/sandboxes/test_remote.py b/tests/onboarding/sandboxes/test_remote.py new file mode 100644 index 0000000000..2e2b39183b --- /dev/null +++ b/tests/onboarding/sandboxes/test_remote.py @@ -0,0 +1,105 @@ +"""Tests for the external sandbox runtime controller adapter.""" + +from __future__ import annotations + +import json +from typing import Any +from urllib.request import Request + +import pytest + +from omnigent.onboarding.sandboxes.remote import RemoteSandboxLauncher + + +class _Response: + def __init__(self, body: dict[str, object] | None = None) -> None: + self._body = json.dumps(body).encode() if body is not None else b"" + + def __enter__(self) -> _Response: + return self + + def __exit__(self, *_args: object) -> None: + return None + + def read(self) -> bytes: + return self._body + + +def test_provision_sends_launch_context_and_returns_stable_runtime_id( + monkeypatch: pytest.MonkeyPatch, +) -> None: + requests: list[Request] = [] + + def _urlopen(request: Request, *, timeout: int) -> _Response: + assert timeout == 15 * 60 + requests.append(request) + return _Response({"runtime": {"id": "runtime_abc", "state": "running"}}) + + monkeypatch.setenv("OMNIGENT_REMOTE_SANDBOX_TOKEN", "runtime-secret") + monkeypatch.setenv("PLATFORM_GIT_BROKER_URL", "https://platform.example.com/git") + monkeypatch.setattr("omnigent.onboarding.sandboxes.remote.urlopen", _urlopen) + launcher = RemoteSandboxLauncher( + url="https://platform.example.com", + env=["PLATFORM_GIT_BROKER_URL"], + ) + launcher.set_launch_context(owner="alice@example.com", session_id="conv_alice") + + assert launcher.provision("managed-abcd1234") == "runtime_abc" + payload = json.loads(requests[0].data or b"{}") + assert payload == { + "name": "managed-abcd1234", + "owner": "alice@example.com", + "sessionId": "conv_alice", + "env": {"PLATFORM_GIT_BROKER_URL": "https://platform.example.com/git"}, + } + assert requests[0].headers["Authorization"] == "Bearer runtime-secret" + assert requests[0].headers["X-sandbox-runtime-api-version"] == "1" + + +def test_run_uses_controller_command_endpoint_and_preserves_output( + monkeypatch: pytest.MonkeyPatch, +) -> None: + seen: dict[str, Any] = {} + + def _urlopen(request: Request, *, timeout: int) -> _Response: + del timeout + seen["url"] = request.full_url + seen["payload"] = json.loads(request.data or b"{}") + return _Response({"result": {"exitCode": 0, "stdout": "hello\n", "stderr": ""}}) + + monkeypatch.setenv("OMNIGENT_REMOTE_SANDBOX_TOKEN", "runtime-secret") + monkeypatch.setattr("omnigent.onboarding.sandboxes.remote.urlopen", _urlopen) + launcher = RemoteSandboxLauncher(url="https://platform.example.com") + + result = launcher.run("runtime_abc", "printf hello") + + assert seen["url"].endswith("/api/v1/sandbox-runtimes/runtime_abc/commands") + assert seen["payload"] == {"command": "printf hello", "timeoutSeconds": 15 * 60} + assert result.returncode == 0 + assert result.stdout == "hello\n" + + +def test_stopped_runtime_is_resumed_through_the_controller( + monkeypatch: pytest.MonkeyPatch, +) -> None: + requests: list[tuple[str, str]] = [] + + def _urlopen(request: Request, *, timeout: int) -> _Response: + del timeout + requests.append((request.method, request.full_url)) + if request.method == "GET": + return _Response({"runtime": {"id": "runtime_abc", "state": "stopped"}}) + return _Response({"runtime": {"id": "runtime_abc", "state": "running"}}) + + monkeypatch.setenv("OMNIGENT_REMOTE_SANDBOX_TOKEN", "runtime-secret") + monkeypatch.setattr("omnigent.onboarding.sandboxes.remote.urlopen", _urlopen) + launcher = RemoteSandboxLauncher(url="https://platform.example.com") + + assert launcher.is_running("runtime_abc") is False + assert launcher.exists("runtime_abc") is True + launcher.resume("runtime_abc") + + assert requests[-1] == ( + "POST", + "https://platform.example.com/api/v1/sandbox-runtimes/runtime_abc/resume", + ) diff --git a/tests/server/test_managed_hosts.py b/tests/server/test_managed_hosts.py index 7c80a498d3..34a5f5c15f 100644 --- a/tests/server/test_managed_hosts.py +++ b/tests/server/test_managed_hosts.py @@ -24,6 +24,7 @@ KUBERNETES_MANAGED_TOKEN_TTL_S, MODAL_MANAGED_TOKEN_TTL_S, OPENSHELL_MANAGED_TOKEN_TTL_S, + REMOTE_MANAGED_TOKEN_TTL_S, ManagedSandboxConfig, RepoWorkspace, host_resume_supported, @@ -212,6 +213,28 @@ def test_parse_daytona_without_section_defaults( assert fake.env is None +async def test_parse_remote_controller_config() -> None: + """Remote runtimes are a first-class managed provider with a bounded token.""" + cfg = parse_sandbox_config( + { + "provider": "remote", + "server_url": "https://omnigent.example.com", + "remote": { + "url": "https://platform.example.com/", + "token_env": "PLATFORM_SANDBOX_RUNTIME_TOKEN", + "env": ["PLATFORM_GIT_BROKER_URL"], + }, + } + ) + assert cfg is not None + assert cfg.provider == "remote" + assert cfg.managed_launch_supported is True + assert cfg.token_ttl_s == REMOTE_MANAGED_TOKEN_TTL_S + launcher = cfg.launcher_factory() + assert launcher.provider == "remote" + assert launcher.can_resume is True + + def test_parse_valid_boxlite_cloud_config_builds_parameterized_factory( monkeypatch: pytest.MonkeyPatch, ) -> None: From 82854660ba63e9ed3675d2a4b225b3084c221af0 Mon Sep 17 00:00:00 2001 From: Sean Knowles Date: Wed, 22 Jul 2026 20:48:08 +0100 Subject: [PATCH 2/4] Poll remote sandbox wake state Signed-off-by: Sean Knowles --- omnigent/onboarding/sandboxes/remote.py | 44 ++++++++++++++++++---- tests/onboarding/sandboxes/test_remote.py | 46 +++++++++++++++++++++-- 2 files changed, 80 insertions(+), 10 deletions(-) diff --git a/omnigent/onboarding/sandboxes/remote.py b/omnigent/onboarding/sandboxes/remote.py index 7a44e96a0d..d3cb4289cc 100644 --- a/omnigent/onboarding/sandboxes/remote.py +++ b/omnigent/onboarding/sandboxes/remote.py @@ -4,6 +4,7 @@ import json import os +import time from collections.abc import Mapping, Sequence from typing import ClassVar from urllib.error import HTTPError, URLError @@ -14,6 +15,8 @@ from omnigent.onboarding.sandboxes.base import RemoteCommandResult, SandboxLauncher DEFAULT_TOKEN_ENV = "OMNIGENT_REMOTE_SANDBOX_TOKEN" +_RESUME_TIMEOUT_S = 15 * 60 +_RESUME_POLL_INTERVAL_S = 2 class RemoteSandboxLauncher(SandboxLauncher): @@ -80,6 +83,7 @@ def provision(self, name: str) -> str: return runtime_id def run(self, sandbox_id: str, command: str, *, check: bool = True) -> RemoteCommandResult: + self._ensure_running(sandbox_id) body = self._request( "POST", f"/api/v1/sandbox-runtimes/{sandbox_id}/commands", @@ -111,6 +115,7 @@ def run_background( log_path: str = "/tmp/omnigent-host.log", ) -> RemoteCommandResult: del log_path + self._ensure_running(sandbox_id) body = self._request( "POST", f"/api/v1/sandbox-runtimes/{sandbox_id}/commands", @@ -128,7 +133,29 @@ def terminate(self, sandbox_id: str) -> None: self._request("DELETE", f"/api/v1/sandbox-runtimes/{sandbox_id}") def resume(self, sandbox_id: str) -> None: - self._request("POST", f"/api/v1/sandbox-runtimes/{sandbox_id}/resume") + self._request( + "POST", + f"/api/v1/sandbox-runtimes/{sandbox_id}/resume", + timeout=30, + ) + deadline = time.monotonic() + _RESUME_TIMEOUT_S + while time.monotonic() < deadline: + runtime = self._runtime(sandbox_id) + if runtime is None: + raise click.ClickException( + f"remote sandbox runtime '{sandbox_id}' disappeared while waking" + ) + state = runtime.get("state") + if state == "running": + return + if state in {"deleted", "error"}: + raise click.ClickException( + f"remote sandbox runtime '{sandbox_id}' could not wake (state: {state})" + ) + time.sleep(_RESUME_POLL_INTERVAL_S) + raise click.ClickException( + f"remote sandbox runtime '{sandbox_id}' did not wake within 15 minutes" + ) def is_running(self, sandbox_id: str) -> bool | None: runtime = self._runtime(sandbox_id) @@ -138,6 +165,13 @@ def exists(self, sandbox_id: str) -> bool | None: runtime = self._runtime(sandbox_id) return runtime is not None and runtime.get("state") != "deleted" + def _ensure_running(self, sandbox_id: str) -> None: + runtime = self._runtime(sandbox_id) + if runtime is None: + raise click.ClickException(f"remote sandbox runtime '{sandbox_id}' was not found") + if runtime.get("state") != "running": + self.resume(sandbox_id) + def _runtime(self, sandbox_id: str) -> Mapping[str, object] | None: try: body = self._request("GET", f"/api/v1/sandbox-runtimes/{sandbox_id}") @@ -181,17 +215,13 @@ def _request( f"remote sandbox controller request failed ({exc.code}): {detail}" ) from exc except (URLError, TimeoutError, OSError) as exc: - raise click.ClickException( - f"remote sandbox controller is unavailable: {exc}" - ) from exc + raise click.ClickException(f"remote sandbox controller is unavailable: {exc}") from exc if not raw: return {} try: value = json.loads(raw) except ValueError as exc: - raise click.ClickException( - "remote sandbox controller returned invalid JSON" - ) from exc + raise click.ClickException("remote sandbox controller returned invalid JSON") from exc return self._mapping(value, "response") @staticmethod diff --git a/tests/onboarding/sandboxes/test_remote.py b/tests/onboarding/sandboxes/test_remote.py index 2e2b39183b..8fa677e86c 100644 --- a/tests/onboarding/sandboxes/test_remote.py +++ b/tests/onboarding/sandboxes/test_remote.py @@ -65,6 +65,8 @@ def _urlopen(request: Request, *, timeout: int) -> _Response: del timeout seen["url"] = request.full_url seen["payload"] = json.loads(request.data or b"{}") + if request.method == "GET": + return _Response({"runtime": {"id": "runtime_abc", "state": "running"}}) return _Response({"result": {"exitCode": 0, "stdout": "hello\n", "stderr": ""}}) monkeypatch.setenv("OMNIGENT_REMOTE_SANDBOX_TOKEN", "runtime-secret") @@ -83,23 +85,61 @@ def test_stopped_runtime_is_resumed_through_the_controller( monkeypatch: pytest.MonkeyPatch, ) -> None: requests: list[tuple[str, str]] = [] + get_states = iter(["stopped", "stopped", "running"]) def _urlopen(request: Request, *, timeout: int) -> _Response: del timeout requests.append((request.method, request.full_url)) if request.method == "GET": - return _Response({"runtime": {"id": "runtime_abc", "state": "stopped"}}) - return _Response({"runtime": {"id": "runtime_abc", "state": "running"}}) + return _Response({"runtime": {"id": "runtime_abc", "state": next(get_states)}}) + return _Response({"runtime": {"id": "runtime_abc", "state": "provisioning"}}) monkeypatch.setenv("OMNIGENT_REMOTE_SANDBOX_TOKEN", "runtime-secret") monkeypatch.setattr("omnigent.onboarding.sandboxes.remote.urlopen", _urlopen) + monkeypatch.setattr("omnigent.onboarding.sandboxes.remote.time.sleep", lambda _seconds: None) launcher = RemoteSandboxLauncher(url="https://platform.example.com") assert launcher.is_running("runtime_abc") is False assert launcher.exists("runtime_abc") is True launcher.resume("runtime_abc") - assert requests[-1] == ( + assert ( "POST", "https://platform.example.com/api/v1/sandbox-runtimes/runtime_abc/resume", + ) in requests + assert requests[-1] == ( + "GET", + "https://platform.example.com/api/v1/sandbox-runtimes/runtime_abc", ) + + +def test_first_command_polls_a_stopped_runtime_before_execution( + monkeypatch: pytest.MonkeyPatch, +) -> None: + requests: list[tuple[str, str]] = [] + states = iter(["stopped", "provisioning", "running"]) + + def _urlopen(request: Request, *, timeout: int) -> _Response: + del timeout + requests.append((request.method, request.full_url)) + if request.method == "GET": + return _Response({"runtime": {"id": "runtime_abc", "state": next(states)}}) + if request.full_url.endswith("/resume"): + return _Response({"runtime": {"id": "runtime_abc", "state": "provisioning"}}) + return _Response({"result": {"exitCode": 0, "stdout": "awake\n", "stderr": ""}}) + + monkeypatch.setenv("OMNIGENT_REMOTE_SANDBOX_TOKEN", "runtime-secret") + monkeypatch.setattr("omnigent.onboarding.sandboxes.remote.urlopen", _urlopen) + monkeypatch.setattr("omnigent.onboarding.sandboxes.remote.time.sleep", lambda _seconds: None) + launcher = RemoteSandboxLauncher(url="https://platform.example.com") + + result = launcher.run("runtime_abc", "printf awake") + + assert result.stdout == "awake\n" + assert requests == [ + ("GET", "https://platform.example.com/api/v1/sandbox-runtimes/runtime_abc"), + ("POST", "https://platform.example.com/api/v1/sandbox-runtimes/runtime_abc/resume"), + ("GET", "https://platform.example.com/api/v1/sandbox-runtimes/runtime_abc"), + ("GET", "https://platform.example.com/api/v1/sandbox-runtimes/runtime_abc"), + ("POST", "https://platform.example.com/api/v1/sandbox-runtimes/runtime_abc/commands"), + ] From e66de00adf96c140e7079664408d5f3c5ddabbf9 Mon Sep 17 00:00:00 2001 From: Sean Knowles Date: Wed, 22 Jul 2026 21:33:51 +0100 Subject: [PATCH 3/4] Launch remote sandbox hosts durably Signed-off-by: Sean Knowles --- omnigent/onboarding/sandboxes/remote.py | 22 ---------------- tests/onboarding/sandboxes/test_remote.py | 32 +++++++++++++++++++++++ 2 files changed, 32 insertions(+), 22 deletions(-) diff --git a/omnigent/onboarding/sandboxes/remote.py b/omnigent/onboarding/sandboxes/remote.py index d3cb4289cc..238260b9e4 100644 --- a/omnigent/onboarding/sandboxes/remote.py +++ b/omnigent/onboarding/sandboxes/remote.py @@ -107,28 +107,6 @@ def run(self, sandbox_id: str, command: str, *, check: bool = True) -> RemoteCom ) return completed - def run_background( - self, - sandbox_id: str, - command: str, - *, - log_path: str = "/tmp/omnigent-host.log", - ) -> RemoteCommandResult: - del log_path - self._ensure_running(sandbox_id) - body = self._request( - "POST", - f"/api/v1/sandbox-runtimes/{sandbox_id}/commands", - {"command": command, "detached": True}, - timeout=30, - ) - result = self._mapping(body.get("result"), "command result") - return RemoteCommandResult( - returncode=int(result.get("exitCode") or 0), - stdout=str(result.get("stdout") or "launched\n"), - stderr=str(result.get("stderr") or ""), - ) - def terminate(self, sandbox_id: str) -> None: self._request("DELETE", f"/api/v1/sandbox-runtimes/{sandbox_id}") diff --git a/tests/onboarding/sandboxes/test_remote.py b/tests/onboarding/sandboxes/test_remote.py index 8fa677e86c..d4ee6aeab6 100644 --- a/tests/onboarding/sandboxes/test_remote.py +++ b/tests/onboarding/sandboxes/test_remote.py @@ -81,6 +81,38 @@ def _urlopen(request: Request, *, timeout: int) -> _Response: assert result.stdout == "hello\n" +def test_background_run_uses_the_shared_durable_shell_wrapper( + monkeypatch: pytest.MonkeyPatch, +) -> None: + payloads: list[dict[str, object]] = [] + + def _urlopen(request: Request, *, timeout: int) -> _Response: + del timeout + if request.method == "GET": + return _Response({"runtime": {"id": "runtime_abc", "state": "running"}}) + payloads.append(json.loads(request.data or b"{}")) + return _Response({"result": {"exitCode": 0, "stdout": "launched\n", "stderr": ""}}) + + monkeypatch.setenv("OMNIGENT_REMOTE_SANDBOX_TOKEN", "runtime-secret") + monkeypatch.setattr("omnigent.onboarding.sandboxes.remote.urlopen", _urlopen) + launcher = RemoteSandboxLauncher(url="https://platform.example.com") + + result = launcher.run_background( + "runtime_abc", + "FOO=bar omnigent host --server https://omnigent.example.com", + ) + + assert result.stdout == "launched\n" + assert payloads == [ + { + "command": "setsid nohup sh -c " + "'FOO=bar omnigent host --server https://omnigent.example.com' " + "> /tmp/omnigent-host.log 2>&1 < /dev/null & echo launched", + "timeoutSeconds": 15 * 60, + } + ] + + def test_stopped_runtime_is_resumed_through_the_controller( monkeypatch: pytest.MonkeyPatch, ) -> None: From 7ac46f07c630ec7c1e3de25ab0d10638812c0a8b Mon Sep 17 00:00:00 2001 From: Sean Knowles Date: Wed, 22 Jul 2026 23:51:27 +0100 Subject: [PATCH 4/4] Harden remote sandbox lifecycle activity --- omnigent/onboarding/sandboxes/base.py | 13 +++++ omnigent/onboarding/sandboxes/remote.py | 66 +++++++++++++++++---- omnigent/server/managed_hosts.py | 36 ++++++++++++ omnigent/server/routes/sessions.py | 13 +++++ tests/onboarding/sandboxes/test_remote.py | 71 ++++++++++++++++++++++- tests/server/test_managed_hosts.py | 33 +++++++++++ 6 files changed, 218 insertions(+), 14 deletions(-) diff --git a/omnigent/onboarding/sandboxes/base.py b/omnigent/onboarding/sandboxes/base.py index ced4c18a04..29e2f3f659 100644 --- a/omnigent/onboarding/sandboxes/base.py +++ b/omnigent/onboarding/sandboxes/base.py @@ -665,6 +665,19 @@ def keep_alive(self, sandbox_id: str) -> None: """ raise self._capability_error("configure keep-alive") + def set_activity(self, sandbox_id: str, *, active: bool) -> None: + """Tell a lifecycle-aware provider whether the sandbox has active work. + + The default is deliberately a no-op: most providers either infer + activity themselves or do not expose a policy hook. Managed remote + controllers can override this to protect a long-running turn from an + idle timer and restore the normal idle policy when the turn finishes. + + :param sandbox_id: Target sandbox. + :param active: ``True`` while a turn or background task is running. + """ + del sandbox_id, active + @abstractmethod def run(self, sandbox_id: str, command: str, *, check: bool = True) -> RemoteCommandResult: """ diff --git a/omnigent/onboarding/sandboxes/remote.py b/omnigent/onboarding/sandboxes/remote.py index 238260b9e4..e183b3ac97 100644 --- a/omnigent/onboarding/sandboxes/remote.py +++ b/omnigent/onboarding/sandboxes/remote.py @@ -17,6 +17,10 @@ DEFAULT_TOKEN_ENV = "OMNIGENT_REMOTE_SANDBOX_TOKEN" _RESUME_TIMEOUT_S = 15 * 60 _RESUME_POLL_INTERVAL_S = 2 +_MAX_RESPONSE_BYTES = 1024 * 1024 +_MAX_ERROR_BYTES = 4096 +_RETRYABLE_STATUS_CODES = frozenset({429, 502, 503, 504}) +_RETRY_DELAYS_S = (0.25, 1.0) class RemoteSandboxLauncher(SandboxLauncher): @@ -75,6 +79,7 @@ def provision(self, name: str) -> str: "env": env, }, timeout=15 * 60, + retryable=True, ) runtime = self._mapping(body.get("runtime"), "runtime") runtime_id = runtime.get("id") @@ -108,13 +113,23 @@ def run(self, sandbox_id: str, command: str, *, check: bool = True) -> RemoteCom return completed def terminate(self, sandbox_id: str) -> None: - self._request("DELETE", f"/api/v1/sandbox-runtimes/{sandbox_id}") + self._request("DELETE", f"/api/v1/sandbox-runtimes/{sandbox_id}", retryable=True) + + def set_activity(self, sandbox_id: str, *, active: bool) -> None: + self._request( + "POST", + f"/api/v1/sandbox-runtimes/{sandbox_id}/activity", + {"active": active}, + timeout=30, + retryable=True, + ) def resume(self, sandbox_id: str) -> None: self._request( "POST", f"/api/v1/sandbox-runtimes/{sandbox_id}/resume", timeout=30, + retryable=True, ) deadline = time.monotonic() + _RESUME_TIMEOUT_S while time.monotonic() < deadline: @@ -152,7 +167,7 @@ def _ensure_running(self, sandbox_id: str) -> None: def _runtime(self, sandbox_id: str) -> Mapping[str, object] | None: try: - body = self._request("GET", f"/api/v1/sandbox-runtimes/{sandbox_id}") + body = self._request("GET", f"/api/v1/sandbox-runtimes/{sandbox_id}", retryable=True) except click.ClickException as exc: if "(404)" in exc.message: return None @@ -166,6 +181,7 @@ def _request( body: Mapping[str, object] | None = None, *, timeout: int = 90, + retryable: bool = False, ) -> Mapping[str, object]: token = os.environ.get(self._token_env) if not token: @@ -184,16 +200,30 @@ def _request( "X-Sandbox-Runtime-API-Version": "1", }, ) - try: - with urlopen(request, timeout=timeout) as response: - raw = response.read() - except HTTPError as exc: - detail = exc.read().decode("utf-8", errors="replace")[:500] - raise click.ClickException( - f"remote sandbox controller request failed ({exc.code}): {detail}" - ) from exc - except (URLError, TimeoutError, OSError) as exc: - raise click.ClickException(f"remote sandbox controller is unavailable: {exc}") from exc + attempts = len(_RETRY_DELAYS_S) + 1 if retryable else 1 + raw = b"" + for attempt in range(attempts): + try: + with urlopen(request, timeout=timeout) as response: + raw = self._read_bounded(response, _MAX_RESPONSE_BYTES) + break + except HTTPError as exc: + if exc.code in _RETRYABLE_STATUS_CODES and attempt + 1 < attempts: + time.sleep(_RETRY_DELAYS_S[attempt]) + continue + detail = self._read_bounded(exc, _MAX_ERROR_BYTES).decode( + "utf-8", errors="replace" + ) + raise click.ClickException( + f"remote sandbox controller request failed ({exc.code}): {detail}" + ) from exc + except (URLError, TimeoutError, OSError) as exc: + if attempt + 1 < attempts: + time.sleep(_RETRY_DELAYS_S[attempt]) + continue + raise click.ClickException( + f"remote sandbox controller is unavailable: {exc}" + ) from exc if not raw: return {} try: @@ -202,6 +232,18 @@ def _request( raise click.ClickException("remote sandbox controller returned invalid JSON") from exc return self._mapping(value, "response") + @staticmethod + def _read_bounded(response: object, limit: int) -> bytes: + read = getattr(response, "read", None) + if not callable(read): + raise click.ClickException("remote sandbox controller returned an invalid response") + raw = read(limit + 1) + if len(raw) > limit: + raise click.ClickException( + f"remote sandbox controller response exceeded {limit} bytes" + ) + return raw + @staticmethod def _mapping(value: object, name: str) -> Mapping[str, object]: if not isinstance(value, dict): diff --git a/omnigent/server/managed_hosts.py b/omnigent/server/managed_hosts.py index b07c2eb61a..26d2edb705 100644 --- a/omnigent/server/managed_hosts.py +++ b/omnigent/server/managed_hosts.py @@ -2253,6 +2253,42 @@ def host_sandbox_exists( return launcher.exists(host.sandbox_id) +async def set_managed_host_activity( + host_id: str, + host_store: HostStore, + config: ManagedSandboxConfig | None, + *, + active: bool, +) -> None: + """Propagate turn activity to a lifecycle-aware managed sandbox. + + This is best-effort by design: publishing a model status edge must never + fail because an optional provider policy endpoint is unavailable. The + platform-owned remote launcher uses the signal to extend the provider + safety timer while work is active and restore the normal idle timer once + the session is quiescent. + """ + host = await asyncio.to_thread(host_store.get_host, host_id) + if host is None or host.sandbox_id is None: + return + launcher = _launcher_for_teardown(host, config) + if launcher is None: + return + try: + await asyncio.to_thread( + launcher.set_activity, + host.sandbox_id, + active=active, + ) + except Exception: # noqa: BLE001 -- optional provider boundary must soft-fail + _logger.warning( + "Could not update activity policy for managed host %s (sandbox %s)", + host_id, + host.sandbox_id, + exc_info=True, + ) + + # ── Managed-host wake (resume a dormant host on demand) ───────────────────── # Per-host resume single-flight: one in-flight resume per host_id on this diff --git a/omnigent/server/routes/sessions.py b/omnigent/server/routes/sessions.py index 8d269213b5..4b639e681a 100644 --- a/omnigent/server/routes/sessions.py +++ b/omnigent/server/routes/sessions.py @@ -181,6 +181,7 @@ RepoWorkspace, host_resume_supported, host_sandbox_is_running, + set_managed_host_activity, ) from omnigent.server.mcp_pool import ServerMcpPool from omnigent.server.permissions import check_session_access @@ -20839,6 +20840,18 @@ async def post_event( response_id=response_id, background_task_count=bg_count, ) + host_store = getattr(request.app.state, "host_store", None) + sandbox_config = getattr(request.app.state, "sandbox_config", None) + if conv.host_id is not None and host_store is not None: + await set_managed_host_activity( + conv.host_id, + host_store, + sandbox_config, + active=( + status in {"running", "waiting"} + or _session_background_task_count_cache.get(session_id, 0) > 0 + ), + ) forward_body = body.model_dump() forward_body["data"] = await _enrich_idle_status_with_subagent_output( forward_body["data"], status, session_id, conversation_store diff --git a/tests/onboarding/sandboxes/test_remote.py b/tests/onboarding/sandboxes/test_remote.py index d4ee6aeab6..772e7a4fc6 100644 --- a/tests/onboarding/sandboxes/test_remote.py +++ b/tests/onboarding/sandboxes/test_remote.py @@ -4,8 +4,10 @@ import json from typing import Any +from urllib.error import URLError from urllib.request import Request +import click import pytest from omnigent.onboarding.sandboxes.remote import RemoteSandboxLauncher @@ -21,8 +23,8 @@ def __enter__(self) -> _Response: def __exit__(self, *_args: object) -> None: return None - def read(self) -> bytes: - return self._body + def read(self, amount: int = -1) -> bytes: + return self._body if amount < 0 else self._body[:amount] def test_provision_sends_launch_context_and_returns_stable_runtime_id( @@ -175,3 +177,68 @@ def _urlopen(request: Request, *, timeout: int) -> _Response: ("GET", "https://platform.example.com/api/v1/sandbox-runtimes/runtime_abc"), ("POST", "https://platform.example.com/api/v1/sandbox-runtimes/runtime_abc/commands"), ] + + +def test_activity_signal_uses_versioned_controller_endpoint( + monkeypatch: pytest.MonkeyPatch, +) -> None: + requests: list[tuple[str, str, dict[str, object]]] = [] + + def _urlopen(request: Request, *, timeout: int) -> _Response: + assert timeout == 30 + requests.append((request.method, request.full_url, json.loads(request.data or b"{}"))) + return _Response({"runtime": {"id": "runtime_abc", "active": True}}) + + monkeypatch.setenv("OMNIGENT_REMOTE_SANDBOX_TOKEN", "runtime-secret") + monkeypatch.setattr("omnigent.onboarding.sandboxes.remote.urlopen", _urlopen) + launcher = RemoteSandboxLauncher(url="https://platform.example.com") + + launcher.set_activity("runtime_abc", active=True) + + assert requests == [ + ( + "POST", + "https://platform.example.com/api/v1/sandbox-runtimes/runtime_abc/activity", + {"active": True}, + ) + ] + + +def test_retryable_status_lookup_recovers_from_transient_transport_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + attempts = 0 + + def _urlopen(request: Request, *, timeout: int) -> _Response: + del request, timeout + nonlocal attempts + attempts += 1 + if attempts == 1: + raise URLError("temporary outage") + return _Response({"runtime": {"id": "runtime_abc", "state": "running"}}) + + monkeypatch.setenv("OMNIGENT_REMOTE_SANDBOX_TOKEN", "runtime-secret") + monkeypatch.setattr("omnigent.onboarding.sandboxes.remote.urlopen", _urlopen) + monkeypatch.setattr("omnigent.onboarding.sandboxes.remote.time.sleep", lambda _seconds: None) + launcher = RemoteSandboxLauncher(url="https://platform.example.com") + + assert launcher.is_running("runtime_abc") is True + assert attempts == 2 + + +def test_response_body_is_bounded_before_json_parsing( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class _OversizedResponse(_Response): + def __init__(self) -> None: + self._body = b"x" * (1024 * 1024 + 1) + + monkeypatch.setenv("OMNIGENT_REMOTE_SANDBOX_TOKEN", "runtime-secret") + monkeypatch.setattr( + "omnigent.onboarding.sandboxes.remote.urlopen", + lambda _request, *, timeout: _OversizedResponse(), + ) + launcher = RemoteSandboxLauncher(url="https://platform.example.com") + + with pytest.raises(click.ClickException, match="response exceeded 1048576 bytes"): + launcher.is_running("runtime_abc") diff --git a/tests/server/test_managed_hosts.py b/tests/server/test_managed_hosts.py index 34a5f5c15f..786cbbd68f 100644 --- a/tests/server/test_managed_hosts.py +++ b/tests/server/test_managed_hosts.py @@ -33,6 +33,7 @@ parse_sandbox_config, relaunch_managed_host, resume_managed_host, + set_managed_host_activity, terminate_managed_host, ) from omnigent.stores.agent_store.sqlalchemy_store import SqlAlchemyAgentStore @@ -1828,6 +1829,38 @@ class _IsloFakeLauncher(FakeSandboxLauncher): provider: ClassVar[str] = "islo" +class _ActivityFakeLauncher(_IsloFakeLauncher): + """Capture lifecycle activity signals from the managed-host seam.""" + + def __init__(self) -> None: + super().__init__() + self.activity: list[tuple[str, bool]] = [] + + def set_activity(self, sandbox_id: str, *, active: bool) -> None: + self.activity.append((sandbox_id, active)) + + +async def test_managed_host_activity_targets_its_bound_sandbox(db_uri: str) -> None: + """A status edge updates only the sandbox recorded on that host row.""" + host_store = HostStore(db_uri) + host_id = "618f50482ee943a99eae16fcf1cc9158" + host_store.register_managed_host( + host_id=host_id, + name="managed-activity", + user_id=_OWNER, + token="tok-activity", + provider="islo", + sandbox_id="sb-activity", + token_expires_at=now_epoch() + 3600, + ) + fake = _ActivityFakeLauncher() + + await set_managed_host_activity(host_id, host_store, _injected_config(fake), active=True) + await set_managed_host_activity(host_id, host_store, _injected_config(fake), active=False) + + assert fake.activity == [("sb-activity", True), ("sb-activity", False)] + + async def test_host_resume_supported_requires_resumable_matching_launcher(db_uri: str) -> None: """The wake gate requires matching provider, sandbox id, and ``can_resume``.""" host_store = HostStore(db_uri)