feat: restore sandbox tools and proactive compaction

agent/context_manager/manager.py

@@ -90,11 +90,11 @@ class ContextManager:
         compact_size: float = 0.1,
         untouched_messages: int = 5,
         tool_specs: list[dict[str, Any]] | None = None,
-        prompt_file_suffix: str = "system_prompt_v2.yaml",
+        prompt_file_suffix: str = "system_prompt_v3.yaml",
     ):
         self.system_prompt = self._load_system_prompt(
             tool_specs or [],
-            prompt_file_suffix="system_prompt_v2.yaml",
+            prompt_file_suffix="system_prompt_v3.yaml",
         )
         self.max_context = max_context
         self.compact_size = int(max_context * compact_size)
@@ -144,7 +144,9 @@ class ContextManager:
         """Get all messages for sending to LLM"""
         return self.items
 
-    async def compact(self, model_name: str) -> None:
+    async def compact(
+        self, model_name: str, tool_specs: list[dict] | None = None
+    ) -> None:
         """Remove old messages to keep history under target size"""
         if (self.context_length <= self.max_context) or not self.items:
             return
@@ -179,6 +181,7 @@ class ContextManager:
             model=model_name,
             messages=messages_to_summarize,
             max_completion_tokens=self.compact_size,
+            tools=tool_specs,
             api_key=hf_key
             if hf_key and model_name.startswith("huggingface/")
             else None,

agent/core/agent_loop.py

@@ -8,6 +8,7 @@ import logging
 import os
 
 from litellm import ChatCompletionMessageToolCall, Message, acompletion
+from litellm.exceptions import ContextWindowExceededError
 from lmnr import observe
 
 from agent.config import Config
@@ -88,6 +89,9 @@ def _needs_approval(
     if not args_valid:
         return False
 
+    if tool_name == "sandbox_create":
+        return True
+
     if tool_name == "hf_jobs":
         operation = tool_args.get("operation", "")
         if operation not in ["run", "uv", "scheduled run", "scheduled uv"]:
@@ -142,6 +146,24 @@ def _needs_approval(
     return False
 
 
+async def _compact_and_notify(session: Session) -> None:
+    """Run compaction and send event if context was reduced."""
+    old_length = session.context_manager.context_length
+    tool_specs = session.tool_router.get_tool_specs_for_llm()
+    await session.context_manager.compact(
+        model_name=session.config.model_name,
+        tool_specs=tool_specs,
+    )
+    new_length = session.context_manager.context_length
+    if new_length != old_length:
+        await session.send_event(
+            Event(
+                event_type="compacted",
+                data={"old_tokens": old_length, "new_tokens": new_length},
+            )
+        )
+
+
 class Handlers:
     """Handler functions for each operation type"""
 
@@ -184,7 +206,7 @@ class Handlers:
     @staticmethod
     @observe(name="run_agent")
     async def run_agent(
-        session: Session, text: str, max_iterations: int = 10
+        session: Session, text: str, max_iterations: int = 300
     ) -> str | None:
         """
         Handle user input (like user_input_or_turn in codex.rs:1291)
@@ -216,6 +238,9 @@ class Handlers:
         final_response = None
 
         while iteration < max_iterations:
+            # Compact before calling the LLM if context is near the limit
+            await _compact_and_notify(session)
+
             messages = session.context_manager.get_messages()
             tools = session.tool_router.get_tool_specs_for_llm()
             try:
@@ -449,6 +474,14 @@ class Handlers:
 
                 iteration += 1
 
+            except ContextWindowExceededError:
+                # Force compact and retry this iteration
+                session.context_manager.context_length = (
+                    session.context_manager.max_context + 1
+                )
+                await _compact_and_notify(session)
+                continue
+
             except Exception as e:
                 import traceback
 
@@ -460,18 +493,6 @@ class Handlers:
                 )
                 break
 
-        old_length = session.context_manager.context_length
-        await session.context_manager.compact(model_name=session.config.model_name)
-        new_length = session.context_manager.context_length
-
-        if new_length != old_length:
-            await session.send_event(
-                Event(
-                    event_type="compacted",
-                    data={"old_tokens": old_length, "new_tokens": new_length},
-                )
-            )
-
         await session.send_event(
             Event(
                 event_type="turn_complete",
@@ -491,20 +512,6 @@ class Handlers:
         session.interrupt()
         await session.send_event(Event(event_type="interrupted"))
 
-    @staticmethod
-    async def compact(session: Session) -> None:
-        """Handle compact (like compact in codex.rs:1317)"""
-        old_length = session.context_manager.context_length
-        await session.context_manager.compact(model_name=session.config.model_name)
-        new_length = session.context_manager.context_length
-
-        await session.send_event(
-            Event(
-                event_type="compacted",
-                data={"removed": old_length, "remaining": new_length},
-            )
-        )
-
     @staticmethod
     async def undo(session: Session) -> None:
         """Remove the last complete turn (user msg + all assistant/tool msgs that follow).
@@ -769,7 +776,7 @@ async def process_submission(session: Session, submission) -> bool:
         return True
 
     if op.op_type == OpType.COMPACT:
-        await Handlers.compact(session)
+        await _compact_and_notify(session)
         return True
 
     if op.op_type == OpType.UNDO:

agent/core/session.py

@@ -99,6 +99,7 @@ class Session:
         self.pending_approval: Optional[dict[str, Any]] = None
         # User's HF OAuth token — set by session_manager after construction
         self.hf_token: Optional[str] = None
+        self.sandbox = None
 
         # Session trajectory logging
         self.logged_events: list[dict] = []

agent/core/tools.py

@@ -48,6 +48,7 @@ from agent.tools.hf_repo_git_tool import (
 )
 from agent.tools.jobs_tool import HF_JOBS_TOOL_SPEC, hf_jobs_handler
 from agent.tools.plan_tool import PLAN_TOOL_SPEC, plan_tool_handler
+from agent.tools.sandbox_tool import get_sandbox_tools
 
 # NOTE: Private HF repo tool disabled - replaced by hf_repo_files and hf_repo_git
 # from agent.tools.private_hf_repo_tools import (
@@ -334,6 +335,9 @@ def create_builtin_tools() -> list[ToolSpec]:
         ),
     ]
 
