feat: add managed tool gateway and Nous subscription support

- add managed modal and gateway-backed tool integrations\n- improve CLI setup, auth, and configuration for subscriber flows\n- expand tests and docs for managed tool support
This commit is contained in:
Robin Fernandes
2026-03-26 15:27:27 -07:00
parent cbf195e806
commit 95dc9aaa75
44 changed files with 4567 additions and 423 deletions

View File

@@ -0,0 +1,282 @@
"""Managed Modal environment backed by tool-gateway."""
from __future__ import annotations
import json
import logging
import os
import requests
import time
import uuid
from typing import Any, Dict, Optional
from tools.environments.base import BaseEnvironment
from tools.interrupt import is_interrupted
from tools.managed_tool_gateway import resolve_managed_tool_gateway
logger = logging.getLogger(__name__)
def _request_timeout_env(name: str, default: float) -> float:
try:
value = float(os.getenv(name, str(default)))
return value if value > 0 else default
except (TypeError, ValueError):
return default
class ManagedModalEnvironment(BaseEnvironment):
"""Gateway-owned Modal sandbox with Hermes-compatible execute/cleanup."""
_CONNECT_TIMEOUT_SECONDS = _request_timeout_env("TERMINAL_MANAGED_MODAL_CONNECT_TIMEOUT_SECONDS", 1.0)
_POLL_READ_TIMEOUT_SECONDS = _request_timeout_env("TERMINAL_MANAGED_MODAL_POLL_READ_TIMEOUT_SECONDS", 5.0)
_CANCEL_READ_TIMEOUT_SECONDS = _request_timeout_env("TERMINAL_MANAGED_MODAL_CANCEL_READ_TIMEOUT_SECONDS", 5.0)
def __init__(
self,
image: str,
cwd: str = "/root",
timeout: int = 60,
modal_sandbox_kwargs: Optional[Dict[str, Any]] = None,
persistent_filesystem: bool = True,
task_id: str = "default",
):
super().__init__(cwd=cwd, timeout=timeout)
gateway = resolve_managed_tool_gateway("modal")
if gateway is None:
raise ValueError("Managed Modal requires a configured tool gateway and Nous user token")
self._gateway_origin = gateway.gateway_origin.rstrip("/")
self._nous_user_token = gateway.nous_user_token
self._task_id = task_id
self._persistent = persistent_filesystem
self._image = image
self._sandbox_kwargs = dict(modal_sandbox_kwargs or {})
self._create_idempotency_key = str(uuid.uuid4())
self._sandbox_id = self._create_sandbox()
def execute(self, command: str, cwd: str = "", *,
timeout: int | None = None,
stdin_data: str | None = None) -> dict:
exec_command, sudo_stdin = self._prepare_command(command)
# When a sudo password is present, inject it via a shell-level pipe
# (same approach as the direct ModalEnvironment) since the gateway
# cannot pipe subprocess stdin directly.
if sudo_stdin is not None:
import shlex
exec_command = (
f"printf '%s\\n' {shlex.quote(sudo_stdin.rstrip())} | {exec_command}"
)
exec_cwd = cwd or self.cwd
effective_timeout = timeout or self.timeout
exec_id = str(uuid.uuid4())
payload: Dict[str, Any] = {
"execId": exec_id,
"command": exec_command,
"cwd": exec_cwd,
"timeoutMs": int(effective_timeout * 1000),
}
if stdin_data is not None:
payload["stdinData"] = stdin_data
try:
response = self._request(
"POST",
f"/v1/sandboxes/{self._sandbox_id}/execs",
json=payload,
timeout=10,
)
except Exception as exc:
return {
"output": f"Managed Modal exec failed: {exc}",
"returncode": 1,
}
if response.status_code >= 400:
return {
"output": self._format_error("Managed Modal exec failed", response),
"returncode": 1,
}
body = response.json()
status = body.get("status")
if status in {"completed", "failed", "cancelled", "timeout"}:
return {
"output": body.get("output", ""),
"returncode": body.get("returncode", 1),
}
if body.get("execId") != exec_id:
return {
"output": "Managed Modal exec start did not return the expected exec id",
"returncode": 1,
}
poll_interval = 0.25
deadline = time.monotonic() + effective_timeout + 10
while time.monotonic() < deadline:
if is_interrupted():
self._cancel_exec(exec_id)
return {
"output": "[Command interrupted - Modal sandbox exec cancelled]",
"returncode": 130,
}
try:
status_response = self._request(
"GET",
f"/v1/sandboxes/{self._sandbox_id}/execs/{exec_id}",
timeout=(self._CONNECT_TIMEOUT_SECONDS, self._POLL_READ_TIMEOUT_SECONDS),
)
except Exception as exc:
return {
"output": f"Managed Modal exec poll failed: {exc}",
"returncode": 1,
}
if status_response.status_code == 404:
return {
"output": "Managed Modal exec not found",
"returncode": 1,
}
if status_response.status_code >= 400:
return {
"output": self._format_error("Managed Modal exec poll failed", status_response),
"returncode": 1,
}
status_body = status_response.json()
status = status_body.get("status")
if status in {"completed", "failed", "cancelled", "timeout"}:
return {
"output": status_body.get("output", ""),
"returncode": status_body.get("returncode", 1),
}
time.sleep(poll_interval)
self._cancel_exec(exec_id)
return {
"output": f"Managed Modal exec timed out after {effective_timeout}s",
"returncode": 124,
}
def cleanup(self):
if not getattr(self, "_sandbox_id", None):
return
try:
self._request(
"POST",
f"/v1/sandboxes/{self._sandbox_id}/terminate",
json={
"snapshotBeforeTerminate": self._persistent,
},
timeout=60,
)
except Exception as exc:
logger.warning("Managed Modal cleanup failed: %s", exc)
finally:
self._sandbox_id = None
def _create_sandbox(self) -> str:
cpu = self._coerce_number(self._sandbox_kwargs.get("cpu"), 1)
memory = self._coerce_number(
self._sandbox_kwargs.get("memoryMiB", self._sandbox_kwargs.get("memory")),
5120,
)
disk = self._coerce_number(
self._sandbox_kwargs.get("ephemeral_disk", self._sandbox_kwargs.get("diskMiB")),
None,
)
create_payload = {
"image": self._image,
"cwd": self.cwd,
"cpu": cpu,
"memoryMiB": memory,
"timeoutMs": 3_600_000,
"idleTimeoutMs": max(300_000, int(self.timeout * 1000)),
"persistentFilesystem": self._persistent,
"logicalKey": self._task_id,
}
if disk is not None:
create_payload["diskMiB"] = disk
response = self._request(
"POST",
"/v1/sandboxes",
json=create_payload,
timeout=60,
extra_headers={
"x-idempotency-key": self._create_idempotency_key,
},
)
if response.status_code >= 400:
raise RuntimeError(self._format_error("Managed Modal create failed", response))
body = response.json()
sandbox_id = body.get("id")
if not isinstance(sandbox_id, str) or not sandbox_id:
raise RuntimeError("Managed Modal create did not return a sandbox id")
return sandbox_id
def _request(self, method: str, path: str, *,
json: Dict[str, Any] | None = None,
timeout: int = 30,
extra_headers: Dict[str, str] | None = None) -> requests.Response:
headers = {
"Authorization": f"Bearer {self._nous_user_token}",
"Content-Type": "application/json",
}
if extra_headers:
headers.update(extra_headers)
return requests.request(
method,
f"{self._gateway_origin}{path}",
headers=headers,
json=json,
timeout=timeout,
)
def _cancel_exec(self, exec_id: str) -> None:
try:
self._request(
"POST",
f"/v1/sandboxes/{self._sandbox_id}/execs/{exec_id}/cancel",
timeout=(self._CONNECT_TIMEOUT_SECONDS, self._CANCEL_READ_TIMEOUT_SECONDS),
)
except Exception as exc:
logger.warning("Managed Modal exec cancel failed: %s", exc)
@staticmethod
def _coerce_number(value: Any, default: float) -> float:
try:
if value is None:
return default
return float(value)
except (TypeError, ValueError):
return default
@staticmethod
def _format_error(prefix: str, response: requests.Response) -> str:
try:
payload = response.json()
if isinstance(payload, dict):
message = payload.get("error") or payload.get("message") or payload.get("code")
if isinstance(message, str) and message:
return f"{prefix}: {message}"
return f"{prefix}: {json.dumps(payload, ensure_ascii=False)}"
except Exception:
pass
text = response.text.strip()
if text:
return f"{prefix}: {text}"
return f"{prefix}: HTTP {response.status_code}"

