-
Notifications
You must be signed in to change notification settings - Fork 615
Support restricted shared agent runtimes #2125
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
46d7bfe
a9aea26
b92d9c0
67e02ba
deebb91
1bede2a
e276e6f
f2ff094
c741127
bd752cf
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -9,6 +9,7 @@ | |
| import uuid | ||
| import weakref | ||
| from abc import ABC, abstractmethod | ||
| from collections.abc import AsyncIterator | ||
| from dataclasses import dataclass | ||
| from pathlib import PurePosixPath | ||
| from typing import ClassVar, Self | ||
|
|
@@ -40,6 +41,57 @@ | |
| f"|| {{ {_INSTALL_CURL}; {_DOWNLOAD_UV}; }}" | ||
| ) | ||
|
|
||
| # A reused restricted box must stop processes and discard framework setup caches | ||
| # left by the prior untrusted agent before trusted setup reopens egress. | ||
| # PID + kernel start time survives PID reuse. | ||
| _PROCESS_SNAPSHOT = r""" | ||
| for proc in /proc/[0-9]*; do | ||
| pid=${proc##*/} | ||
| stat=$(cat "$proc/stat" 2>/dev/null) || continue | ||
| rest=${stat##*) } | ||
| set -- $rest | ||
| printf '%s:%s\n' "$pid" "${20}" | ||
| done | ||
| """ | ||
| _KILL_AFTER_SNAPSHOT = r""" | ||
| baseline=" | ||
| $VF_PROCESS_BASELINE | ||
| " | ||
| protected=" " | ||
| pid=$$ | ||
| while [ "$pid" -gt 1 ] 2>/dev/null; do | ||
| protected="$protected$pid " | ||
| stat=$(cat "/proc/$pid/stat" 2>/dev/null) || break | ||
| rest=${stat##*) } | ||
| set -- $rest | ||
| pid=$2 | ||
| done | ||
| while :; do | ||
| targeted= | ||
| for proc in /proc/[0-9]*; do | ||
| pid=${proc##*/} | ||
| case "$protected" in *" $pid "*) continue ;; esac | ||
| stat=$(cat "$proc/stat" 2>/dev/null) || continue | ||
| rest=${stat##*) } | ||
| set -- $rest | ||
| case "$1" in Z|X) continue ;; esac | ||
| if [ "$1" = D ]; then | ||
| echo "process $pid is stuck in uninterruptible sleep" >&2 | ||
| exit 1 | ||
| fi | ||
| identity="$pid:${20}" | ||
| case "$baseline" in *" | ||
| $identity | ||
| "*) continue ;; esac | ||
| if kill -STOP "$pid" 2>/dev/null; then | ||
| kill -KILL "$pid" 2>/dev/null || true | ||
| targeted=1 | ||
| fi | ||
| done | ||
| [ -z "$targeted" ] && break | ||
| done | ||
| """ | ||
|
cursor[bot] marked this conversation as resolved.
|
||
|
|
||
| # The single port a self-publishing runtime (modal/prime) forwards to a public URL for a server | ||
| # hosted in its sandbox. A server placed in such a runtime binds this (on 0.0.0.0) and is reached | ||
| # at the runtime's public URL. | ||
|
|
@@ -143,6 +195,9 @@ def __init__(self, name: str | None = None) -> None: | |
| self._uv_interpreters: dict[str, str] = {} | ||
| self._uv_script_locks: dict[str, asyncio.Lock] = {} | ||
| self._setup_claimed = False | ||
| self._process_baseline: str | None = None | ||
| self.execution_prepared = False | ||
| """Whether a rollout successfully activated this runtime's execution policy.""" | ||
| self.stopped = False | ||
| """Whether teardown has begun (set by `stop`). A stopped runtime is dead: a rollout | ||
| refuses to borrow one — the owner tore it down, so any use is a lifetime bug in the | ||
|
|
@@ -272,20 +327,70 @@ def host_url(self, url: str) -> str: | |
| """The URL a program inside this runtime uses to reach a host-bound `url`.""" | ||
| return url | ||
|
|
||
| async def prepare_setup(self) -> None: | ||
| """Claim the runtime for trusted setup; restricted runtimes may reject reuse.""" | ||
| @contextlib.asynccontextmanager | ||
| async def rollout(self) -> AsyncIterator[None]: | ||
| """Claim one restricted rollout, reopening trusted setup only between runs.""" | ||
| if not self.network_restricted: | ||
| yield | ||
| return | ||
| if self._setup_claimed: | ||
| raise SandboxError( | ||
| f"network-filtered {self.type} runtimes are single-rollout; " | ||
| "provision a fresh runtime instead of reusing this one" | ||
| f"network-filtered {self.type} runtime {self.name!r} is already in " | ||
| "use; wait for its rollout to finish before reusing it" | ||
| ) | ||
| self._setup_claimed = True | ||
| try: | ||
| if self._process_baseline is not None: | ||
| result = await self.run( | ||
| ["sh", "-c", _KILL_AFTER_SNAPSHOT], | ||
| {"VF_PROCESS_BASELINE": self._process_baseline}, | ||
|
Comment on lines
+344
to
+346
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
When the first agent runs as root, as it does in the default Docker Python image, it can replace Useful? React with 👍 / 👎. |
||
| ) | ||
| if result.exit_code != 0: | ||
| raise SandboxError( | ||
| f"failed to isolate reused {self.type} runtime processes: " | ||
| f"{result.stderr.strip()[-500:]}" | ||
| ) | ||
| uv_envs = [ | ||
| str(PurePosixPath(path).parent.parent) | ||
| for path in self._uv_interpreters.values() | ||
| ] | ||
| reset = await self.run( | ||
| ["sh", "-c", 'rm -rf "$@" /tmp/vf-*', "reset", *uv_envs], | ||
| {}, | ||
| ) | ||
| if reset.exit_code != 0: | ||
| raise SandboxError( | ||
| f"failed to reset reused {self.type} runtime setup caches: " | ||
| f"{reset.stderr.strip()[-500:]}" | ||
| ) | ||
| self._uv_interpreters.clear() | ||
| snapshot = await self.run(["sh", "-c", _PROCESS_SNAPSHOT], {}) | ||
| if snapshot.exit_code != 0: | ||
| raise SandboxError( | ||
| f"failed to snapshot {self.type} runtime processes: " | ||
| f"{snapshot.stderr.strip()[-500:]}" | ||
| ) | ||
| self._process_baseline = snapshot.stdout | ||
| if self.execution_prepared: | ||
| await self._apply_network_policy(None) | ||
|
xeophon marked this conversation as resolved.
|
||
| self.execution_prepared = False | ||
|
cursor[bot] marked this conversation as resolved.
xeophon marked this conversation as resolved.
|
||
| yield | ||
| finally: | ||
| self._setup_claimed = False | ||
|
|
||
| async def prepare_execution(self, routes: list[str]) -> None: | ||
| """Last setup step, right before the agent starts. Restricted runtimes enforce | ||
| their policy here while keeping the interception and MCP `routes` reachable.""" | ||
| if not self.network_restricted: | ||
| return | ||
| await self._apply_network_policy(routes) | ||
| self.execution_prepared = True | ||
|
|
||
| async def _apply_network_policy(self, routes: list[str] | None) -> None: | ||
| """Apply execution routes, or restore unrestricted trusted setup for None.""" | ||
| raise NotImplementedError( | ||
| f"{type(self).__name__} does not support restricted networking" | ||
| ) | ||
|
|
||
| @property | ||
| def network_restricted(self) -> bool: | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -233,11 +233,14 @@ def host_url(self, url: str) -> str: | |
| return url.replace(host, "host.docker.internal", 1) | ||
| return url | ||
|
|
||
| async def prepare_execution(self, routes: list[str]) -> None: | ||
| async def _apply_network_policy(self, routes: list[str] | None) -> None: | ||
| """Allow the declared framework routes, then leave the proxy as the only route.""" | ||
| if not self.network_restricted: | ||
| return | ||
| assert self._proxy is not None | ||
| if routes is None: | ||
| self._proxy.policy = NetworkPolicy( | ||
| ["*"], [], [HOST_ALIAS], allow_non_global=True | ||
| ) | ||
| return | ||
|
Comment on lines
+239
to
+243
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
When a restricted Docker box is reused after the first rollout, this branch only makes the HTTP proxy permissive; the container routes are still cut because Useful? React with 👍 / 👎. |
||
| framework = [ | ||
| urlsplit(url)._replace(path="", query="", fragment="").geturl() | ||
| for url in routes | ||
|
|
@@ -247,6 +250,8 @@ async def prepare_execution(self, routes: list[str]) -> None: | |
| self.config.block, | ||
| framework, | ||
| ) | ||
| if self._cut: | ||
| return | ||
| script = ( | ||
| "set -eu; HOST=$1; " | ||
| "PORT=$2; " | ||
|
|
||
Uh oh!
There was an error while loading. Please reload this page.