+    # Sandbox tools (highest priority)
+    tools = get_sandbox_tools() + tools
+
     tool_names = ", ".join([t.name for t in tools])
     logger.info(f"Loaded {len(tools)} built-in tools: {tool_names}")
 

agent/prompts/system_prompt_v3.yaml

@@ -0,0 +1,118 @@
+system_prompt: |
+  You are Hugging Face Agent, an ML engineering assistant with {{ num_tools }} tools for training, fine-tuning, data processing, inference, and evaluation on the Hugging Face ecosystem.
+
+  _Current Time: **{{ current_date }} {{ current_time }} ({{ current_timezone }})**_
+  {% if hf_user_info %}_Authenticated as: **{{ hf_user_info }}**_{% endif %}
+
+  Your goal is to complete what the user requested with zero errors. You are fully autonomous — research, validate, implement, and deliver results without asking for unnecessary confirmation.
+
+  # Your knowledge of HF libraries is outdated
+
+  You do not know current APIs for TRL, Transformers, PEFT, Trackio, or other HF libraries. Your internal knowledge WILL produce wrong imports, wrong argument names, and wrong trainer configurations.
+
+  Before writing any ML implementation code (training, fine-tuning, inference, data processing), ground yourself in current working code:
+
+    github_find_examples → github_read_file → explore_hf_docs + fetch_hf_docs
+
+  Skip research only for trivial non-code operations.
+
+  # Mistakes you WILL make without research
+
+  HALLUCINATED IMPORTS: You will import from modules that were renamed or removed. Example: old TRL trainer class names, deprecated Transformers APIs, wrong trackio parameter names (e.g. `run_name` instead of `name`). Fix: read a current example script first.
+
+  WRONG TRAINER ARGUMENTS: You will pass configuration arguments that don't exist in current trainer versions. Fix: fetch the actual trainer/config docs via explore_hf_docs + fetch_hf_docs.
+
+  WRONG DATASET FORMAT: You will assume column names without checking. Training fails with KeyError. Fix: call hf_inspect_dataset or hub_repo_details and verify columns match the training method.
+
+  DEFAULT TIMEOUT KILLS JOBS: You will leave timeout at the default 30m for training jobs. Training takes hours. The job gets killed and all progress is lost. Fix: set timeout based on model size (minimum 2h for any training).
+
+  LOST MODELS: You will forget push_to_hub=True and hub_model_id in training config. Job storage is ephemeral — the filesystem is deleted when the job ends. Without push_to_hub, the trained model is permanently lost.
+
+  BATCH FAILURES: You will submit all ablation/batch jobs at once without testing that one works first. All will fail for the same bug. Fix: submit ONE job first, verify it completes successfully, then submit the rest.
+
+  SILENT DATASET SUBSTITUTION: When a requested dataset fails to load, you will silently switch to a different one without telling the user. Fix: if the requested dataset isn't available, tell the user and ask what to do.
+
+  HARDCODED UNAVAILABLE PACKAGES: You will forget to install necessary packages like 'flash-attn' for flash_attention_2 or other packages that aren't automatically installed in the job environment. Fix: install necessary packages before running the job.
+
+  SCOPE-CHANGING FIXES: Avoid at all costs! When you hit an error (especially OOM), you will try "creative" workarounds that change what the user asked for and/or change the training task itself — switching full SFT to LoRA on OOM, reducing max_length (silently truncates training data and changes what the model learns), disabling monitoring instead of fixing it. Do not do this. Fix errors with the minimal change that preserves the user's original request and are grounded in research and examples. If the original approach genuinely cannot work, explain why and ask the user for input before changing methods, sequence length, training approach or any other part of the task.
+
+  # When writing ML code
+
+  Required sequence before any training/fine-tuning/inference script:
+  1. Find working examples: github_find_examples (discover) → github_read_file (study)
+  2. Check documentation: explore_hf_docs + fetch_hf_docs for trainer configs and parameters
+  3. Validate dataset details: hf_inspect_dataset to confirm column names and format.
+  4. Validate model details: hub_repo_details to confirm model exists, it's the correct architecture/size/tokenizer etc.
+
+  Dataset format requirements by training method:
+    SFT: "messages", "text", or "prompt"/"completion"
+    DPO: "prompt", "chosen", "rejected"
+    GRPO: "prompt"
+
+  # When submitting a training job
+
+  Before calling hf_jobs, output a pre-flight check:
+    - Reference implementation: [which example you based this on]
+    - Dataset format verified: [columns confirmed via hf_inspect_dataset/hub_repo_details]
+    - push_to_hub=True and hub_model_id set
+    - timeout: [value] (based on: [model size] on [hardware])
+    - Trackio monitoring included and working
+
+  If you cannot fill in all items, stop and complete the missing steps first.
+
+  For batch/ablation jobs: submit ONE job first. Check logs to confirm it starts training successfully. Only then submit the remaining jobs. Never submit all at once.
+
+  Hardware sizing:
+    1-3B params: a10g-largex2
+    7-13B params: a100-large
+    30B+ params: l40sx4 or a100x4
+    70B+ params: a100x8
+  Note: a10g-small and a10g-large have the SAME 24GB GPU memory. The difference is CPU/RAM only.
+
+  # Sandbox-first development
+
+  For non-trivial scripts, develop and test in a sandbox before launching via hf_jobs:
+    sandbox_create → install deps → write script → test with small run → fix errors → launch via hf_jobs at scale
+
+  Use GPU sandbox (t4-small minimum) when testing code that uses CUDA, bf16, or model loading. CPU sandboxes cannot test GPU code paths.
+
+
+  # When a task has 3+ steps
+
+  Use plan_tool to track progress. One task in_progress at a time. Mark completed immediately after finishing. Update frequently to show the user what you're doing.
+
+  # Error recovery
+
+  When something fails:
+  - Diagnose the actual error. Read the full error message and logs.
+  - Do not retry the exact same thing. Identify what needs to change.
+  - If an API/import error: check documentation for the correct API.
+  - If an OOM error: (1) reduce per_device_train_batch_size and increase gradient_accumulation_steps proportionally to keep effective batch size identical, (2) enable gradient_checkpointing=True, (3) upgrade to larger GPU (a10gx4→a100→a100x4→a100x8). Do NOT switch training methods (e.g. SFT→LoRA) or reduce max_length — those change what the user gets. If OOM happens in sandbox, create a new sandbox with larger GPU hardware.
+  - Never change the user's requested approach (training method, dataset, model, sequence length) without explicit approval.
+  - If a tool call fails repeatedly for the same reason: stop and try a different approach.
+  - Never silently substitute resources (datasets, models) — tell the user if something isn't available.
+
+  # Task completion
+
+  Before ending your turn, verify:
+  - Did you actually DO what the user asked, not just explain what you would do?
+  - If something failed: did you diagnose and fix it, or at minimum explain what went wrong and ask for user input?
+  - For training jobs: did you include a working Trackio dashboard URL?
+
+  Do not stop after describing what you plan to do. Continue calling tools until the task is verifiably done.
+  Do not mark plan tasks as completed if they failed or are only partially done.
+
+  # Communication
+
+  - Be concise and direct. No filler, no restating what the user said.
+  - One-word answers when appropriate for simple questions.
+  - Always include direct Hub URLs when referencing models, datasets, Spaces, or jobs.
+  - For errors: state what went wrong, why, and what you're doing to fix it.
+  - Do not over-explain or present elaborate option menus for simple tasks. When the user's intent is clear, act on it. Present options only when there's genuine ambiguity.
+
+  # Tool usage
+
+  - Execute multiple independent tool calls in parallel when possible.
+  - HF_TOKEN is automatically available in job secrets — no need to include it extra.
+  - For training monitoring: include Trackio in the script and provide the dashboard URL.
+  - For private/gated datasets: HF_TOKEN is needed — it's auto-loaded into job secrets.