View File

@@ -20,6 +20,7 @@ from tools.interrupt import is_interrupted
logger = logging.getLogger(__name__)
_SNAPSHOT_STORE = get_hermes_home() / "modal_snapshots.json"
_DIRECT_SNAPSHOT_NAMESPACE = "direct"
def _load_snapshots() -> Dict[str, str]:
@@ -38,12 +39,72 @@ def _save_snapshots(data: Dict[str, str]) -> None:
_SNAPSHOT_STORE.write_text(json.dumps(data, indent=2))
class _AsyncWorker:
"""Background thread with its own event loop for async-safe swe-rex calls.
def _direct_snapshot_key(task_id: str) -> str:
return f"{_DIRECT_SNAPSHOT_NAMESPACE}:{task_id}"
Allows sync code to submit async coroutines and block for results,
even when called from inside another running event loop (e.g. Atropos).
"""
def _get_snapshot_restore_candidate(task_id: str) -> tuple[str | None, bool]:
"""Return a snapshot id for direct Modal restore and whether the key is legacy."""
snapshots = _load_snapshots()
namespaced_key = _direct_snapshot_key(task_id)
snapshot_id = snapshots.get(namespaced_key)
if isinstance(snapshot_id, str) and snapshot_id:
return snapshot_id, False
legacy_snapshot_id = snapshots.get(task_id)
if isinstance(legacy_snapshot_id, str) and legacy_snapshot_id:
return legacy_snapshot_id, True
return None, False
def _store_direct_snapshot(task_id: str, snapshot_id: str) -> None:
"""Persist the direct Modal snapshot id under the direct namespace."""
snapshots = _load_snapshots()
snapshots[_direct_snapshot_key(task_id)] = snapshot_id
snapshots.pop(task_id, None)
_save_snapshots(snapshots)
def _delete_direct_snapshot(task_id: str, snapshot_id: str | None = None) -> None:
"""Remove direct Modal snapshot entries for a task, including legacy keys."""
snapshots = _load_snapshots()
updated = False
for key in (_direct_snapshot_key(task_id), task_id):
value = snapshots.get(key)
if value is None:
continue
if snapshot_id is None or value == snapshot_id:
snapshots.pop(key, None)
updated = True
if updated:
_save_snapshots(snapshots)
def _resolve_modal_image(image_spec: Any) -> Any:
"""Convert registry references or snapshot ids into Modal image objects."""
import modal as _modal
if not isinstance(image_spec, str):
return image_spec
if image_spec.startswith("im-"):
return _modal.Image.from_id(image_spec)
return _modal.Image.from_registry(
image_spec,
setup_dockerfile_commands=[
"RUN rm -rf /usr/local/lib/python*/site-packages/pip* 2>/dev/null; "
"python -m ensurepip --upgrade --default-pip 2>/dev/null || true",
],
)
class _AsyncWorker:
"""Background thread with its own event loop for async-safe swe-rex calls."""
def __init__(self):
self._loop: Optional[asyncio.AbstractEventLoop] = None
@@ -101,42 +162,20 @@ class ModalEnvironment(BaseEnvironment):
sandbox_kwargs = dict(modal_sandbox_kwargs or {})
# If persistent, try to restore from a previous snapshot
restored_image = None
restored_snapshot_id = None
restored_from_legacy_key = False
if self._persistent:
snapshot_id = _load_snapshots().get(self._task_id)
if snapshot_id:
try:
import modal
restored_image = modal.Image.from_id(snapshot_id)
logger.info("Modal: restoring from snapshot %s", snapshot_id[:20])
except Exception as e:
logger.warning("Modal: failed to restore snapshot, using base image: %s", e)
restored_image = None
restored_snapshot_id, restored_from_legacy_key = _get_snapshot_restore_candidate(self._task_id)
if restored_snapshot_id:
logger.info("Modal: restoring from snapshot %s", restored_snapshot_id[:20])
effective_image = restored_image if restored_image else image
# Pre-build a modal.Image with pip fix for Modal's legacy image builder.
# Some task images have broken pip; fix via ensurepip before Modal uses it.
import modal as _modal
if isinstance(effective_image, str):
effective_image = _modal.Image.from_registry(
effective_image,
setup_dockerfile_commands=[
"RUN rm -rf /usr/local/lib/python*/site-packages/pip* 2>/dev/null; "
"python -m ensurepip --upgrade --default-pip 2>/dev/null || true",
],
)
# Start the async worker thread and create the deployment on it
# so all gRPC channels are bound to the worker's event loop.
self._worker.start()
from swerex.deployment.modal import ModalDeployment
async def _create_and_start():
async def _create_and_start(image_spec: Any):
deployment = ModalDeployment(
image=effective_image,
image=image_spec,
startup_timeout=180.0,
runtime_timeout=3600.0,
deployment_timeout=3600.0,
@@ -146,7 +185,30 @@ class ModalEnvironment(BaseEnvironment):
await deployment.start()
return deployment
self._deployment = self._worker.run_coroutine(_create_and_start())
try:
target_image_spec = restored_snapshot_id or image
try:
effective_image = _resolve_modal_image(target_image_spec)
self._deployment = self._worker.run_coroutine(_create_and_start(effective_image))
except Exception as exc:
if not restored_snapshot_id:
raise
logger.warning(
"Modal: failed to restore snapshot %s, retrying with base image: %s",
restored_snapshot_id[:20],
exc,
)
_delete_direct_snapshot(self._task_id, restored_snapshot_id)
base_image = _resolve_modal_image(image)
self._deployment = self._worker.run_coroutine(_create_and_start(base_image))
else:
if restored_snapshot_id and restored_from_legacy_key:
_store_direct_snapshot(self._task_id, restored_snapshot_id)
logger.info("Modal: migrated legacy snapshot entry for task %s", self._task_id)
except Exception:
self._worker.stop()
raise
def execute(self, command: str, cwd: str = "", *,
timeout: int | None = None,
@@ -160,7 +222,7 @@ class ModalEnvironment(BaseEnvironment):
exec_command, sudo_stdin = self._prepare_command(command)
# Modal sandboxes execute commands via the Modal SDK and cannot pipe
# subprocess stdin directly the way a local Popen can. When a sudo
# subprocess stdin directly the way a local Popen can. When a sudo
# password is present, use a shell-level pipe from printf so that the
# password feeds sudo -S without appearing as an echo argument embedded
# in the shell string.
@@ -175,7 +237,6 @@ class ModalEnvironment(BaseEnvironment):
effective_cwd = cwd or self.cwd
effective_timeout = timeout or self.timeout
# Run in a background thread so we can poll for interrupts
result_holder = {"value": None, "error": None}
def _run():
@@ -191,6 +252,7 @@ class ModalEnvironment(BaseEnvironment):
merge_output_streams=True,
)
)
output = self._worker.run_coroutine(_do_execute())
result_holder["value"] = {
"output": output.stdout,
@@ -227,7 +289,7 @@ class ModalEnvironment(BaseEnvironment):
if self._persistent:
try:
sandbox = getattr(self._deployment, '_sandbox', None)
sandbox = getattr(self._deployment, "_sandbox", None)
if sandbox:
async def _snapshot():
img = await sandbox.snapshot_filesystem.aio()
@@ -239,11 +301,12 @@ class ModalEnvironment(BaseEnvironment):
snapshot_id = None
if snapshot_id:
snapshots = _load_snapshots()
snapshots[self._task_id] = snapshot_id
_save_snapshots(snapshots)
logger.info("Modal: saved filesystem snapshot %s for task %s",
snapshot_id[:20], self._task_id)
_store_direct_snapshot(self._task_id, snapshot_id)
logger.info(
"Modal: saved filesystem snapshot %s for task %s",
snapshot_id[:20],
self._task_id,
)
except Exception as e:
logger.warning("Modal: filesystem snapshot failed: %s", e)