stream sandbox creation logs to tool output, fix hf_jobs final status race

agent/tools/jobs_tool.py

@@ -466,11 +466,17 @@ class HfJobsTool:
                     await asyncio.sleep(retry_delay)
                     continue
 
-        # Fetch final job status
-        job_info = await _async_call(
-            self.api.inspect_job, job_id=job_id, namespace=namespace
-        )
-        final_status = job_info.status.stage
+        # Fetch final job status — retry briefly if still RUNNING
+        # (the API may lag a few seconds behind the log stream ending)
+        final_status = "UNKNOWN"
+        for _ in range(6):
+            job_info = await _async_call(
+                self.api.inspect_job, job_id=job_id, namespace=namespace
+            )
+            final_status = job_info.status.stage
+            if final_status in terminal_states:
+                break
+            await asyncio.sleep(2.5)
 
         return final_status, all_logs
 
@@ -548,6 +554,20 @@ class HfJobsTool:
                 namespace=self.namespace,
             )
 
+            # Notify frontend of final status
+            if self.session and self.tool_call_id:
+                await self.session.send_event(
+                    Event(
+                        event_type="tool_state_change",
+                        data={
+                            "tool_call_id": self.tool_call_id,
+                            "tool": "hf_jobs",
+                            "state": final_status.lower(),
+                            "jobUrl": job.url,
+                        },
+                    )
+                )
+
             # Filter out UV package installation output
             filtered_logs = _filter_uv_install_output(all_logs)
 

agent/tools/sandbox_client.py

@@ -41,7 +41,7 @@ import sys
 import time
 import uuid
 from dataclasses import dataclass, field
-from typing import Any
+from typing import Any, Callable
 
 import httpx
 from huggingface_hub import CommitOperationAdd, HfApi
@@ -265,6 +265,7 @@ class Sandbox:
         sleep_time: int | None = None,
         token: str | None = None,
         wait_timeout: int = WAIT_TIMEOUT,
+        log: "Callable[[str], object] | None" = None,
     ) -> Sandbox:
         """
         Create a new sandbox by duplicating the template Space.
@@ -286,13 +287,14 @@ class Sandbox:
         Returns:
             A Sandbox instance connected to the running Space.
         """
+        _log = log or print
         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})...")
+        _log(f"Creating sandbox: {space_id} (from {template})...")
 
         kwargs = {
             "from_id": template,
@@ -304,25 +306,25 @@ class Sandbox:
             kwargs["sleep_time"] = sleep_time
 
         api.duplicate_space(**kwargs)
-        print(f"Space created: https://huggingface.co/spaces/{space_id}")
+        _log(f"Space created: https://huggingface.co/spaces/{space_id}")
 
         # Upload sandbox server and Dockerfile (triggers rebuild)
-        cls._setup_server(space_id, api)
+        cls._setup_server(space_id, api, log=_log)
 
         # Wait for it to come online (rebuild + start)
-        print(f"Waiting for Space to start (timeout: {wait_timeout}s)...")
+        _log(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})")
+                _log(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}...")
+            _log(f"  {runtime.stage}...")
             time.sleep(WAIT_INTERVAL)
         else:
             raise TimeoutError(
@@ -333,17 +335,17 @@ class Sandbox:
         # 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)
+            sb._wait_for_api(timeout=API_WAIT_TIMEOUT, log=_log)
         except TimeoutError as e:
-            print(
+            _log(
                 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:
+    def _setup_server(space_id: str, api: HfApi, *, log: Callable[[str], object] = print) -> None:
         """Upload embedded sandbox server + Dockerfile to the Space (single commit)."""
-        print(f"Uploading sandbox server to {space_id}...")
+        log(f"Uploading sandbox server to {space_id}...")
         api.create_commit(
             repo_id=space_id,
             repo_type="space",
@@ -359,7 +361,7 @@ class Sandbox:
             ],
             commit_message="Setup sandbox server",
         )
-        print("Server files uploaded, rebuild triggered.")
+        log("Server files uploaded, rebuild triggered.")
 
     @classmethod
     def connect(cls, space_id: str, *, token: str | None = None) -> Sandbox:
@@ -372,7 +374,7 @@ class Sandbox:
         sb._wait_for_api(timeout=60)
         return sb
 
-    def _wait_for_api(self, timeout: int = API_WAIT_TIMEOUT):
+    def _wait_for_api(self, timeout: int = API_WAIT_TIMEOUT, log: Callable[[str], object] = print):
         """Poll the health endpoint until the server responds."""
         deadline = time.time() + timeout
         last_err = None
@@ -382,7 +384,7 @@ class Sandbox:
                 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}")
+                    log(f"API is responsive at {self._base_url}")
                     return
             except Exception as e:
                 last_err = e

agent/tools/sandbox_tool.py

@@ -89,7 +89,16 @@ async def _ensure_sandbox(
         )
     )
 
