diff --git a/README.md b/README.md index 7c8af38d9..a66d0ac8a 100644 --- a/README.md +++ b/README.md @@ -35,7 +35,7 @@ limitations under the License. ✨ **Simplicity**: the logic for agents fits in ~1,000 lines of code (see [agents.py](https://github.com/huggingface/smolagents/blob/main/src/smolagents/agents.py)). We kept abstractions to their minimal shape above raw code! -🧑‍💻 **First-class support for Code Agents**. Our [`CodeAgent`](https://huggingface.co/docs/smolagents/reference/agents#smolagents.CodeAgent) writes its actions in code (as opposed to "agents being used to write code"). To make it secure, we support executing in sandboxed environments via [Blaxel](https://blaxel.ai), [E2B](https://e2b.dev/), [Modal](https://modal.com/), or Docker. +🧑‍💻 **First-class support for Code Agents**. Our [`CodeAgent`](https://huggingface.co/docs/smolagents/reference/agents#smolagents.CodeAgent) writes its actions in code (as opposed to "agents being used to write code"). To make it secure, we support executing in sandboxed environments via [Blaxel](https://blaxel.ai), [E2B](https://e2b.dev/), [Modal](https://modal.com/), [Tenki](https://tenki.cloud), or Docker. 🤗 **Hub integrations**: you can [share/pull tools or agents to/from the Hub](https://huggingface.co/docs/smolagents/reference/tools#smolagents.Tool.from_hub) for instant sharing of the most efficient agents! @@ -239,7 +239,7 @@ for request in requests_to_search: Writing actions as code snippets is demonstrated to work better than the current industry practice of letting the LLM output a dictionary of the tools it wants to call: [uses 30% fewer steps](https://huggingface.co/papers/2402.01030) (thus 30% fewer LLM calls) and [reaches higher performance on difficult benchmarks](https://huggingface.co/papers/2411.01747). Head to [our high-level intro to agents](https://huggingface.co/docs/smolagents/conceptual_guides/intro_agents) to learn more on that. Since code execution can be a serious security concern (arbitrary code execution!), **you should run agent code in a sandbox**. We support several options: - - [E2B](https://e2b.dev/), [Blaxel](https://blaxel.ai), [Modal](https://modal.com/) — managed cloud sandboxes, simplest to set up + - [E2B](https://e2b.dev/), [Blaxel](https://blaxel.ai), [Modal](https://modal.com/), [Tenki](https://tenki.cloud) — managed cloud sandboxes, simplest to set up - [Docker](https://www.docker.com/) — self-hosted container isolation The built-in `LocalPythonExecutor` is **not a security sandbox**. It applies some restrictions but can be bypassed and must not be used as a security boundary. diff --git a/docs/source/en/guided_tour.md b/docs/source/en/guided_tour.md index 6f6bb68f8..80b702b76 100644 --- a/docs/source/en/guided_tour.md +++ b/docs/source/en/guided_tour.md @@ -85,7 +85,7 @@ This could also be authorized by using `numpy.*`, which will allow `numpy` as we The execution will stop at any code trying to perform an illegal operation or if there is a regular Python error with the code generated by the agent. -You can also use [Blaxel](https://blaxel.ai), [E2B](https://e2b.dev/docs#what-is-e2-b), or Docker instead of a local Python interpreter. For Blaxel, first [set the `BL_API_KEY` and `BL_WORKSPACE` environment variables](https://app.blaxel.ai/profile/security) and then pass `executor_type="blaxel"` upon agent initialization. For E2B, first [set the `E2B_API_KEY` environment variable](https://e2b.dev/dashboard?tab=keys) and then pass `executor_type="e2b"`. For Docker, pass `executor_type="docker"`. +You can also use [Blaxel](https://blaxel.ai), [E2B](https://e2b.dev/docs#what-is-e2-b), [Tenki](https://tenki.cloud), or Docker instead of a local Python interpreter. For Blaxel, first [set the `BL_API_KEY` and `BL_WORKSPACE` environment variables](https://app.blaxel.ai/profile/security) and then pass `executor_type="blaxel"` upon agent initialization. For E2B, first [set the `E2B_API_KEY` environment variable](https://e2b.dev/dashboard?tab=keys) and then pass `executor_type="e2b"`. For Tenki, first set the `TENKI_API_KEY` environment variable (plus `TENKI_PROJECT_ID` if your account has more than one project) and then pass `executor_type="tenki"`. For Docker, pass `executor_type="docker"`. > [!TIP] diff --git a/docs/source/en/index.md b/docs/source/en/index.md index 90621715f..397560706 100644 --- a/docs/source/en/index.md +++ b/docs/source/en/index.md @@ -12,7 +12,7 @@ Key features of `smolagents` include: ✨ **Simplicity**: The logic for agents fits in ~thousand lines of code. We kept abstractions to their minimal shape above raw code! -🧑‍💻 **First-class support for Code Agents**: [`CodeAgent`](reference/agents#smolagents.CodeAgent) writes its actions in code (as opposed to "agents being used to write code") to invoke tools or perform computations, enabling natural composability (function nesting, loops, conditionals). To make it secure, we support [executing in sandboxed environment](tutorials/secure_code_execution) via [Modal](https://modal.com/), [Blaxel](https://blaxel.ai), [E2B](https://e2b.dev/), or Docker. +🧑‍💻 **First-class support for Code Agents**: [`CodeAgent`](reference/agents#smolagents.CodeAgent) writes its actions in code (as opposed to "agents being used to write code") to invoke tools or perform computations, enabling natural composability (function nesting, loops, conditionals). To make it secure, we support [executing in sandboxed environment](tutorials/secure_code_execution) via [Modal](https://modal.com/), [Blaxel](https://blaxel.ai), [E2B](https://e2b.dev/), [Tenki](https://tenki.cloud), or Docker. 📡 **Common Tool-Calling Agent Support**: In addition to CodeAgents, [`ToolCallingAgent`](reference/agents#smolagents.ToolCallingAgent) supports usual JSON/text-based tool-calling for scenarios where that paradigm is preferred. diff --git a/docs/source/en/installation.md b/docs/source/en/installation.md index 44e8f3bf7..5249387d4 100644 --- a/docs/source/en/installation.md +++ b/docs/source/en/installation.md @@ -188,6 +188,10 @@ Extras for executing code remotely: ```bash pip install "smolagents[e2b]" ``` +- **tenki**: Enable Tenki support for remote execution. + ```bash + pip install "smolagents[tenki]" + ``` - **docker**: Add support for executing code in Docker containers. ```bash pip install "smolagents[docker]" @@ -202,6 +206,10 @@ Extras for executing code remotely: ```bash uv pip install "smolagents[e2b]" ``` +- **tenki**: Enable Tenki support for remote execution. + ```bash + uv pip install "smolagents[tenki]" + ``` - **docker**: Add support for executing code in Docker containers. ```bash uv pip install "smolagents[docker]" diff --git a/docs/source/en/reference/python_executors.md b/docs/source/en/reference/python_executors.md index 47c4ff937..bc9923f60 100644 --- a/docs/source/en/reference/python_executors.md +++ b/docs/source/en/reference/python_executors.md @@ -32,6 +32,10 @@ available executor implementations. [[autodoc]] smolagents.remote_executors.ModalExecutor +### TenkiExecutor + +[[autodoc]] smolagents.remote_executors.TenkiExecutor + ### DockerExecutor [[autodoc]] smolagents.remote_executors.DockerExecutor diff --git a/docs/source/en/tutorials/secure_code_execution.md b/docs/source/en/tutorials/secure_code_execution.md index 7ab247f0a..92b56655a 100644 --- a/docs/source/en/tutorials/secure_code_execution.md +++ b/docs/source/en/tutorials/secure_code_execution.md @@ -118,7 +118,7 @@ When working with AI agents that execute code, security is paramount. There are ![Sandbox approaches comparison](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/smolagents/sandboxed_execution.png) -1. **Running individual code snippets in a sandbox**: This approach (left side of diagram) only executes the agent-generated Python code snippets in a sandbox while keeping the rest of the agentic system in your local environment. It's simpler to set up using `executor_type="blaxel"`, `executor_type="e2b"`, `executor_type="modal"`, or +1. **Running individual code snippets in a sandbox**: This approach (left side of diagram) only executes the agent-generated Python code snippets in a sandbox while keeping the rest of the agentic system in your local environment. It's simpler to set up using `executor_type="blaxel"`, `executor_type="e2b"`, `executor_type="modal"`, `executor_type="tenki"`, or `executor_type="docker"`, but it doesn't support multi-agents and still requires passing state data between your environment and the sandbox. 2. **Running the entire agentic system in a sandbox**: This approach (right side of diagram) runs the entire agentic system, including the agent, model, and tools, within a sandbox environment. This provides better isolation but requires more manual setup and may require passing sensitive credentials (like API keys) to the sandbox environment. @@ -279,6 +279,41 @@ with CodeAgent(model=InferenceClientModel(), tools=[], executor_type="modal") as The agent state and generated code from the `InferenceClientModel` are sent to a Modal sandbox, which can securely execute code inside them. +### Tenki setup + +#### Installation + +1. Create a Tenki account at [tenki.cloud](https://tenki.cloud) +2. Install the required packages: +```bash +pip install 'smolagents[tenki]' +``` +3. Set your API key as an environment variable: +```bash +export TENKI_API_KEY="your-api-key" +``` +If your account has more than one project, also set `TENKI_PROJECT_ID` (or pass `executor_kwargs={"create_kwargs": {"project_id": ...}}`); +with a single project it is detected automatically. + +#### Running your agent in Tenki: quick start + +We provide a simple way to use a Tenki Sandbox: simply add `executor_type="tenki"` to the agent initialization, as follows: + +```py +from smolagents import InferenceClientModel, CodeAgent + +with CodeAgent(model=InferenceClientModel(), tools=[], executor_type="tenki") as agent: + agent.run("What is the 42th Fibonacci number?") +``` + +> [!TIP] +> Using the agent as a context manager (with the `with` statement) ensures that the Tenki sandbox is cleaned up immediately after the agent completes its task. +> Alternatively, you can manually call the agent's `cleanup()` method. + +Each agent run executes in a disposable Linux microVM with kernel-level isolation, booted from a pre-warmed pool. +You can customize the sandbox via `executor_kwargs`, e.g. `executor_kwargs={"create_kwargs": {"cpu_cores": 4, "memory_mb": 8192}}`; +the sandbox image only needs to provide `python3` (`pip` and the Jupyter Kernel Gateway are installed on startup if missing). + ### Docker setup #### Installation @@ -434,7 +469,7 @@ finally: ### Best practices for sandboxes -These key practices apply to Blaxel, E2B, and Docker sandboxes: +These key practices apply to Blaxel, E2B, Tenki, and Docker sandboxes: - Resource management - Set memory and CPU limits diff --git a/examples/sandboxed_execution.py b/examples/sandboxed_execution.py index 30b89d5db..3dc05c1d7 100644 --- a/examples/sandboxed_execution.py +++ b/examples/sandboxed_execution.py @@ -22,3 +22,8 @@ with CodeAgent(tools=[WebSearchTool()], model=model, executor_type="modal") as agent: output = agent.run("How many seconds would it take for a leopard at full speed to run through Pont des Arts?") print("Modal executor result:", output) + +# Tenki executor example +with CodeAgent(tools=[WebSearchTool()], model=model, executor_type="tenki") as agent: + output = agent.run("How many seconds would it take for a leopard at full speed to run through Pont des Arts?") +print("Tenki executor result:", output) diff --git a/pyproject.toml b/pyproject.toml index 61ead4273..c2987f43c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -63,6 +63,10 @@ modal = [ "modal>=1.1.3", "websocket-client", ] +tenki = [ + "tenki-sandbox>=0.4.0", + "websocket-client", +] openai = [ "openai>=1.58.1" ] @@ -90,7 +94,7 @@ vllm = [ "torch" ] all = [ - "smolagents[audio,blaxel,docker,e2b,gradio,litellm,mcp,mlx-lm,modal,openai,telemetry,toolkit,transformers,vision,bedrock]", + "smolagents[audio,blaxel,docker,e2b,gradio,litellm,mcp,mlx-lm,modal,openai,telemetry,tenki,toolkit,transformers,vision,bedrock]", ] quality = [ "ruff>=0.9.0", diff --git a/src/smolagents/agents.py b/src/smolagents/agents.py index a8f6e52de..3f46552e7 100644 --- a/src/smolagents/agents.py +++ b/src/smolagents/agents.py @@ -77,7 +77,7 @@ Monitor, TokenUsage, ) -from .remote_executors import BlaxelExecutor, DockerExecutor, E2BExecutor, ModalExecutor +from .remote_executors import BlaxelExecutor, DockerExecutor, E2BExecutor, ModalExecutor, TenkiExecutor from .tools import BaseTool, Tool, validate_tool_arguments from .utils import ( AgentError, @@ -1513,7 +1513,7 @@ class CodeAgent(MultiStepAgent): additional_authorized_imports (`list[str]`, *optional*): Additional authorized imports for the agent. planning_interval (`int`, *optional*): Interval at which the agent will run a planning step. executor ([`PythonExecutor`], *optional*): Custom Python code executor. If not provided, a default executor will be created based on `executor_type`. - executor_type (`Literal["local", "blaxel", "e2b", "modal", "docker"]`, default `"local"`): Type of code executor. + executor_type (`Literal["local", "blaxel", "e2b", "modal", "docker", "tenki"]`, default `"local"`): Type of code executor. executor_kwargs (`dict`, *optional*): Additional arguments to pass to initialize the executor. max_print_outputs_length (`int`, *optional*): Maximum length of the print outputs. stream_outputs (`bool`, *optional*, default `False`): Whether to stream outputs during execution. @@ -1532,7 +1532,7 @@ def __init__( additional_authorized_imports: list[str] | None = None, planning_interval: int | None = None, executor: PythonExecutor = None, - executor_type: Literal["local", "blaxel", "e2b", "modal", "docker"] = "local", + executor_type: Literal["local", "blaxel", "e2b", "modal", "docker", "tenki"] = "local", executor_kwargs: dict[str, Any] | None = None, max_print_outputs_length: int | None = None, stream_outputs: bool = False, @@ -1596,7 +1596,7 @@ def cleanup(self): self.python_executor.cleanup() def create_python_executor(self) -> PythonExecutor: - if self.executor_type not in {"local", "blaxel", "e2b", "modal", "docker"}: + if self.executor_type not in {"local", "blaxel", "e2b", "modal", "docker", "tenki"}: raise ValueError(f"Unsupported executor type: {self.executor_type}") if self.executor_type == "local": @@ -1612,6 +1612,7 @@ def create_python_executor(self) -> PythonExecutor: "e2b": E2BExecutor, "docker": DockerExecutor, "modal": ModalExecutor, + "tenki": TenkiExecutor, } return remote_executors[self.executor_type]( self.additional_authorized_imports, self.logger, **self.executor_kwargs diff --git a/src/smolagents/remote_executors.py b/src/smolagents/remote_executors.py index 1704a148b..d439ae0a3 100644 --- a/src/smolagents/remote_executors.py +++ b/src/smolagents/remote_executors.py @@ -17,6 +17,7 @@ import base64 import inspect import json +import os import pickle import re import secrets @@ -39,7 +40,7 @@ from .utils import AgentError -__all__ = ["BlaxelExecutor", "E2BExecutor", "ModalExecutor", "DockerExecutor"] +__all__ = ["BlaxelExecutor", "E2BExecutor", "ModalExecutor", "DockerExecutor", "TenkiExecutor"] try: @@ -1074,3 +1075,204 @@ def __del__(self): self.cleanup() except Exception: pass # Silently ignore errors during cleanup + + +class TenkiExecutor(RemotePythonExecutor): + """ + Remote Python code executor in a Tenki sandbox. + + Tenki provides disposable Linux microVMs with kernel-level isolation that boot from + pre-warmed pools in a few seconds. Authentication is read from the `TENKI_API_KEY` + (or `TENKI_AUTH_TOKEN`) environment variable, or can be passed via `create_kwargs={"auth_token": ...}`. + + Args: + additional_imports (`list[str]`): Additional Python packages to install. + logger (`Logger`): Logger to use for output and errors. + allow_pickle (`bool`, default `False`): Whether to allow pickle serialization for objects that cannot be safely serialized to JSON. + - `False` (default, recommended): Only safe JSON serialization is used. Raises error if object cannot be safely serialized. + - `True` (legacy mode): Tries safe JSON serialization first, falls back to pickle with warning if needed. + + **Security Warning:** Pickle deserialization can execute arbitrary code. Only set `allow_pickle=True` + if you fully trust the execution environment and need backward compatibility with custom types. + sandbox_name (`str`, *optional*): Name for the sandbox. Defaults to "smolagent-executor-" followed by a random suffix. + port (`int`, default `8888`): Port for the Jupyter Kernel Gateway to bind to inside the sandbox. + create_kwargs (`dict`, *optional*): Additional keyword arguments to pass to the Tenki Sandbox create command, + e.g. `project_id`, `image`, `cpu_cores`, `memory_mb`. The sandbox image must provide + `python3`; the Jupyter Kernel Gateway (and `pip`, if missing) is installed on startup, so an image with + them preinstalled boots faster. If `project_id` is not provided, it is read from the `TENKI_PROJECT_ID` + environment variable, or resolved automatically when the account (or the given `workspace_id`) has a + single project. + """ + + # Bootstraps the Jupyter Kernel Gateway on any python3-capable image: no-op if already installed. + # Everything is installed to the user site so that the kernel and the agent's in-kernel `!pip install` + # share one interpreter without touching the externally-managed system site. The user-site directory is + # created upfront because CPython only adds it to sys.path when it exists at interpreter startup. + _SETUP_KERNEL_GATEWAY_COMMAND = dedent("""\ + mkdir -p "$(python3 -m site --user-site)" + python3 -c 'import kernel_gateway' 2>/dev/null || { + python3 -m pip --version >/dev/null 2>&1 \\ + || python3 -m ensurepip --user >/dev/null 2>&1 \\ + || { sudo apt-get update -qq && sudo apt-get install -y -qq python3-pip >/dev/null; } + PIP_BREAK_SYSTEM_PACKAGES=1 python3 -m pip install --quiet --user jupyter_kernel_gateway ipykernel + }""") + + def __init__( + self, + additional_imports: list[str], + logger, + allow_pickle: bool = False, + sandbox_name: str | None = None, + port: int = 8888, + create_kwargs: Optional[dict] = None, + ): + super().__init__(additional_imports, logger, allow_pickle) + try: + from tenki_sandbox import Client + except ModuleNotFoundError: + raise ModuleNotFoundError( + """Please install 'tenki' extra to use TenkiExecutor: `pip install 'smolagents[tenki]'`""" + ) + + self.port = int(port) + self._cleaned_up = False + token = secrets.token_urlsafe(16) + create_kwargs = { + "name": sandbox_name or f"smolagent-executor-{uuid.uuid4().hex[:8]}", + **(create_kwargs or {}), + } + client_kwargs = { + key: create_kwargs.pop(key) + for key in ("auth_token", "base_url", "gateway_url", "cookie_name", "timeout") + if key in create_kwargs + } + self.client = Client(**client_kwargs) + try: + if not create_kwargs.get("project_id"): + create_kwargs["project_id"] = self._resolve_project_id(self.client, create_kwargs.get("workspace_id")) + + self.logger.log("Starting Tenki sandbox", level=LogLevel.INFO) + # Allocate without the server-side readiness wait, assign the handle, THEN wait: if + # wait_ready() fails, self.sandbox is already set so cleanup() can terminate the session. + # create(wait=True) would raise mid-wait before this assignment, leaking the allocated session. + create_kwargs["wait"] = False + self.sandbox = self.client.create(**create_kwargs) + self.sandbox.wait_ready() + setup_result = self.sandbox.shell(self._SETUP_KERNEL_GATEWAY_COMMAND, timeout=600) + if not setup_result.ok: + raise RuntimeError( + f"Kernel gateway setup failed with exit code {setup_result.exit_code}: " + f"{setup_result.stderr_text.strip()[-2000:]}" + ) + self.logger.log("Starting Jupyter Kernel Gateway", level=LogLevel.INFO) + self._kernel_gateway_process = self.sandbox.start( + "bash", + "-lc", + 'export PATH="$HOME/.local/bin:$PATH";' + f" exec jupyter kernelgateway --KernelGatewayApp.ip=0.0.0.0 --KernelGatewayApp.port={self.port}", + # PIP_USER/PIP_BREAK_SYSTEM_PACKAGES let the kernel's `!pip install` work on externally-managed images + env={"KG_AUTH_TOKEN": token, "PIP_USER": "1", "PIP_BREAK_SYSTEM_PACKAGES": "1"}, + ) + base_url = self.sandbox.expose_port(self.port).url.rstrip("/") + self._wait_for_server(base_url, token) + + kernel_id = _create_kernel_http(f"{base_url}/api/kernels?token={token}", self.logger) + ws_scheme = "wss" if base_url.startswith("https") else "ws" + ws_base = base_url.replace("https://", "").replace("http://", "") + self.ws_url = f"{ws_scheme}://{ws_base}/api/kernels/{kernel_id}/channels?token={token}" + + self.installed_packages = self.install_packages(additional_imports) + self.logger.log("Tenki sandbox is running", level=LogLevel.INFO) + except ValueError: + self.cleanup() + raise + except Exception as e: + self.cleanup() + raise RuntimeError(f"Failed to initialize Tenki sandbox: {e}") from e + except BaseException: + # KeyboardInterrupt / cancellation inherit from BaseException, not Exception; without this + # a Ctrl+C during the multi-second startup would leak the sandbox and client. + self.cleanup() + raise + + @staticmethod + def _resolve_project_id(client, workspace_id: str | None = None) -> str: + """Resolve the Tenki project id from the environment or the account's sole project.""" + project_id = os.getenv("TENKI_PROJECT_ID") + if project_id: + return project_id + identity = client.who_am_i() + workspaces = [ws for ws in identity.workspaces if workspace_id is None or ws.id == workspace_id] + projects = [project for workspace in workspaces for project in workspace.projects] + if len(projects) != 1: + raise ValueError( + f"Could not determine which Tenki project to use ({len(projects)} projects found): " + "set the TENKI_PROJECT_ID environment variable or pass create_kwargs={'project_id': ...} " + '(with CodeAgent: executor_kwargs={"create_kwargs": {"project_id": ...}})' + ) + return projects[0].id + + def run_code_raise_errors(self, code: str) -> CodeOutput: + """ + Execute Python code in the Tenki sandbox and return the result. + + Args: + code (`str`): Python code to execute. + + Returns: + `CodeOutput`: Code output containing the result, logs, and whether it is the final answer. + """ + from websocket import create_connection + + with closing(create_connection(self.ws_url)) as ws: + return _websocket_run_code_raise_errors(code, ws, self.logger, self.allow_pickle) + + def cleanup(self): + """Clean up the Tenki sandbox by terminating it.""" + if self._cleaned_up: + return + if hasattr(self, "sandbox"): + try: + self.logger.log("Shutting down sandbox...", level=LogLevel.INFO) + self.sandbox.close_if_open() + self.logger.log("Sandbox cleanup completed", level=LogLevel.INFO) + except Exception as e: + # Keep the sandbox handle and client so a later cleanup() can retry, rather than + # letting a transient termination failure become a permanent leak. + self.logger.log_error(f"Error during cleanup, will retry on next attempt: {e}") + return + del self.sandbox + if hasattr(self, "client"): + try: + self.client.close() + except Exception as e: + self.logger.log_error(f"Error closing Tenki client: {e}") + self._cleaned_up = True + + def delete(self): + """Ensure cleanup on deletion.""" + self.cleanup() + + def _wait_for_server(self, base_url: str, token: str): + """Wait for the Jupyter Kernel Gateway to start up.""" + n_retries = 0 + while True: + try: + resp = requests.get(f"{base_url}/api/kernelspecs?token={token}", timeout=2) + if resp.status_code == 200: + break + except RequestException: + pass + n_retries += 1 + if n_retries % 10 == 0: + self.logger.log("Waiting for server to startup, retrying...", level=LogLevel.INFO) + if n_retries > 60: + error_message = "Unable to connect to the Jupyter Kernel Gateway in the Tenki sandbox" + try: + stderr = self._kernel_gateway_process.stderr.read_text().strip() + except Exception: + stderr = "" + if stderr: + error_message += f". Kernel gateway output:\n{stderr[-2000:]}" + raise RuntimeError(error_message) + time.sleep(1.0) diff --git a/tests/test_remote_executors.py b/tests/test_remote_executors.py index 740eee59d..bdb672a98 100644 --- a/tests/test_remote_executors.py +++ b/tests/test_remote_executors.py @@ -17,6 +17,7 @@ E2BExecutor, ModalExecutor, RemotePythonExecutor, + TenkiExecutor, ) from smolagents.serialization import SerializationError from smolagents.utils import AgentError @@ -615,3 +616,261 @@ def test_blaxel_executor_cleanup(self, mock_settings, mock_sandbox_instance, moc assert mock_delete_sandbox.sync.called # Verify sandbox reference was cleaned up assert not hasattr(executor, "sandbox") + + +@require_run_all +class TestTenkiExecutorIntegration(CommonDockerExecutorIntegration): + @pytest.fixture + def custom_executor(self): + executor = TenkiExecutor( + additional_imports=["pillow", "numpy"], + logger=AgentLogger(LogLevel.INFO, Console(force_terminal=False, file=io.StringIO())), + ) + yield executor + executor.delete() + + +class TestTenkiExecutorUnit: + def test_tenki_executor_instantiation_without_tenki_sdk(self): + """Test that TenkiExecutor raises appropriate error when tenki-sandbox SDK is not installed.""" + logger = MagicMock() + with patch.dict("sys.modules", {"tenki_sandbox": None}): + with pytest.raises(ModuleNotFoundError) as excinfo: + TenkiExecutor(additional_imports=[], logger=logger) + assert "Please install 'tenki' extra" in str(excinfo.value) + + @staticmethod + def _make_mock_client(mock_client_cls): + """Wire the mocked tenki_sandbox.Client with a sandbox whose init calls succeed.""" + mock_client = mock_client_cls.return_value + mock_sandbox = mock_client.create.return_value + mock_sandbox.shell.return_value = MagicMock(ok=True) + mock_sandbox.expose_port.return_value = MagicMock(url="https://test-sandbox.tenki.cloud") + return mock_client, mock_sandbox + + @patch.dict("os.environ", {"TENKI_PROJECT_ID": "proj-123"}) + @patch("smolagents.remote_executors.TenkiExecutor._wait_for_server") + @patch("smolagents.remote_executors._create_kernel_http") + @patch("tenki_sandbox.Client") + def test_tenki_executor_instantiation(self, mock_client_cls, mock_create_kernel, mock_wait_for_server): + """Test TenkiExecutor instantiation with mocked Tenki SDK.""" + logger = MagicMock() + mock_client, mock_sandbox = self._make_mock_client(mock_client_cls) + mock_create_kernel.return_value = "kernel-123" + + executor = TenkiExecutor(additional_imports=[], logger=logger) + + create_kwargs = mock_client.create.call_args.kwargs + assert create_kwargs["name"].startswith("smolagent-executor-") + assert create_kwargs["project_id"] == "proj-123" + # The sandbox is allocated with wait=False, then wait_ready() is awaited on the retained handle + assert create_kwargs["wait"] is False + mock_sandbox.wait_ready.assert_called_once() + # The kernel gateway bootstrap ran in the sandbox + assert "jupyter_kernel_gateway" in mock_sandbox.shell.call_args.args[0] + # The gateway process gets the auth token, and pip env vars for in-kernel `!pip install` + start_env = mock_sandbox.start.call_args.kwargs["env"] + assert start_env["KG_AUTH_TOKEN"] + assert start_env["PIP_USER"] == "1" + assert start_env["PIP_BREAK_SYSTEM_PACKAGES"] == "1" + mock_sandbox.expose_port.assert_called_once_with(8888) + assert executor.ws_url.startswith("wss://test-sandbox.tenki.cloud/api/kernels/kernel-123/channels?token=") + + @patch.dict("os.environ", {"TENKI_PROJECT_ID": "proj-123"}) + @patch("smolagents.remote_executors.TenkiExecutor.install_packages") + @patch("smolagents.remote_executors.TenkiExecutor._wait_for_server") + @patch("smolagents.remote_executors._create_kernel_http") + @patch("tenki_sandbox.Client") + def test_tenki_executor_custom_parameters( + self, mock_client_cls, mock_create_kernel, mock_wait_for_server, mock_install_packages + ): + """Test TenkiExecutor with custom parameters, including client-kwarg routing.""" + logger = MagicMock() + mock_client, mock_sandbox = self._make_mock_client(mock_client_cls) + mock_create_kernel.return_value = "kernel-123" + mock_install_packages.return_value = ["numpy"] + + executor = TenkiExecutor( + additional_imports=["numpy"], + logger=logger, + sandbox_name="test-sandbox", + port=9999, + create_kwargs={ + "project_id": "proj-456", + "image": "custom-image:latest", + "cpu_cores": 4, + "auth_token": "test-token", + }, + ) + + # Client-level kwargs are routed to the Client, not to sandbox creation + mock_client_cls.assert_called_once_with(auth_token="test-token") + create_kwargs = mock_client.create.call_args.kwargs + assert "auth_token" not in create_kwargs + assert create_kwargs["name"] == "test-sandbox" + assert create_kwargs["project_id"] == "proj-456" + assert create_kwargs["image"] == "custom-image:latest" + assert create_kwargs["cpu_cores"] == 4 + mock_sandbox.expose_port.assert_called_once_with(9999) + assert executor.port == 9999 + assert mock_install_packages.called + + @patch.dict("os.environ", {"TENKI_PROJECT_ID": "proj-123"}) + @patch("smolagents.remote_executors.TenkiExecutor._wait_for_server") + @patch("smolagents.remote_executors._create_kernel_http") + @patch("tenki_sandbox.Client") + def test_tenki_executor_cleanup(self, mock_client_cls, mock_create_kernel, mock_wait_for_server): + """Test TenkiExecutor cleanup method and double-cleanup guard.""" + logger = MagicMock() + mock_client, mock_sandbox = self._make_mock_client(mock_client_cls) + mock_create_kernel.return_value = "kernel-123" + + executor = TenkiExecutor(additional_imports=[], logger=logger) + executor.cleanup() + + assert mock_sandbox.close_if_open.call_count == 1 + assert mock_client.close.call_count == 1 + assert not hasattr(executor, "sandbox") + + # Second cleanup should be a no-op + executor.cleanup() + assert mock_sandbox.close_if_open.call_count == 1 + assert mock_client.close.call_count == 1 + + @patch.dict("os.environ", {"TENKI_PROJECT_ID": "proj-123"}) + @patch("smolagents.remote_executors.TenkiExecutor._wait_for_server") + @patch("smolagents.remote_executors._create_kernel_http") + @patch("tenki_sandbox.Client") + def test_tenki_executor_cleanup_retries_after_transient_failure( + self, mock_client_cls, mock_create_kernel, mock_wait_for_server + ): + """A transient close_if_open failure must not be marked as cleaned up: the handle is retained and retried.""" + logger = MagicMock() + mock_client, mock_sandbox = self._make_mock_client(mock_client_cls) + mock_create_kernel.return_value = "kernel-123" + mock_sandbox.close_if_open.side_effect = [RuntimeError("transient"), None] + + executor = TenkiExecutor(additional_imports=[], logger=logger) + + # First cleanup: termination fails, handle + client retained for a retry + executor.cleanup() + assert mock_sandbox.close_if_open.call_count == 1 + assert mock_client.close.call_count == 0 + assert hasattr(executor, "sandbox") + assert executor._cleaned_up is False + + # Second cleanup: termination succeeds, resources released + executor.cleanup() + assert mock_sandbox.close_if_open.call_count == 2 + assert mock_client.close.call_count == 1 + assert not hasattr(executor, "sandbox") + assert executor._cleaned_up is True + + @patch.dict("os.environ", {"TENKI_PROJECT_ID": "proj-123"}) + @patch("smolagents.remote_executors.TenkiExecutor._wait_for_server") + @patch("smolagents.remote_executors._create_kernel_http") + @patch("tenki_sandbox.Client") + def test_tenki_executor_cleanup_on_init_failure(self, mock_client_cls, mock_create_kernel, mock_wait_for_server): + """Test that the sandbox and client are cleaned up when initialization fails.""" + logger = MagicMock() + mock_client, mock_sandbox = self._make_mock_client(mock_client_cls) + mock_create_kernel.side_effect = RuntimeError("Failed to create kernel") + + with pytest.raises(RuntimeError, match="Failed to initialize Tenki sandbox"): + TenkiExecutor(additional_imports=[], logger=logger) + + assert mock_sandbox.close_if_open.call_count == 1 + assert mock_client.close.call_count == 1 + + @patch.dict("os.environ", {"TENKI_PROJECT_ID": "proj-123"}) + @patch("smolagents.remote_executors.TenkiExecutor._wait_for_server") + @patch("smolagents.remote_executors._create_kernel_http") + @patch("tenki_sandbox.Client") + def test_tenki_executor_cleanup_on_keyboard_interrupt( + self, mock_client_cls, mock_create_kernel, mock_wait_for_server + ): + """A KeyboardInterrupt during init (BaseException, not Exception) must still clean up and re-raise unchanged.""" + logger = MagicMock() + mock_client, mock_sandbox = self._make_mock_client(mock_client_cls) + mock_create_kernel.side_effect = KeyboardInterrupt() + + with pytest.raises(KeyboardInterrupt): + TenkiExecutor(additional_imports=[], logger=logger) + + assert mock_sandbox.close_if_open.call_count == 1 + assert mock_client.close.call_count == 1 + + @patch.dict("os.environ", {"TENKI_PROJECT_ID": "proj-123"}) + @patch("smolagents.remote_executors.TenkiExecutor._wait_for_server") + @patch("smolagents.remote_executors._create_kernel_http") + @patch("tenki_sandbox.Client") + def test_tenki_executor_bootstrap_failure_surfaces_stderr( + self, mock_client_cls, mock_create_kernel, mock_wait_for_server + ): + """Test that a failed kernel gateway bootstrap reports the command's stderr.""" + logger = MagicMock() + mock_client, mock_sandbox = self._make_mock_client(mock_client_cls) + mock_sandbox.shell.return_value = MagicMock(ok=False, exit_code=127, stderr_text="apt-get: not found") + + with pytest.raises(RuntimeError, match="apt-get: not found"): + TenkiExecutor(additional_imports=[], logger=logger) + + assert mock_sandbox.close_if_open.call_count == 1 + + @patch("smolagents.remote_executors.TenkiExecutor._wait_for_server") + @patch("smolagents.remote_executors._create_kernel_http") + @patch("tenki_sandbox.Client") + def test_tenki_executor_project_id_auto_resolution( + self, mock_client_cls, mock_create_kernel, mock_wait_for_server + ): + """Test that a sole project is auto-resolved when TENKI_PROJECT_ID is not set.""" + logger = MagicMock() + mock_client, mock_sandbox = self._make_mock_client(mock_client_cls) + mock_create_kernel.return_value = "kernel-123" + mock_client.who_am_i.return_value = MagicMock(workspaces=[MagicMock(projects=[MagicMock(id="proj-auto")])]) + + with patch.dict("os.environ", {}, clear=True): + executor = TenkiExecutor(additional_imports=[], logger=logger) + + assert mock_client.create.call_args.kwargs["project_id"] == "proj-auto" + assert isinstance(executor, TenkiExecutor) + + @patch("smolagents.remote_executors.TenkiExecutor._wait_for_server") + @patch("smolagents.remote_executors._create_kernel_http") + @patch("tenki_sandbox.Client") + def test_tenki_executor_project_id_workspace_filter( + self, mock_client_cls, mock_create_kernel, mock_wait_for_server + ): + """Test that a workspace_id in create_kwargs disambiguates project resolution.""" + logger = MagicMock() + mock_client, mock_sandbox = self._make_mock_client(mock_client_cls) + mock_create_kernel.return_value = "kernel-123" + mock_client.who_am_i.return_value = MagicMock( + workspaces=[ + MagicMock(id="ws-1", projects=[MagicMock(id="proj-1")]), + MagicMock(id="ws-2", projects=[MagicMock(id="proj-2")]), + ] + ) + + with patch.dict("os.environ", {}, clear=True): + TenkiExecutor(additional_imports=[], logger=logger, create_kwargs={"workspace_id": "ws-2"}) + + create_kwargs = mock_client.create.call_args.kwargs + assert create_kwargs["project_id"] == "proj-2" + assert create_kwargs["workspace_id"] == "ws-2" + + @patch("tenki_sandbox.Client") + def test_tenki_executor_project_id_ambiguous(self, mock_client_cls): + """Test that an ambiguous project resolution raises a helpful error without creating a sandbox.""" + logger = MagicMock() + mock_client = mock_client_cls.return_value + mock_client.who_am_i.return_value = MagicMock( + workspaces=[MagicMock(projects=[MagicMock(id="proj-1"), MagicMock(id="proj-2")])] + ) + + with patch.dict("os.environ", {}, clear=True): + with pytest.raises(ValueError, match="TENKI_PROJECT_ID"): + TenkiExecutor(additional_imports=[], logger=logger) + + mock_client.create.assert_not_called() + assert mock_client.close.call_count == 1