Sandbox implementation

agent/context_manager/manager.py

@@ -23,11 +23,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)
@@ -78,7 +78,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
@@ -112,6 +114,7 @@ class ContextManager:
             model=model_name,
             messages=messages_to_summarize,
             max_completion_tokens=self.compact_size,
+            tools=tool_specs,
         )
         summarized_message = Message(
             role="assistant", content=response.choices[0].message.content

agent/core/agent_loop.py

@@ -6,6 +6,7 @@ import asyncio
 import json
 
 from litellm import ChatCompletionMessageToolCall, Message, ModelResponse, acompletion
+from litellm.exceptions import ContextWindowExceededError
 from lmnr import observe
 
 from agent.config import Config
@@ -38,7 +39,9 @@ def _validate_tool_args(tool_args: dict) -> tuple[bool, str | None]:
     return True, None
 
 
-def _needs_approval(tool_name: str, tool_args: dict, config: Config | None = None) -> bool:
+def _needs_approval(
+    tool_name: str, tool_args: dict, config: Config | None = None
+) -> bool:
     """Check if a tool call requires user approval before execution."""
     # Yolo mode: skip all approvals
     if config and config.yolo_mode:
@@ -49,23 +52,31 @@ def _needs_approval(tool_name: str, tool_args: dict, config: Config | None = Non
     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"]:
             return False
-        
+
         # Check if this is a CPU-only job
         # hardware_flavor is at top level of tool_args, not nested in args
-        hardware_flavor = tool_args.get("hardware_flavor") or tool_args.get("flavor") or tool_args.get("hardware") or "cpu-basic"
+        hardware_flavor = (
+            tool_args.get("hardware_flavor")
+            or tool_args.get("flavor")
+            or tool_args.get("hardware")
+            or "cpu-basic"
+        )
         is_cpu_job = hardware_flavor in CPU_FLAVORS
-        
+
         if is_cpu_job:
             if config and not config.confirm_cpu_jobs:
                 return False
             return True
-        
+
         return True
-    
+
     # Check for file upload operations (hf_private_repos or other tools)
     if tool_name == "hf_private_repos":
         operation = tool_args.get("operation", "")
@@ -86,19 +97,43 @@ def _needs_approval(tool_name: str, tool_args: dict, config: Config | None = Non
     # hf_repo_git: destructive operations require approval
     if tool_name == "hf_repo_git":
         operation = tool_args.get("operation", "")
-        if operation in ["delete_branch", "delete_tag", "merge_pr", "create_repo", "update_repo"]:
+        if operation in [
+            "delete_branch",
+            "delete_tag",
+            "merge_pr",
+            "create_repo",
+            "update_repo",
+        ]:
             return True
 
     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"""
 
     @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)
@@ -125,6 +160,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()
 
@@ -261,6 +299,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
 
@@ -272,18 +318,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",
@@ -303,20 +337,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:
         """Handle undo (like undo in codex.rs:1314)"""
@@ -489,7 +509,8 @@ async def process_submission(session: Session, submission) -> bool:
         return True
 
     if op.op_type == OpType.COMPACT:
-        await Handlers.compact(session)
+        # compact from the frontend
+        await _compact_and_notify(session)
         return True
 
     if op.op_type == OpType.UNDO:

agent/core/session.py

@@ -59,6 +59,7 @@ class Session:
         self.is_running = True
         self.current_task: asyncio.Task | None = None
         self.pending_approval: Optional[dict[str, Any]] = None
+        self.sandbox = None
 
         # Session trajectory logging
         self.logged_events: list[dict] = []

agent/core/tools.py

@@ -45,6 +45,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 (
@@ -327,6 +328,9 @@ def create_builtin_tools() -> list[ToolSpec]:
         ),
     ]
 
+    # Sandbox tools
+    tools = get_sandbox_tools() + tools
+
     tool_names = ", ".join([t.name for t in tools])
     print(f"Loaded {len(tools)} built-in tools: {tool_names}")
 

agent/prompts/system_prompt_v2.yaml

@@ -186,61 +186,59 @@ system_prompt: |
   3. ✅ Determine optimal processing approach based on requirements
   4. ✅ Plan output format and destination
 
-  ## PHASE 3: IMPLEMENT (Execute with Researched Approaches)
-
-  ### For Training Tasks
-
-  ⚠️ **TRAINING REQUIREMENTS CHECKLIST:**
-
-  **Before Submission:**
-  - [ ] Researched current TRL documentation
-  - [ ] Found and verified base model
-  - [ ] Found dataset and VALIDATED columns and conversational format matches method
-  - [ ] Selected optimal model + dataset + hardware configuration
-  - [ ] Created plan with plan_tool
-  - [ ] Researched Trackio monitoring setup
-
-  **Training Script MUST Include:**
-  - [ ] Imports from researched documentation (current APIs)
-  - [ ] Trackio initialization with project/run_name/config
-  - [ ] Model and tokenizer loading
-  - [ ] Dataset loading with verified columns and conversational format
-  - [ ] Training config with ALL critical settings:
+  ## PHASE 3: IMPLEMENT (Develop in Sandbox, Launch via Jobs)
+
+  ⚠️ **CRITICAL WORKFLOW: Sandbox First, Jobs Second**
+
+  For ANY implementation task (training, data processing, inference), follow this pattern:
+
+  **Step 1: Create a sandbox** — `sandbox_create` with appropriate hardware (cpu-basic for scripting, t4-small for GPU testing)
+  **Step 2: Develop & iterate** — Write scripts, install dependencies, test with small runs, fix errors interactively
+  **Step 3: Launch via hf_jobs** — Once the script works, pass the sandbox file path directly: `hf_jobs(operation="run", script="/app/train.py", ...)`
+
+  This is the CORRECT pattern:
+  ```
+  sandbox_create(hardware="t4-small")     # interactive dev environment
+  bash("pip install trl transformers")     # install deps
+  write("/app/train.py", "...")            # write training script
+  bash("cd /app && python train.py --max_steps 10")  # test run
+  edit("/app/train.py", ...)               # fix issues
+  bash("cd /app && python train.py --max_steps 10")  # verify fix
+  hf_jobs(operation="run", script="/app/train.py", hardware_flavor="a10g-large", timeout="4h")  # launch at scale
+  ```
+
+  Do NOT write long inline scripts directly in hf_jobs if necessary — develop in sandbox first.
+
+  ### Training Script Requirements
+
+  **Script MUST Include:**
+  - Imports from researched documentation (current APIs)
+  - Trackio initialization with project/run_name/config
+  - Model and tokenizer loading
+  - Dataset loading with verified columns and conversational format
+  - Training config with ALL critical settings:
     - `push_to_hub=True` ⚠️ MANDATORY
     - `hub_model_id="username/model-name"` ⚠️ MANDATORY
     - `report_to=["trackio"]` (for monitoring)
     - `output_dir="./output"`
     - `num_train_epochs`, `per_device_train_batch_size`, `learning_rate`
     - `logging_steps`, `save_steps`
-    - `max_length` if needed (default 1024 usually fine)
-  - [ ] Trainer initialization with model, args, dataset, tokenizer
-  - [ ] `trainer.train()` call
-  - [ ] `trainer.push_to_hub()` at end ⚠️ MANDATORY
-  - [ ] `tracker.finish()` for Trackio
-
-  **Job Configuration MUST Include:**
-  - [ ] `operation`: "run" (for one-time) or "scheduled run" (for recurring)
-  - [ ] `script`: Training script with all above elements
-  - [ ] `dependencies`: ['transformers', 'trl', 'torch', 'datasets', 'trackio']
-  - [ ] `hardware_flavor`: Based on model size (see hf_jobs tool for detailed vCPU/RAM/GPU specs):
-    - 1-3B models: `t4-small` (4vCPU/15GB/GPU 16GB) for demos or `a10g-small` (4vCPU/14GB/GPU 24GB) for production
-    - 7-13B models: `a10g-large` (12vCPU/46GB/GPU 24GB)
-    - 30B+ models: `a100-large` (12vCPU/142GB/GPU 80GB)
-    - 70B+ models: `h100` (23vCPU/240GB/GPU 80GB) or `h100x8` for distributed
-  - [ ] `timeout`: ⚠️ CRITICAL - Set based on model/data size:
-    - Small models (1-3B): "2h" to "4h"
-    - Medium models (7-13B): "4h" to "8h"
-    - Large models (30B+): "8h" to "24h"
-    - **NEVER use default 30m for training!**
+  - `trainer.train()` call
+  - `trainer.push_to_hub()` at end ⚠️ MANDATORY
+
+  **hf_jobs Launch Configuration:**
+  - `script`: Path to sandbox file (e.g. "/app/train.py") or inline code
+  - `dependencies`: ['transformers', 'trl', 'torch', 'datasets', 'trackio']
+  - `hardware_flavor`: Based on model size:
+    - 1-3B models: `t4-small` or `a10g-small`
+    - 7-13B models: `a10g-large`
+    - 30B+ models: `a100-large`
+    - 70B+ models: `h100` or `h100x8`
+  - `timeout`: ⚠️ CRITICAL — Small (2-4h), Medium (4-8h), Large (8-24h). NEVER default 30m for training.
 
   ### For Data Processing Tasks
 
-  **Script Requirements:**
-  - Load dataset with `load_dataset`
-  - Process according to user requirements
-  - Push results with `push_to_hub()` or upload to `hf_private_repos`
-
-  **Job Configuration:**
+  **Same pattern:** develop script in sandbox, test on subset, launch via hf_jobs.
   - Use `cpu-upgrade` or `cpu-performance` for most data tasks
   - Set timeout based on dataset size (1-4 hours typical)
 
@@ -341,6 +339,21 @@ system_prompt: |
   - ⚠️ Include HF_TOKEN for Hub operations
   - ⚠️ Storage is EPHEMERAL - must push_to_hub
 
+  ## Sandbox (Interactive Development Environment)
+
+  **sandbox_create:**
+  - ⚠️ **Create a sandbox FIRST for any implementation task** — develop and test before launching jobs
+  - Persistent remote Linux environment on HF Spaces
+  - First call sandbox_create with hardware choice, then use bash/read/write/edit freely
+  - Hardware: cpu-basic (free tier), cpu-upgrade (8vCPU/32GB), t4-small (16GB GPU), a10g-small (24GB GPU), a10g-large (24GB GPU + 46GB RAM), a100-large (80GB GPU)
+  - `pip install` works out of the box — no special flags needed
+  - Workflow: sandbox_create → write script → test → fix → hf_jobs(script="/app/script.py") to launch at scale
+
+  **bash / read / write / edit:**
+  - Available after sandbox_create — no additional approvals needed
+  - Same semantics as local file/shell operations, but run on the remote sandbox
+  - bash: run shell commands; read/write/edit: file operations
+
   **hf_private_repos:**
   - Store job outputs persistently in datasets with push_to_hub (jobs lose files after completion)
   - Upload logs, scripts, results that can't push_to_hub

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/dataset_tools.py

@@ -388,22 +388,15 @@ def _format_parquet_files(data: dict, max_rows: int = 10) -> str | None:
 HF_INSPECT_DATASET_TOOL_SPEC = {
     "name": "hf_inspect_dataset",
     "description": (
-        "Inspect a Hugging Face dataset comprehensively in one call.\n\n"
-        "## What you get\n"
-        "- Status check (validates dataset works without errors)\n"
-        "- All configs and splits (row counts/shares may be '?' when metadata is missing)\n"
-        "- Column names and types (schema)\n"
-        "- Sample rows to understand data format\n"
-        "- Parquet file structure and sizes\n\n"
-        "## CRITICAL\n"
-        "**Always inspect datasets before writing training code** to understand:\n"
-        "- Column names for your dataloader\n"
-        "- Data types and format\n"
-        "- Available splits (train/test/validation)\n\n"
-        "Supports private/gated datasets when HF_TOKEN is set.\n\n"
-        "## Examples\n"
-        '{"dataset": "stanfordnlp/imdb"}\n'
-        '{"dataset": "nyu-mll/glue", "config": "mrpc", "sample_rows": 5}\n'
+        "Inspect a HF dataset in one call: status, configs/splits, schema, sample rows, parquet info.\n\n"
+        "REQUIRED before any training job to verify dataset format matches training method:\n"
+        "  SFT: needs 'messages', 'text', or 'prompt'/'completion'\n"
+        "  DPO: needs 'prompt', 'chosen', 'rejected'\n"
+        "  GRPO: needs 'prompt'\n"
+        "All datasets used for training have to be in conversational ChatML format to be compatible with HF libraries.'\n"
+        "Training will fail with KeyError if columns don't match.\n\n"
+        "Also use to get example datapoints, understand column names, data types, and available splits before writing any data loading code. "
+        "Supports private/gated datasets when HF_TOKEN is set."
     ),
     "parameters": {
         "type": "object",

agent/tools/docs_tools.py

@@ -845,17 +845,12 @@ DOC_ENDPOINTS = [
 EXPLORE_HF_DOCS_TOOL_SPEC = {
     "name": "explore_hf_docs",
     "description": (
-        "Explore Hugging Face documentation structure and discover available pages with 200-character previews. "
-        "⚠️ MANDATORY: ALWAYS use this BEFORE implementing any ML task (training, fine-tuning, data processing, inference). "
-        "Your training data may be outdated - current documentation is the source of truth. "
-        "**Use when:** (1) Starting any implementation task, (2) User asks 'how to' questions, "
-        "(3) Before writing training/processing code, (4) Researching library capabilities, "
-        "(5) Verifying API syntax and parameters. "
-        "**Pattern:** explore (discover structure) → fetch_hf_docs (get details) → implement with researched approach. "
-        "Returns: Sidebar navigation with titles, URLs, and glimpses of all pages in the selected documentation. "
-        "**Then:** Use fetch_hf_docs with specific URLs from results to get full content. "
-        "**Critical for reliability:** Never implement based on internal knowledge without checking current docs first - APIs change frequently."
-        " By default returns the top 20 results; set max_results (max 50) to adjust."
+        "Browse HF documentation structure — discover all available documentation with 200-char previews.\n\n"
+        "Use this to find relevant documentation and/or examples with detailed parameter docs and API reference. "
+        "To be used together with github_find_examples and github_read_file to find working examples and documentation.\n\n"
+        "Pattern: explore_hf_docs (find relevant pages) → fetch_hf_docs (get full content).\n\n"
+        "For training tasks: fetch the trainer config docs (SFTConfig, DPOConfig, GRPOConfig) to verify parameter names. "
+        "Returns top 20 results by default; set max_results (max 50) to adjust."
     ),
     "parameters": {
         "type": "object",
@@ -928,16 +923,10 @@ EXPLORE_HF_DOCS_TOOL_SPEC = {
 HF_DOCS_FETCH_TOOL_SPEC = {
     "name": "fetch_hf_docs",
     "description": (
-        "Fetch full markdown content of a specific HF documentation page. "
-        "⚠️ CRITICAL: Use this after explore_hf_docs to get detailed implementation guidance. "
-        "**Use when:** (1) Found relevant page in explore_hf_docs results, (2) Need complete API documentation, "
-        "(3) Need training method details (SFT/DPO/GRPO), (4) Need configuration examples, "
-        "(5) Need parameter descriptions and usage patterns. "
-        "**Pattern:** explore_hf_docs (find relevant page) → fetch_hf_docs (get full content) → implement using documented approach. "
-        "Provide full URL from explore_hf_docs results (e.g., 'https://huggingface.co/docs/trl/sft_trainer'). "
-        "Returns: Complete markdown documentation with examples, parameters, and usage patterns. "
-        "**For training tasks:** ALWAYS fetch trainer docs (SFTConfig, DPOConfig, etc.) before creating training scripts. "
-        "**Critical for reliability:** This ensures you use current APIs and best practices."
+        "Fetch full markdown content of an HF documentation page. Use after explore_hf_docs.\n\n"
+        "Critical for finding documentation e.g. current trainer configuration parameters (SFTConfig, DPOConfig, etc.) "
+        "Use for researching solutions and before writing training scripts. Your internal knowledge is outdated.\n\n"
+        "Provide the full URL from explore_hf_docs results. The .md extension is added automatically."
     ),
     "parameters": {
         "type": "object",

agent/tools/github_find_examples.py

@@ -405,55 +405,16 @@ def find_examples(
 GITHUB_FIND_EXAMPLES_TOOL_SPEC = {
     "name": "github_find_examples",
     "description": (
-        "Discover working code examples, tutorials, scripts, and demos in GitHub repositories. "
-        "⚠️ CRITICAL: ALWAYS use this BEFORE implementing ML tasks - find working reference code first. "
-        "Your training data may be outdated; real repository examples show current best practices. "
-        "**Use when:** (1) Starting any ML implementation (training, inference, evaluation), "
-        "(2) User asks 'how to' questions about libraries, (3) Need reference implementations, "
-        "(4) Exploring library capabilities, (5) Before writing training/processing scripts. "
-        "**Pattern:** github_find_examples (discover) → github_read_file (study code) → implement with researched approach. "
-        "Returns: List of example files (scripts/notebooks/tutorials) with paths and URLs, sorted by relevance. "
-        "**Then:** Use github_read_file to read the actual implementation code. "
-        "**Critical for reliability:** Real examples prevent outdated API usage and show proven patterns. "
-        "## How it works\n\n"
-        "1. Fetches all example files (examples/, scripts/, tutorials/, demos/, notebooks/, etc.) from repository\n"
-        "2. If keyword provided, scores files against keyword using fuzzy matching\n"
-        "3. Returns best matches sorted by relevance and pattern priority\n"
-        "4. Provides copyable parameters for github_read_file tool\n\n"
-        "## Examples\n\n"
-        "<example>\n"
-        "// ML Workflow Step: Find GRPO training examples before implementation\n"
-        "// Task: Starting GRPO fine-tuning project, need reference implementation\n"
-        "{\n"
-        "  keyword: 'grpo',\n"
-        "  repo: 'trl',\n"
-        "  org: 'huggingface'\n"
-        "}\n"
-        "// Returns: examples/scripts/grpo_agent.py, examples/scripts/grpo_vlm.py\n"
-        "// Next step: github_read_file to study working implementation\n"
-        "</example>\n\n"
-        "<example>\n"
-        "// ML Workflow Step: Discover all available training methods\n"
-        "// Task: Exploring TRL training options before choosing approach\n"
-        "{\n"
-        "  repo: 'trl',\n"
-        "  org: 'huggingface',\n"
-        "  max_results: 20\n"
-        "}\n"
-        "// Lists: SFT, DPO, GRPO, PPO, reward modeling examples\n"
-        "// Helps user choose appropriate method\n"
-        "</example>\n\n"
-        "<example>\n"
-        "// ML Workflow Step: Find LoRA fine-tuning examples\n"
-        "// Task: Learning parameter-efficient fine-tuning patterns\n"
-        "{\n"
-        "  keyword: 'lora',\n"
-        "  repo: 'peft',\n"
-        "  org: 'huggingface'\n"
-        "}\n"
-        "// Discovers LoRA configuration and training examples\n"
-        "// Shows current PEFT API usage patterns\n"
-        "</example>"
+        "Find working example scripts in GitHub repositories (from a list of predetermined directories e.g. examples/, scripts/, tutorials/, etc.). "
+        "Uses fuzzy keyword matching.\n\n"
+        "MANDATORY before writing any ML training, fine-tuning, or inference code. "
+        "Your internal knowledge of library APIs is outdated — working examples show current API patterns.\n\n"
+        "Sequence: github_find_examples → github_read_file (study the example) → implement based on what you found.\n\n"
+        "Skip this only for: simple data queries, status checks, non-code tasks.\n\n"
+        "Examples:\n"
+        "  {keyword: 'sft', repo: 'trl'} → finds examples/scripts/sft.py\n"
+        "  {keyword: 'grpo', repo: 'trl'} → finds GRPO training examples\n"
+        "  {repo: 'trl', max_results: 20} → lists all available training method examples"
     ),
     "parameters": {
         "type": "object",

agent/tools/github_read_file.py

@@ -250,59 +250,13 @@ def read_file(
 GITHUB_READ_FILE_TOOL_SPEC = {
     "name": "github_read_file",
     "description": (
-        "Read file contents from GitHub repositories with line range support (default 300 lines). "
-        "⚠️ CRITICAL: Use AFTER github_find_examples to study working implementation code. "
-        "**Use when:** (1) Found example file via github_find_examples and need full code, "
-        "(2) Need to read trainer class implementation, (3) Study configuration patterns, "
-        "(4) Read specific code sections with line ranges, (5) Review code from specific branches/commits. "
-        "**Pattern:** github_find_examples (discover files) → github_read_file (read code) → implement using researched patterns. "
-        "Returns: File contents with line numbers, formatted for LLM reading. Auto-converts Jupyter notebooks to markdown. "
-        "**Then:** Implement using patterns and APIs from the example code. "
-        "**Critical for reliability:** Reading working examples prevents API errors and shows current best practices. "
+        "Read file contents from GitHub repositories. Returns first 300 lines by default. "
+        "Auto-converts Jupyter notebooks to markdown.\n\n"
+        "Use AFTER github_find_examples to study the working implementation. "
+        "The purpose is to learn current API patterns — imports, trainer configs, dataset handling — "
+        "so your implementation uses correct, up-to-date code.\n\n"
         "Use line_start/line_end for large files (>300 lines) to read specific sections.\n\n"
-        "## When to use this tool\n\n"
-        "- When reading example code, trainer implementations, or configuration files\n"
-        "- After github_find_examples returns file paths you want to study\n"
-        "- When investigating specific code sections with line ranges\n"
-        "- When reading from specific branches, tags, or commits (use ref parameter)\n\n"
-        "## When NOT to use this tool\n\n"
-        "- When you don't know exact file path (use github_find_examples or github_search_code first)\n"
-        "- When searching for code patterns across repos (use github_search_code instead)\n\n"
-        "## Examples\n\n"
-        "<example>\n"
-        "// ML Workflow Step: Read GRPO trainer class after finding via github_find_examples\n"
-        "// Use case: Understand GRPOTrainer API, parameters, and methods\n"
-        "{\n"
-        "  repo: 'huggingface/trl',\n"
-        "  path: 'trl/trainer/grpo_trainer.py',\n"
-        "  line_start: 1,\n"
-        "  line_end: 200\n"
-        "}\n"
-        "// Read class definition and constructor to understand current API\n"
-        "// Shows: __init__ parameters, configuration, required arguments\n"
-        "</example>\n\n"
-        "<example>\n"
-        "// ML Workflow Step: Study complete training script from examples\n"
-        "// Use case: Learn end-to-end VLM fine-tuning workflow\n"
-        "{\n"
-        "  repo: 'huggingface/trl',\n"
-        "  path: 'examples/scripts/grpo_vlm.py'\n"
-        "}\n"
-        "// Returns first 300 lines - shows full training setup\n"
-        "// Use line_start/line_end if need to read more\n"
-        "</example>\n\n"
-        "<example>\n"
-        "// ML Workflow Step: Check TrainingArguments configuration patterns\n"
-        "// Use case: Learn how to structure training configs correctly\n"
-        "{\n"
-        "  repo: 'huggingface/transformers',\n"
-        "  path: 'examples/pytorch/language-modeling/run_clm.py',\n"
-        "  line_start: 50,\n"
-        "  line_end: 150\n"
-        "}\n"
-        "// Read argument parsing and config setup section\n"
-        "// Shows: current parameter names, default values, best practices\n"
-        "</example>"
+        "When NOT to use: when you don't know the file path (use github_find_examples first)."
     ),
     "parameters": {
         "type": "object",

agent/tools/jobs_tool.py

@@ -9,7 +9,7 @@ import base64
 import http.client
 import os
 import re
-from typing import Any, Dict, Literal, Optional, Callable, Awaitable
+from typing import Any, Awaitable, Callable, Dict, Literal, Optional
 
 import httpx
 from huggingface_hub import HfApi
@@ -25,38 +25,33 @@ from agent.tools.utilities import (
 )
 
 # Hardware flavors
-CPU_FLAVORS = ["cpu-basic", "cpu-upgrade", "cpu-performance", "cpu-xl"]
+CPU_FLAVORS = ["cpu-basic", "cpu-upgrade"]
 GPU_FLAVORS = [
-    "sprx8",
-    "zero-a10g",
     "t4-small",
     "t4-medium",
-    "l4x1",
-    "l4x4",
-    "l40sx1",
-    "l40sx4",
-    "l40sx8",
     "a10g-small",
     "a10g-large",
     "a10g-largex2",
     "a10g-largex4",
     "a100-large",
-    "h100",
-    "h100x8",
+    "a100x4",
+    "a100x8",
+    "l4x1",
+    "l4x4",
+    "l40sx1",
+    "l40sx4",
+    "l40sx8",
 ]
 
 # Detailed specs for display (vCPU/RAM/GPU VRAM)
-CPU_FLAVORS_DESC = (
-    "cpu-basic(2vCPU/16GB), cpu-upgrade(8vCPU/32GB), cpu-performance, cpu-xl"
-)
+CPU_FLAVORS_DESC = "cpu-basic(2vCPU/16GB), cpu-upgrade(8vCPU/32GB)"
 GPU_FLAVORS_DESC = (
     "t4-small(4vCPU/15GB/GPU 16GB), t4-medium(8vCPU/30GB/GPU 16GB), "
-    "l4x1(8vCPU/30GB/GPU 24GB), l4x4(48vCPU/186GB/GPU 96GB), "
-    "l40sx1(8vCPU/62GB/GPU 48GB), l40sx4(48vCPU/382GB/GPU 192GB), l40sx8(192vCPU/1534GB/GPU 384GB), "
-    "a10g-small(4vCPU/14GB/GPU 24GB), a10g-large(12vCPU/46GB/GPU 24GB), "
+    "a10g-small(4vCPU/15GB/GPU 24GB), a10g-large(12vCPU/46GB/GPU 24GB), "
     "a10g-largex2(24vCPU/92GB/GPU 48GB), a10g-largex4(48vCPU/184GB/GPU 96GB), "
-    "a100-large(12vCPU/142GB/GPU 80GB), h100(23vCPU/240GB/GPU 80GB), h100x8(184vCPU/1920GB/GPU 640GB), "
-    "zero-a10g(dynamic alloc)"
+    "a100-large(12vCPU/142GB/GPU 80GB), a100x4(48vCPU/568GB/GPU 320GB), a100x8(96vCPU/1136GB/GPU 640GB), "
+    "l4x1(8vCPU/30GB/GPU 24GB), l4x4(48vCPU/186GB/GPU 96GB), "
+    "l40sx1(8vCPU/62GB/GPU 48GB), l40sx4(48vCPU/382GB/GPU 192GB), l40sx8(192vCPU/1534GB/GPU 384GB)"
 )
 SPECIALIZED_FLAVORS = ["inf2x6"]
 ALL_FLAVORS = CPU_FLAVORS + GPU_FLAVORS + SPECIALIZED_FLAVORS
@@ -118,6 +113,21 @@ def _filter_uv_install_output(logs: list[str]) -> list[str]:
     return logs
 
 
+_DEFAULT_ENV = {
+    "HF_HUB_DISABLE_PROGRESS_BARS": "1",
+    "TQDM_DISABLE": "1",
+    "TRANSFORMERS_VERBOSITY": "warning",
+    "HF_HUB_ENABLE_HF_TRANSFER": "1",
+}
+
+
+def _add_default_env(params: Dict[str, Any] | None) -> Dict[str, Any]:
+    """Inject default env vars for clean, agent-friendly output."""
+    result = dict(_DEFAULT_ENV)
+    result.update(params or {})  # user-provided values override defaults
+    return result
+
+
 def _add_environment_variables(params: Dict[str, Any] | None) -> Dict[str, Any]:
     token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN") or ""
 
@@ -374,7 +384,9 @@ class HfJobsTool:
                 def log_producer():
                     try:
                         # fetch_job_logs is a blocking sync generator
-                        logs_gen = self.api.fetch_job_logs(job_id=job_id, namespace=namespace)
+                        logs_gen = self.api.fetch_job_logs(
+                            job_id=job_id, namespace=namespace
+                        )
                         for line in logs_gen:
                             # Push line to queue thread-safely
                             loop.call_soon_threadsafe(queue.put_nowait, line)
@@ -497,7 +509,7 @@ class HfJobsTool:
                 self.api.run_job,
                 image=image,
                 command=command,
-                env=args.get("env"),
+                env=_add_default_env(args.get("env")),
                 secrets=_add_environment_variables(args.get("secrets")),
                 flavor=args.get("hardware_flavor", "cpu-basic"),
                 timeout=args.get("timeout", "30m"),
@@ -715,7 +727,7 @@ To verify, call this tool with `{{"operation": "inspect", "job_id": "{job_id}"}}
                 image=image,
                 command=command,
                 schedule=schedule,
-                env=args.get("env"),
+                env=_add_default_env(args.get("env")),
                 secrets=_add_environment_variables(args.get("secrets")),
                 flavor=args.get("hardware_flavor", "cpu-basic"),
                 timeout=args.get("timeout", "30m"),
@@ -875,56 +887,31 @@ To inspect, call this tool with `{{"operation": "scheduled inspect", "scheduled_
 HF_JOBS_TOOL_SPEC = {
     "name": "hf_jobs",
     "description": (
-        "Execute Python scripts or Docker containers on HF cloud infrastructure (CPUs/GPUs) in one of two modes. "
-        "\n\n"
-        "**Two Modes (mutually exclusive):**\n"
-        "1. Python mode: using 'script' arg (REQUIRED) + 'dependencies'\n"
-        "2. Docker mode: using 'command' arg (REQUIRED) + 'image'\n\n"
-        "🚨 **REQUIRED:** You MUST provide exactly ONE of: 'script' (Python code as string) OR 'command' (Docker command as array). "
-        "They are mutually exclusive - provide one or the other, never both, never neither. "
-        "Do NOT call with just {'operation': 'run'} - always include your code. Example: {'operation': 'run', 'script': 'import torch; print(torch.cuda.is_available())', 'dependencies': ['torch']} or {'operation': 'run', 'command': ['duckdb', '-c', 'select 1 + 2']', 'image': 'duckdb/duckdb'}\n\n"
-        "⚠️ CRITICAL for reliability: (1) Jobs run ASYNC - provide monitoring URL immediately, don't poll; "
-        "(2) Set timeout >30min (default too short - training needs 2-8h); "
-        "(3) HF_TOKEN auto-loaded to secrets for Hub ops (push_to_hub, private repos); "
-        "(4) Job storage EPHEMERAL - MUST push_to_hub() or ALL work is LOST. "
-        "**Use when:** User wants cloud compute, training models, data processing, batch inference, GPU workloads, scheduled tasks. "
-        "ALWAYS use this tool (✓), never bash 'hf jobs' commands (✗). Pass script content inline (✓), don't save to files unless requested (✗). "
-        "\n\n"
-        "**Operations:** run, ps, logs, inspect, cancel, scheduled run, scheduled ps, scheduled inspect, scheduled delete, scheduled suspend, scheduled resume. "
-        "**Available Hardware (vCPU/RAM/GPU):**\n"
-        f"• CPU: {CPU_FLAVORS_DESC}\n"
-        f"• GPU: {GPU_FLAVORS_DESC}\n"
-        "  ◦ Common: t4-small ($0.60/hr, demos/1-3B models), a10g-small ($1/hr), a10g-large ($2/hr, production 7-13B), a100-large ($4/hr, 30B+), h100 ($6/hr, 70B+)\n\n"
-        "**After Submission Ground Rules:**\n"
-        "✓ Return immediately with job ID and monitoring URL\n"
-        "✓ Provide expected completion time and cost estimate\n"
-        "✓ For training: Include Trackio dashboard URL\n"
-        "✓ Note user can check status later\n"
-        "✗ DON'T poll logs automatically\n"
-        "✗ DON'T wait for completion\n"
-        "✗ DON'T check status unless user asks\n\n"
-        "**For Training Tasks:**\n"
-        "• ALWAYS research TRL docs first: explore_hf_docs('trl') → fetch_hf_docs(<trainer_url>)\n"
-        "• ALWAYS validate dataset format with hub_repo_details (SFT needs messages/text, DPO needs chosen/rejected)\n"
-        "• ALWAYS include Trackio monitoring in script (explore_hf_docs('trackio'))\n"
-        "• ALWAYS enable push_to_hub=True in training config\n"
-        "• Set timeout 2-8h for training (NOT default 30m)\n"
-        "• Confirm model/dataset choices with user before submitting\n\n"
-        "**Examples:**\n\n"
-        "**Training - Fine-tune LLM:**\n"
-        "{'operation': 'run', 'script': '# Training script with TRL\\nfrom trl import SFTConfig, SFTTrainer\\nfrom transformers import AutoModelForCausalLM\\nmodel = AutoModelForCausalLM.from_pretrained(\"Qwen/Qwen3-4B\")\\n# ... researched implementation from docs ...\\ntrainer.train()\\ntrainer.push_to_hub(\"user-name/my-model\")', 'dependencies': ['transformers', 'trl', 'torch', 'datasets', 'trackio'], 'hardware_flavor': 'a10g-large', 'timeout': '4h'}\n\n"
-        "**Data Processing:**\n"
-        "{'operation': 'run', 'script': 'from datasets import load_dataset\\nds = load_dataset(\"data\")\\n# process...\\nds.push_to_hub(\"user/processed\")', 'dependencies': ['datasets', 'pandas'], 'hardware_flavor': 'cpu-upgrade', 'timeout': '2h'}\n\n"
-        "**Scheduled Daily Job:**\n"
-        "{'operation': 'scheduled run', 'schedule': '@daily', 'script': 'from datasets import Dataset\\nimport pandas as pd\\n# scrape/generate data\\ndf = pd.DataFrame(data)\\nds = Dataset.from_pandas(df)\\nds.push_to_hub(\"user-name/daily-dataset\")', 'dependencies': ['datasets', 'pandas'], 'hardware_flavor': 'cpu-basic'}\n\n"
-        "**Docker Mode:**\n"
-        "{'operation': 'run', 'image': 'pytorch/pytorch:2.0.0-cuda11.7-cudnn8-runtime', 'command': ['python', 'train.py', '--epochs', '10'], 'hardware_flavor': 'a100-large'}\n\n"
-        "**Monitor Operations:**\n"
-        "{'operation': 'ps'} - List all jobs\n"
-        "{'operation': 'logs', 'job_id': 'xxx'} - Stream logs (only when user requests)\n"
-        "{'operation': 'inspect', 'job_id': 'xxx'} - Get job details\n"
-        "{'operation': 'cancel', 'job_id': 'xxx'} - Stop job\n\n"
-        "⚠️ CRITICAL: Files created during execution are DELETED when job finishes. MUST push_to_hub() all outputs (models, datasets, artifacts) in script. For logs/scripts, use hf_private_repos after completion."
+        "Execute Python scripts or Docker containers on HF cloud infrastructure.\n\n"
+        "Two modes (mutually exclusive): Python mode (script + dependencies) or Docker mode (command + image). "
+        "Provide exactly ONE of 'script' or 'command'.\n\n"
+        "BEFORE submitting training/fine-tuning jobs:\n"
+        "- You MUST have called github_find_examples + github_read_file to find a working reference implementation. "
+        "Scripts based on your internal knowledge WILL use outdated APIs and fail.\n"
+        "- You MUST have validated dataset format via hf_inspect_dataset or hub_repo_details.\n"
+        "- Training config MUST include push_to_hub=True and hub_model_id. "
+        "Job storage is EPHEMERAL — all files are deleted when the job ends. Without push_to_hub, trained models are lost permanently.\n"
+        "- Include trackio monitoring and provide the dashboard URL to the user.\n\n"
+        "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 — if there's a bug, all jobs fail.\n\n"
+        "Operations: run, ps, logs, inspect, cancel, scheduled run/ps/inspect/delete/suspend/resume.\n\n"
+        f"Hardware: CPU: {CPU_FLAVORS_DESC}. GPU: {GPU_FLAVORS_DESC}.\n"
+        "Common picks: t4-small ($0.60/hr, 1-3B), a10g-large ($2/hr, 7-13B), a100-large ($4/hr, 30B+), h100 ($6/hr, 70B+). "
+        "Note: a10g-small and a10g-large have the SAME 24GB GPU — the difference is CPU/RAM only.\n\n"
+        "OOM RECOVERY: When a training job fails with CUDA OOM:\n"
+        "1. Reduce per_device_train_batch_size and increase gradient_accumulation_steps proportionally (keep effective batch size identical)\n"
+        "2. Enable gradient_checkpointing=True\n"
+        "3. Upgrade to larger GPU (a10g→a100→h100)\n"
+        "Do NOT switch training methods (e.g. full SFT to LoRA) or reduce max_length — those change what the user gets and require explicit approval.\n\n"
+        "Examples:\n"
+        "Training: {'operation': 'run', 'script': '/app/train.py', 'dependencies': ['transformers', 'trl', 'torch', 'datasets', 'trackio'], 'hardware_flavor': 'a100-large', 'timeout': '8h'}\n"
+        "Monitor: {'operation': 'ps'}, {'operation': 'logs', 'job_id': 'xxx'}, {'operation': 'cancel', 'job_id': 'xxx'}"
+        "Docker: {'operation': 'run', 'command': ['duckdb', '-c', 'select 1 + 2'], 'image': 'duckdb/duckdb', 'hardware_flavor': 'cpu-basic', 'timeout': '1h'}\n"
     ),
     "parameters": {
         "type": "object",
@@ -944,58 +931,65 @@ HF_JOBS_TOOL_SPEC = {
                     "scheduled suspend",
                     "scheduled resume",
                 ],
-                "description": (
-                    "Operation to execute. Valid values: [run, ps, logs, inspect, cancel, "
-                    "scheduled run, scheduled ps, scheduled inspect, scheduled delete, "
-                    "scheduled suspend, scheduled resume]"
-                ),
+                "description": "Operation to execute.",
             },
-            # 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",
                 "items": {"type": "string"},
-                "description": "Pip packages to install. Example: ['trl', 'torch', 'datasets', 'transformers']. Only used with 'script'.",
+                "description": (
+                    "Pip packages to install. Include ALL required packages. "
+                    "Common training set: ['transformers', 'trl', 'torch', 'datasets', 'trackio', 'accelerate']. "
+                    "Only used with 'script'."
+                ),
             },
-            # Docker specific parameters
             "image": {
                 "type": "string",
-                "description": "Docker image. Example: 'pytorch/pytorch:2.0.0-cuda11.7-cudnn8-runtime'. Use with 'run'/'scheduled run'. Optional (auto-selected if not provided).",
+                "description": "Docker image. Optional — auto-selected if not provided. Use with 'command'.",
             },
             "command": {
                 "type": "array",
                 "items": {"type": "string"},
-                "description": "Command to execute as list. Example: ['python', 'train.py', '--epochs', '10']. Triggers Docker mode. Use with 'run'/'scheduled run'. Mutually exclusive with 'script'.",
+                "description": "Command to execute as list. Triggers Docker mode. Mutually exclusive with 'script'.",
             },
-            # Hardware and environment
             "hardware_flavor": {
                 "type": "string",
-                "description": f"Hardware type. Available CPU flavors: {CPU_FLAVORS}. Available GPU flavors: {GPU_FLAVORS}. Use with 'run'/'scheduled run'.",
+                "description": (
+                    "Hardware type. Sizing guide: 1-3B params → t4-small/a10g-small, "
+                    "7-13B → a10g-large, 30B+ → a100-large, 70B+ → h100/h100x8. "
+                    f"All options: CPU: {CPU_FLAVORS}. GPU: {GPU_FLAVORS}."
+                ),
             },
             "timeout": {
                 "type": "string",
-                "description": "Max runtime. Examples: '30m', '2h', '4h'. Default: '30m'. Important for long training jobs. Use with 'run'/'scheduled run'.",
+                "description": (
+                    "Maximum job runtime. MUST be >2h for any training job — default 30m kills training mid-run. "
+                    "Guidelines: 1-3B models: 3-4h, 7-13B: 6-8h, 30B+: 12-24h. "
+                    "Use 30m-1h only for quick data processing or inference tasks. Default: '30m'."
+                ),
             },
             "env": {
                 "type": "object",
-                "description": "Environment variables. Format: {'KEY': 'VALUE'}. HF_TOKEN is automatically included from your auth. Use with 'run'/'scheduled run'.",
+                "description": "Environment variables {'KEY': 'VALUE'}. HF_TOKEN is auto-included.",
             },
-            # Job management parameters
             "job_id": {
                 "type": "string",
-                "description": "Job ID to operate on. Required for: 'logs', 'inspect', 'cancel'.",
+                "description": "Job ID. Required for: logs, inspect, cancel.",
             },
-            # Scheduled job parameters
             "scheduled_job_id": {
                 "type": "string",
-                "description": "Scheduled job ID. Required for: 'scheduled inspect', 'scheduled delete', 'scheduled suspend', 'scheduled resume'.",
+                "description": "Scheduled job ID. Required for: scheduled inspect/delete/suspend/resume.",
             },
             "schedule": {
                 "type": "string",
-                "description": "Schedule for recurring job. Presets: '@hourly', '@daily', '@weekly', '@monthly'. Cron: '0 9 * * 1' (Mon 9am). Required for: 'scheduled run'.",
+                "description": "Cron schedule or preset (@hourly, @daily, @weekly, @monthly). Required for: scheduled run.",
             },
         },
         "required": ["operation"],
@@ -1015,6 +1009,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}
+
         # Get token and namespace from HF token
         hf_token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN")
         namespace = HfApi(token=hf_token).whoami().get("name") if hf_token else None

agent/tools/plan_tool.py

@@ -85,18 +85,11 @@ def get_current_plan() -> List[Dict[str, str]]:
 PLAN_TOOL_SPEC = {
     "name": "plan_tool",
     "description": (
-        "Manage task planning and progress tracking with todo list (pending/in_progress/completed statuses). "
-        "⚠️ CRITICAL: ALWAYS use for multi-step tasks (3+ steps) and MUST update frequently to show progress. "
-        "**Use when:** (1) User provides multiple tasks, (2) Complex workflows (training, evaluation, data processing), "
-        "(3) Tasks requiring multiple tool calls, (4) Need to communicate progress clearly to user, "
-        "(5) Breaking down ambiguous requests into concrete steps. "
-        "**Pattern:** Create plan at start → Mark in_progress when starting task → Mark completed immediately after finishing → User sees clear progress. "
-        "Each call replaces entire plan (full list required). "
-        "**Critical for reliability:** Exactly ONE task in_progress at a time (not zero, not multiple). "
-        "Mark tasks completed IMMEDIATELY after finishing - don't batch completions. "
-        "**For long-running tasks:** Update plan after each major step to keep user informed. "
-        "**Only mark completed when:** Task fully accomplished, no errors, all requirements met. "
-        "Keep tasks pending if blocked/errors occur - create new task to resolve blockers."
+        "Track progress on multi-step tasks with a todo list (pending/in_progress/completed).\n\n"
+        "Use for tasks with 3+ steps. Each call replaces the entire plan (send full list).\n\n"
+        "Rules: exactly ONE task in_progress at a time. Mark completed immediately after finishing. "
+        "Only mark completed when the task fully succeeded — keep in_progress if there are errors. "
+        "Update frequently so the user sees progress."
     ),
     "parameters": {
         "type": "object",

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