-    kwargs = {"owner": owner, "hardware": hardware, "token": token, **create_kwargs}
+    # Thread-safe log callback: posts tool_log events from the worker thread
+    loop = asyncio.get_running_loop()
+
+    def _log(msg: str) -> None:
+        loop.call_soon_threadsafe(
+            session.event_queue.put_nowait,
+            Event(event_type="tool_log", data={"tool": "sandbox", "log": msg}),
+        )
+
+    kwargs = {"owner": owner, "hardware": hardware, "token": token, "log": _log, **create_kwargs}
     if hardware != "cpu-basic":
         kwargs["sleep_time"] = 1500
     sb = await asyncio.to_thread(Sandbox.create, **kwargs)
@@ -122,6 +131,10 @@ SANDBOX_CREATE_TOOL_SPEC = {
         "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"
+        "Before choosing hardware, estimate your VRAM needs (models you run, training data size). Rule of thumb: bf16/fp16 ≈ 2 bytes/param, "
+        "fp32 ≈ 4 bytes/param, plus ~20% overhead for optimizer states during training.\n"
+        "Common picks: t4-small (16GB VRAM, fits ≤1-3B), a10g-small (24GB, ≤7B), a100-large (80GB, ≤30B). "
+        "If the model won't fit, pick larger hardware upfront — OOM on a sandbox wastes time.\n\n"
         "Hardware: " + ", ".join([e.value for e in SpaceHardware]) + ".\n"
     ),
     "parameters": {

frontend/src/components/Chat/ActivityStatusBar.tsx

@@ -9,8 +9,8 @@ const shimmer = keyframes`
 `;
 
 const TOOL_LABELS: Record<string, string> = {
-  sandbox_create: 'Creating sandbox, this might take 1-2 minutes',
-  hf_jobs: 'Running job',
+  sandbox_create: 'Creating sandbox for code development, this might take 1-2 minutes',
+  hf_jobs: 'Running a GPU job, this might take a while',
   hf_repo_files: 'Uploading file',
   hf_repo_git: 'Git operation',
   hf_inspect_dataset: 'Inspecting dataset',

frontend/src/hooks/useAgentChat.ts

@@ -98,12 +98,13 @@ export function useAgentChat({ sessionId, isActive, onReady, onError, onSessionD
       },
       onToolLog: (tool: string, log: string) => {
         if (!isActiveRef.current) return;
-        if (tool === 'hf_jobs') {
+        if (tool === 'hf_jobs' || tool === 'sandbox') {
           const state = useAgentStore.getState();
           const existingOutput = state.panelData?.output?.content || '';
+          const header = tool === 'sandbox' ? '--- Sandbox creation ---' : '--- Job execution started ---';
           const newContent = existingOutput
             ? existingOutput + '\n' + log
-            : '--- Job execution started ---\n' + log;
+            : header + '\n' + log;
 
           setPanelOutput({ content: newContent, language: 'text' });