agent/tools/jobs_tool.py

@@ -979,7 +979,11 @@ HF_JOBS_TOOL_SPEC = {
             # Python/UV specific parameters
             "script": {
                 "type": "string",
-                "description": "Python code to execute. Triggers Python mode (auto pip install). Use with 'run'/'scheduled run'. Mutually exclusive with 'command'.",
+                "description": (
+                    "Python code or sandbox file path (e.g. '/app/train.py') or URL. "
+                    "Triggers Python mode. For ML training: base this on a working example found via github_find_examples, not on internal knowledge. "
+                    "Mutually exclusive with 'command'."
+                ),
             },
             "dependencies": {
                 "type": "array",
@@ -1041,6 +1045,28 @@ async def hf_jobs_handler(
                     Event(event_type="tool_log", data={"tool": "hf_jobs", "log": log})
                 )
 
+        # If script is a sandbox file path, read it from the sandbox
+        script = arguments.get("script", "")
+        sandbox = getattr(session, "sandbox", None) if session else None
+        is_path = (
+            sandbox
+            and isinstance(script, str)
+            and script.strip() == script
+            and not any(c in script for c in "\r\n\0")
+            and (
+                script.startswith("/")
+                or script.startswith("./")
+                or script.startswith("../")
+            )
+        )
+        if is_path:
+            import shlex
+
+            result = await asyncio.to_thread(sandbox.bash, f"cat {shlex.quote(script)}")
+            if not result.success:
+                return f"Failed to read {script} from sandbox: {result.error}", False
+            arguments = {**arguments, "script": result.output}
+
         # Prefer the authenticated user's OAuth token, fall back to global env
         hf_token = (
             (getattr(session, "hf_token", None) if session else None)

agent/tools/sandbox_client.py

@@ -0,0 +1,714 @@
+#!/usr/bin/env python3
+# /// script
+# requires-python = ">=3.10"
+# dependencies = ["huggingface_hub>=0.20.0", "httpx>=0.27.0"]
+# ///
+"""
+Sandbox Tools — Agent-native primitives for HF Space dev-mode sandboxes.
+
+Architecture:
+  - Creates a sandbox by duplicating a template Space (runs sandbox_server.py)
+  - Waits for it to come online
+  - Communicates via HTTPS to the Space's API
+  - Optionally deletes the Space when done
+
+Lifecycle:
+    sb = Sandbox.create(owner="burtenshaw")         # duplicate, wait, connect
+    sb = Sandbox.create(owner="burtenshaw",          # with options
+                        hardware="t4-small",
+                        private=True,
+                        sleep_time=3600)
+    sb = Sandbox.connect("burtenshaw/my-sandbox-abc") # attach to existing
+
+    sb.bash("uv run train.py")
+    sb.read("/app/train.py")
+    sb.edit("/app/train.py", old_str="lr=1e-3", new_str="lr=1e-4")
+
+    sb.delete()                                       # tear down when done
+
+    # Or use as a context manager for automatic cleanup
+    with Sandbox.create(owner="burtenshaw") as sb:
+        sb.bash("python train.py")
+    # Space deleted on exit
+
+Tools: bash, read, write, edit, upload
+"""
+
+from __future__ import annotations
+
+import io
+import os
+import sys
+import time
+import uuid
+from dataclasses import dataclass, field
+from typing import Any
+
+import httpx
+from huggingface_hub import CommitOperationAdd, HfApi
+
+TEMPLATE_SPACE = "burtenshaw/sandbox"
+HARDWARE_OPTIONS = [
+    "cpu-basic",
+    "cpu-upgrade",
+    "t4-small",
+    "t4-medium",
+    "a10g-small",
+    "a10g-large",
+    "a100-large",
+]
+OUTPUT_LIMIT = 30000
+LINE_LIMIT = 2000
+DEFAULT_READ_LIMIT = 2000
+DEFAULT_TIMEOUT = 120
+MAX_TIMEOUT = 600
+WAIT_TIMEOUT = 300
+WAIT_INTERVAL = 5
+API_WAIT_TIMEOUT = 180
+
+_DOCKERFILE = """\
+FROM ghcr.io/astral-sh/uv:python3.12-bookworm-slim
+
+RUN apt-get update && \\
+    apt-get install -y \\
+      bash git git-lfs wget curl procps \\
+      htop vim nano jq tmux \\
+      build-essential && \\
+    rm -rf /var/lib/apt/lists/*
+
+RUN uv pip install --system fastapi uvicorn python-multipart
+
+RUN useradd -m -u 1000 user
+USER user
+
+ENV HOME=/home/user \\
+    PATH=/home/user/.local/bin:$PATH \\
+    PIP_USER=1 \\
+    HF_HUB_DISABLE_PROGRESS_BARS=1 \\
+    TQDM_DISABLE=1 \\
+    TRANSFORMERS_VERBOSITY=warning \\
+    HF_HUB_ENABLE_HF_TRANSFER=1
+
+WORKDIR /app
+COPY --chown=user . /app
+
+EXPOSE 7860
+
+CMD ["python", "sandbox_server.py"]
+"""
+
+_SANDBOX_SERVER = '''\
+"""Minimal FastAPI server for sandbox operations."""
+import os, subprocess, pathlib
+from fastapi import FastAPI
+from pydantic import BaseModel
+from typing import Optional
+import uvicorn
+
+app = FastAPI()
+
+class BashReq(BaseModel):
+    command: str
+    work_dir: str = "/app"
+    timeout: int = 120
+
+class ReadReq(BaseModel):
+    path: str
+    offset: Optional[int] = None
+    limit: Optional[int] = 2000
+
+class WriteReq(BaseModel):
+    path: str
+    content: str
+
+class EditReq(BaseModel):
+    path: str
+    old_str: str
+    new_str: str
+    replace_all: bool = False
+
+class ExistsReq(BaseModel):
+    path: str
+
+@app.get("/api/health")
+def health():
+    return {"status": "ok"}
+
+@app.post("/api/bash")
+def bash(req: BashReq):
+    try:
+        r = subprocess.run(
+            req.command, shell=True, capture_output=True, text=True,
+            cwd=req.work_dir, timeout=req.timeout,
+        )
+        output = r.stdout + r.stderr
+        if len(output) > 30000:
+            output = output[:30000] + "\\n... (truncated)"
+        return {"success": r.returncode == 0, "output": output, "error": "" if r.returncode == 0 else f"Exit code {r.returncode}"}
+    except subprocess.TimeoutExpired:
+        return {"success": False, "output": "", "error": f"Timeout after {req.timeout}s"}
+    except Exception as e:
+        return {"success": False, "output": "", "error": str(e)}
+
+@app.post("/api/read")
+def read(req: ReadReq):
+    try:
+        p = pathlib.Path(req.path)
+        if not p.exists():
+            return {"success": False, "output": "", "error": f"File not found: {req.path}"}
+        if p.is_dir():
+            return {"success": False, "output": "", "error": f"Is a directory: {req.path}"}
+        lines = p.read_text().splitlines()
+        start = (req.offset or 1) - 1
+        end = start + (req.limit or len(lines))
+        selected = lines[start:end]
+        numbered = "\\n".join(f"{start + i + 1}\\t{line}" for i, line in enumerate(selected))
+        return {"success": True, "output": numbered, "error": ""}
+    except Exception as e:
+        return {"success": False, "output": "", "error": str(e)}
+
+@app.post("/api/write")
+def write(req: WriteReq):
+    try:
+        p = pathlib.Path(req.path)
+        p.parent.mkdir(parents=True, exist_ok=True)
+        p.write_text(req.content)
+        return {"success": True, "output": f"Wrote {len(req.content)} bytes to {req.path}", "error": ""}
+    except Exception as e:
+        return {"success": False, "output": "", "error": str(e)}
+
+@app.post("/api/edit")
+def edit(req: EditReq):
+    try:
+        p = pathlib.Path(req.path)
+        if not p.exists():
+            return {"success": False, "output": "", "error": f"File not found: {req.path}"}
+        content = p.read_text()
+        if req.old_str not in content:
+            return {"success": False, "output": "", "error": f"old_str not found in {req.path}"}
+        if not req.replace_all and content.count(req.old_str) > 1:
+            return {"success": False, "output": "", "error": f"old_str appears {content.count(req.old_str)} times. Use replace_all=true or provide more context."}
+        if req.replace_all:
+            new_content = content.replace(req.old_str, req.new_str)
+        else:
+            new_content = content.replace(req.old_str, req.new_str, 1)
+        p.write_text(new_content)
+        return {"success": True, "output": f"Edited {req.path}", "error": ""}
+    except Exception as e:
+        return {"success": False, "output": "", "error": str(e)}
+
+@app.post("/api/exists")
+def exists(req: ExistsReq):
+    return {"success": True, "output": str(pathlib.Path(req.path).exists()).lower(), "error": ""}
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=7860)
+'''
+
+
+@dataclass
+class ToolResult:
+    success: bool
+    output: str = ""
+    error: str = ""
+
+    def __str__(self):
+        if self.success:
+            return self.output or "(no output)"
+        return f"ERROR: {self.error}"
+
+    def to_dict(self) -> dict:
+        return {"success": self.success, "output": self.output, "error": self.error}
+
+
+@dataclass
+class Sandbox:
+    """
+    A handle to an HF Space sandbox.
+
+    Use Sandbox.create() to spin up a new one, or Sandbox.connect() to
+    attach to an existing running Space.
+    """
+
+    space_id: str
+    token: str | None = None
+    work_dir: str = "/app"
+    timeout: int = DEFAULT_TIMEOUT
+    _owns_space: bool = field(default=False, repr=False)
+    _base_url: str = field(init=False, repr=False)
+    _client: httpx.Client = field(init=False, repr=False)
+    _hf_api: HfApi = field(init=False, repr=False)
+    _files_read: set = field(init=False, repr=False, default_factory=set)
+
+    def __post_init__(self):
+        self.token = self.token or os.environ.get("HF_TOKEN")
+        slug = self.space_id.replace("/", "-")
+        # Trailing slash is critical: httpx resolves relative paths against base_url.
+        # Without it, client.get("health") resolves to /health instead of /api/health.
+        self._base_url = f"https://{slug}.hf.space/api/"
+        self._client = httpx.Client(
+            base_url=self._base_url,
+            headers={"Authorization": f"Bearer {self.token}"} if self.token else {},
+            timeout=httpx.Timeout(MAX_TIMEOUT, connect=30),
+            follow_redirects=True,
+        )
+        self._hf_api = HfApi(token=self.token)
+
+    # ── Lifecycle ─────────────────────────────────────────────────
+
+    @classmethod
+    def create(
+        cls,
+        owner: str,
+        *,
+        name: str | None = None,
+        template: str = TEMPLATE_SPACE,
+        hardware: str = "cpu-basic",
+        private: bool = False,
+        sleep_time: int | None = None,
+        token: str | None = None,
+        wait_timeout: int = WAIT_TIMEOUT,
+    ) -> Sandbox:
+        """
+        Create a new sandbox by duplicating the template Space.
+
+        Generates a unique space name, duplicates the template, waits for it
+        to come online, then returns a connected Sandbox.
+
+        Args:
+            owner: HF username or org (e.g. "burtenshaw").
+            name: Base name for the space. Defaults to "sandbox".
+                  A unique suffix is always appended.
+            template: Source Space to duplicate (default: burtenshaw/sandbox).
+            hardware: Hardware tier (cpu-basic, t4-small, etc.).
+            private: Whether the Space should be private.
+            sleep_time: Auto-sleep after N seconds of inactivity.
+            token: HF API token. Falls back to HF_TOKEN env var.
+            wait_timeout: Max seconds to wait for Space to start (default: 300).
+
+        Returns:
+            A Sandbox instance connected to the running Space.
+        """
+        token = token or os.environ.get("HF_TOKEN")
+        api = HfApi(token=token)
+
+        base = name or "sandbox"
+        suffix = uuid.uuid4().hex[:8]
+        space_id = f"{owner}/{base}-{suffix}"
+
+        print(f"Creating sandbox: {space_id} (from {template})...")
+
+        kwargs = {
+            "from_id": template,
+            "to_id": space_id,
+            "private": private,
+            "hardware": hardware,
+        }
+        if sleep_time is not None:
+            kwargs["sleep_time"] = sleep_time
+
+        api.duplicate_space(**kwargs)
+        print(f"Space created: https://huggingface.co/spaces/{space_id}")
+
+        # Upload sandbox server and Dockerfile (triggers rebuild)
+        cls._setup_server(space_id, api)
+
+        # Wait for it to come online (rebuild + start)
+        print(f"Waiting for Space to start (timeout: {wait_timeout}s)...")
+        deadline = time.time() + wait_timeout
+        while time.time() < deadline:
+            runtime = api.get_space_runtime(space_id)
+            if runtime.stage == "RUNNING":
+                print(f"Space is running (hardware: {runtime.hardware})")
+                break
+            if runtime.stage in ("RUNTIME_ERROR", "BUILD_ERROR"):
+                raise RuntimeError(
+                    f"Space failed to start: {runtime.stage}. "
+                    f"Check https://huggingface.co/spaces/{space_id}"
+                )
+            print(f"  {runtime.stage}...")
+            time.sleep(WAIT_INTERVAL)
+        else:
+            raise TimeoutError(
+                f"Space did not start within {wait_timeout}s. "
+                f"Check https://huggingface.co/spaces/{space_id}"
+            )
+
+        # Wait for the API server to be responsive (non-fatal)
+        sb = cls(space_id=space_id, token=token, _owns_space=True)
+        try:
+            sb._wait_for_api(timeout=API_WAIT_TIMEOUT)
+        except TimeoutError as e:
+            print(
+                f"Warning: API health check timed out ({e}), but Space is RUNNING. Continuing."
+            )
+        return sb
+
+    @staticmethod
+    def _setup_server(space_id: str, api: HfApi) -> None:
+        """Upload embedded sandbox server + Dockerfile to the Space (single commit)."""
+        print(f"Uploading sandbox server to {space_id}...")
+        api.create_commit(
+            repo_id=space_id,
+            repo_type="space",
+            operations=[
+                CommitOperationAdd(
+                    path_in_repo="sandbox_server.py",
+                    path_or_fileobj=io.BytesIO(_SANDBOX_SERVER.encode()),
+                ),
+                CommitOperationAdd(
+                    path_in_repo="Dockerfile",
+                    path_or_fileobj=io.BytesIO(_DOCKERFILE.encode()),
+                ),
+            ],
+            commit_message="Setup sandbox server",
+        )
+        print("Server files uploaded, rebuild triggered.")
+
+    @classmethod
+    def connect(cls, space_id: str, *, token: str | None = None) -> Sandbox:
+        """
+        Connect to an existing running Space.
+
+        Does a health check to verify the Space is reachable.
+        """
+        sb = cls(space_id=space_id, token=token, _owns_space=False)
+        sb._wait_for_api(timeout=60)
+        return sb
+
+    def _wait_for_api(self, timeout: int = API_WAIT_TIMEOUT):
+        """Poll the health endpoint until the server responds."""
+        deadline = time.time() + timeout
+        last_err = None
+        last_status = None
+        while time.time() < deadline:
+            try:
+                resp = self._client.get("health", timeout=10)
+                last_status = resp.status_code
+                if resp.status_code == 200:
+                    print(f"API is responsive at {self._base_url}")
+                    return
+            except Exception as e:
+                last_err = e
+            time.sleep(3)
+        raise TimeoutError(
+            f"Sandbox API at {self._base_url} not responding after {timeout}s. "
+            f"Last status: {last_status}, last error: {last_err}"
+        )
+
+    def delete(self):
+        """Delete the Space. Only works if this Sandbox created it."""
+        if not self._owns_space:
+            raise RuntimeError(
+                f"This Sandbox did not create {self.space_id}. "
+                f"Use self._hf_api.delete_repo() directly if you're sure."
+            )
+        print(f"Deleting sandbox: {self.space_id}...")
+        self._hf_api.delete_repo(self.space_id, repo_type="space")
+        self._client.close()
+        print("Deleted.")
+
+    def pause(self):
+        """Pause the Space (stops billing, preserves state)."""
+        self._hf_api.pause_space(self.space_id)
+
+    def restart(self):
+        """Restart the Space."""
+        self._hf_api.restart_space(self.space_id)
+        self._wait_for_api()
+
+    @property
+    def url(self) -> str:
+        """Public URL of the Space."""
+        return f"https://huggingface.co/spaces/{self.space_id}"
+
+    @property
+    def status(self) -> str:
+        """Current Space stage (RUNNING, BUILDING, PAUSED, etc.)."""
+        return self._hf_api.get_space_runtime(self.space_id).stage
+
+    def __enter__(self) -> Sandbox:
+        return self
+
+    def __exit__(self, *exc):
+        if self._owns_space:
+            try:
+                self.delete()
+            except Exception as e:
+                print(f"Warning: failed to delete sandbox: {e}", file=sys.stderr)
+        self._client.close()
+
+    # ── HTTP plumbing ─────────────────────────────────────────────
+
+    def _call(
+        self, endpoint: str, payload: dict, timeout: float | None = None
+    ) -> ToolResult:
+        # Strip leading slash for correct httpx base_url resolution
+        endpoint = endpoint.lstrip("/")
+        try:
+            resp = self._client.post(
+                endpoint,
+                json=payload,
+                timeout=timeout or self.timeout,
+            )
+            data = resp.json()
+            if resp.status_code == 200:
+                return ToolResult(
+                    success=data.get("success", True),
+                    output=data.get("output", ""),
+                    error=data.get("error", ""),
+                )
+            return ToolResult(
+                success=False,
+                error=data.get("error", f"HTTP {resp.status_code}"),
+            )
+        except httpx.TimeoutException:
+            return ToolResult(
+                success=False, error=f"Timeout after {timeout or self.timeout}s"
+            )
+        except httpx.ConnectError:
+            return ToolResult(
+                success=False,
+                error=f"Cannot connect to sandbox. Is {self.space_id} running? Status: {self.status}",
+            )
+        except Exception as e:
+            return ToolResult(success=False, error=str(e))
+
+    # ── Tools ─────────────────────────────────────────────────────
+
+    def bash(
+        self,
+        command: str,
+        *,
+        work_dir: str | None = None,
+        timeout: int | None = None,
+        description: str | None = None,
+    ) -> ToolResult:
+        return self._call(
+            "bash",
+            {
+                "command": command,
+                "work_dir": work_dir or self.work_dir,
+                "timeout": min(timeout or self.timeout, MAX_TIMEOUT),
+            },
+            timeout=timeout,
+        )
+
+    def read(
+        self, path: str, *, offset: int | None = None, limit: int | None = None
+    ) -> ToolResult:
+        self._files_read.add(path)
+        return self._call(
+            "read",
+            {
+                "path": path,
+                "offset": offset,
+                "limit": limit or (DEFAULT_READ_LIMIT if offset is None else None),
+            },
+        )
+
+    def write(self, path: str, content: str) -> ToolResult:
+        if path not in self._files_read:
+            check = self._call("exists", {"path": path})
+            if check.success and check.output == "true":
+                return ToolResult(
+                    success=False,
+                    error=(
+                        f"File {path} exists but has not been read this session. "
+                        f"Read it first, or use sandbox_edit for targeted changes."
+                    ),
+                )
+        result = self._call("write", {"path": path, "content": content})
+        if result.success:
+            self._files_read.add(path)
+        return result
+
+    def edit(
+        self, path: str, old_str: str, new_str: str, *, replace_all: bool = False
+    ) -> ToolResult:
+        if old_str == new_str:
+            return ToolResult(success=False, error="old_str and new_str are identical.")
+        if path not in self._files_read:
+            return ToolResult(
+                success=False,
+                error=f"File {path} has not been read this session. Read it first.",
+            )
+        return self._call(
+            "edit",
+            {
+                "path": path,
+                "old_str": old_str,
+                "new_str": new_str,
+                "replace_all": replace_all,
+            },
+        )
+
+    # ── Tool schemas & dispatch ───────────────────────────────────
+
+    TOOLS = {
+        "bash": {
+            "description": (
+                "Run a shell command in the remote sandbox and return stdout/stderr.\n"
+                "\n"
+                "Commands run in a shell at the working directory (default /app). "
+                "Each invocation is independent — use files in /app to persist state.\n"
+                "\n"
+                "AVOID using bash for operations covered by specialized tools:\n"
+                "- File reading: use read (not cat/head/tail)\n"
+                "- File editing: use edit (not sed/awk)\n"
+                "- File writing: use write (not echo/cat <<EOF)\n"
+                "\n"
+                "For long-running tasks, background them:\n"
+                "  nohup uv run train.py > /app/train.log 2>&1 &\n"
+                "Then check with read on the log file.\n"
+                "\n"
+                "Chain dependent commands with &&. Independent commands should be "
+                "separate bash calls (they can run in parallel).\n"
+                "\n"
+                "Timeout default 120s, max 600s."
+            ),
+            "parameters": {
+                "type": "object",
+                "required": ["command"],
+                "additionalProperties": False,
+                "properties": {
+                    "command": {
+                        "type": "string",
+                        "description": "The shell command to execute.",
+                    },
+                    "description": {
+                        "type": "string",
+                        "description": "Short description (5-10 words, active voice). E.g. 'Install dependencies', 'Run training script'.",
+                    },
+                    "work_dir": {
+                        "type": "string",
+                        "description": "Working directory (default: /app).",
+                    },
+                    "timeout": {
+                        "type": "integer",
+                        "description": "Timeout in seconds (default: 120, max: 600).",
+                    },
+                },
+            },
+        },
+        "read": {
+            "description": (
+                "Read file contents with line numbers (cat -n format).\n"
+                "\n"
+                "Returns the first 2000 lines by default. For large files, use offset/limit "
+                "to read a specific range. Line numbers always match the original file.\n"
+                "\n"
+                "Lines longer than 2000 chars are truncated.\n"
+                "Cannot read directories — use bash with 'ls' instead."
+            ),
+            "parameters": {
+                "type": "object",
+                "required": ["path"],
+                "additionalProperties": False,
+                "properties": {
+                    "path": {
+                        "type": "string",
+                        "description": "Absolute path to the file to read.",
+                    },
+                    "offset": {
+                        "type": "integer",
+                        "description": "Start from this line (1-based). Only if file is too large.",
+                    },
+                    "limit": {
+                        "type": "integer",
+                        "description": "Number of lines to read. Only if file is too large.",
+                    },
+                },
+            },
+        },
+        "write": {
+            "description": (
+                "Create or overwrite a file. Creates parent directories as needed.\n"
+                "\n"
+                "For existing files, you MUST read the file first (system enforced). "
+                "Prefer edit for modifications."
+            ),
+            "parameters": {
+                "type": "object",
+                "required": ["path", "content"],
+                "additionalProperties": False,
+                "properties": {
+                    "path": {
+                        "type": "string",
+                        "description": "Absolute path to the file to write.",
+                    },
+                    "content": {
+                        "type": "string",
+                        "description": "Complete file content.",
+                    },
+                },
+            },
+        },
+        "edit": {
+            "description": (
+                "Targeted edit via exact string replacement.\n"
+                "\n"
+                "Rules:\n"
+                "- old_str must appear EXACTLY once (unless replace_all is true).\n"
+                "- Include enough context in old_str for uniqueness.\n"
+                "- old_str and new_str must differ.\n"
+                "- Preserve indentation exactly.\n"
+                "- To delete code, set new_str to empty string.\n"
+                "- File MUST have been read this session (system enforced).\n"
+                "- Do NOT include line number prefixes in old_str/new_str.\n"
+                "\n"
+                "Use replace_all=true for batch operations like variable renaming."
+            ),
+            "parameters": {
+                "type": "object",
+                "required": ["path", "old_str", "new_str"],
+                "additionalProperties": False,
+                "properties": {
+                    "path": {
+                        "type": "string",
+                        "description": "Absolute path to the file.",
+                    },
+                    "old_str": {
+                        "type": "string",
+                        "description": "Exact text to find (must differ from new_str).",
+                    },
+                    "new_str": {"type": "string", "description": "Replacement text."},
+                    "replace_all": {
+                        "type": "boolean",
+                        "description": "Replace all occurrences (default: false).",
+                        "default": False,
+                    },
+                },
+            },
+        },
+    }
+
+    @classmethod
+    def tool_definitions(cls) -> list[dict]:
+        return [{"name": name, **spec} for name, spec in cls.TOOLS.items()]
+
+    def call_tool(self, name: str, arguments: dict[str, Any]) -> ToolResult:
+        dispatch = {
+            "bash": lambda a: self.bash(
+                a["command"],
+                work_dir=a.get("work_dir"),
+                timeout=a.get("timeout"),
+                description=a.get("description"),
+            ),
+            "read": lambda a: self.read(
+                a["path"],
+                offset=a.get("offset"),
+                limit=a.get("limit"),
+            ),
+            "write": lambda a: self.write(a["path"], a["content"]),
+            "edit": lambda a: self.edit(
+                a["path"],
+                a["old_str"],
+                a["new_str"],
+                replace_all=a.get("replace_all", False),
+            ),
+        }
+        fn = dispatch.get(name)
+        if not fn:
+            return ToolResult(success=False, error=f"Unknown tool: {name}")
+        return fn(arguments)

agent/tools/sandbox_tool.py

@@ -0,0 +1,201 @@
+"""
+Sandbox tools — expose the Sandbox client as agent tools.
+
+5 tools total:
+  sandbox_create — explicit sandbox creation (requires approval)
+  bash, read, write, edit — operations on the sandbox
+
+If any operation tool is called without an active sandbox,
+a cpu-basic sandbox is auto-created (no approval needed).
+"""
+
+from __future__ import annotations
+
+import asyncio
+import os
+from typing import Any
+
+from huggingface_hub import HfApi, SpaceHardware
+
+from agent.core.session import Event
+from agent.tools.sandbox_client import Sandbox
+
+# ── Tool name mapping (short agent names → Sandbox client names) ──────
+
+
+async def _ensure_sandbox(
+    session: Any, hardware: str = "cpu-basic", **create_kwargs
+) -> tuple[Sandbox | None, str | None]:
+    """
+    Ensure a sandbox exists on the session. Auto-creates with given hardware if needed.
+
+    Returns:
+        (sandbox, error_message) — one will be None.
+    """
+    if session and getattr(session, "sandbox", None):
+        return session.sandbox, None
+
+    if not session:
+        return None, "No session available."
+
+    token = os.environ.get("HF_TOKEN")
+    if not token:
+        return None, "HF_TOKEN environment variable not set. Cannot create sandbox."
+
+    api = HfApi(token=token)
+    user_info = api.whoami()
+    owner = user_info.get("name", user_info.get("user", ""))
+    if not owner:
+        return None, "Could not determine HF username from token."
+
+    await session.send_event(
+        Event(
+            event_type="tool_log",
+            data={
+                "tool": "sandbox",
+                "log": f"Auto-creating sandbox for {owner} ({hardware})...",
+            },
+        )
+    )
+
+    kwargs = {"owner": owner, "hardware": hardware, "token": token, **create_kwargs}
+    sb = await asyncio.to_thread(Sandbox.create, **kwargs)
+    session.sandbox = sb
+
+    await session.send_event(
+        Event(
+            event_type="tool_log",
+            data={"tool": "sandbox", "log": f"Sandbox ready: {sb.space_id} ({sb.url})"},
+        )
+    )
+
+    return sb, None
+
+
+# ── sandbox_create tool ──────────────────────────────────────────────
+
+SANDBOX_CREATE_TOOL_SPEC = {
+    "name": "sandbox_create",
+    "description": (
+        "Create a persistent remote Linux environment for developing and testing scripts.\n\n"
+        "Workflow: sandbox_create → write script → pip install → test with small run → fix errors → hf_jobs at scale.\n"
+        "The sandbox persists across tool calls within the session. pip install works out of the box.\n\n"
+        "Use this when: you need to develop, test, and iterate on scripts before launching via hf_jobs. "
+        "Especially for training scripts where you need to verify imports, test on a small subset, and fix errors interactively.\n\n"
+        "Skip this when: the task is a simple one-shot operation (status check, resource search, quick data query), "
+        "or the script is copied from a verified working example with minimal changes.\n\n"
+        "For ML code that uses CUDA, bf16, or model loading: use GPU hardware (t4-small minimum). "
+        "CPU sandboxes cannot run GPU code paths — your test will not catch GPU-related errors.\n\n"
+        "Hardware: " + ", ".join([e.value for e in SpaceHardware]) + ".\n"
+    ),
+    "parameters": {
+        "type": "object",
+        "required": [],
+        "additionalProperties": False,
+        "properties": {
+            "hardware": {
+                "type": "string",
+                "enum": [e.value for e in SpaceHardware],
+                "description": "Hardware tier for the sandbox (default: cpu-basic)",
+            },
+            "private": {
+                "type": "boolean",
+                "description": "If true, create a private Space",
+            },
+        },
+    },
+}
+
+
+async def sandbox_create_handler(
+    args: dict[str, Any], session: Any = None
+) -> tuple[str, bool]:
+    """Handle sandbox_create tool calls."""
+    # If sandbox already exists, return its info
+    if session and getattr(session, "sandbox", None):
+        sb = session.sandbox
+        return (
+            f"Sandbox already active: {sb.space_id}\n"
+            f"URL: {sb.url}\n"
+            f"Use bash/read/write/edit to interact with it."
+        ), True
+
+    hardware = args.get("hardware", "cpu-basic")
+    create_kwargs = {}
+    if "private" in args:
+        create_kwargs["private"] = args["private"]
+
+    try:
+        sb, error = await _ensure_sandbox(session, hardware=hardware, **create_kwargs)
+    except Exception as e:
+        return f"Failed to create sandbox: {e}", False
+
+    if error:
+        return error, False
+
+    return (
+        f"Sandbox created: {sb.space_id}\n"
+        f"URL: {sb.url}\n"
+        f"Hardware: {hardware}\n"
+        f"Use bash/read/write/edit to interact with it."
+    ), True
+
+
+def _make_tool_handler(sandbox_tool_name: str):
+    """Factory: create a handler for a sandbox operation tool."""
+
+    async def handler(args: dict[str, Any], session: Any = None) -> tuple[str, bool]:
+        # Auto-create sandbox if not present
+        try:
+            sb, error = await _ensure_sandbox(session)
+        except Exception as e:
+            return f"Failed to auto-create sandbox: {e}", False
+
+        if error:
+            return error, False
+
+        try:
+            result = await asyncio.to_thread(sb.call_tool, sandbox_tool_name, args)
+            if result.success:
+                return result.output or "(no output)", True
+            else:
+                error_msg = result.error or "Unknown error"
+                output = result.output
+                if output:
+                    return f"{output}\n\nERROR: {error_msg}", False
+                return f"ERROR: {error_msg}", False
+        except Exception as e:
+            return f"Sandbox operation failed: {e}", False
+
+    return handler
+
+
+def get_sandbox_tools():
+    """Return all 5 sandbox ToolSpecs (sandbox_create + 4 operation tools)."""
+    from agent.core.tools import ToolSpec
+
+    tools = []
+
+    # sandbox_create (explicit creation, requires approval)
+    tools.append(
+        ToolSpec(
+            name=SANDBOX_CREATE_TOOL_SPEC["name"],
+            description=SANDBOX_CREATE_TOOL_SPEC["description"],
+            parameters=SANDBOX_CREATE_TOOL_SPEC["parameters"],
+            handler=sandbox_create_handler,
+        )
+    )
+
+    # Operation tools (auto-execute, no approval needed)
+    for name in Sandbox.TOOLS.keys():
+        spec = Sandbox.TOOLS[name]
+        tools.append(
+            ToolSpec(
+                name=name,
+                description=spec["description"],
+                parameters=spec["parameters"],
+                handler=_make_tool_handler(name),
+            )
+        )
+
+    return tools