feat: stream sandbox bash output to frontend in real-time

Add streaming bash endpoint to sandbox server (NDJSON over HTTP),
a bash_stream() generator on the client, and a streaming tool handler
that emits tool_log events per line — same pattern as HF Jobs.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

agent/tools/sandbox_client.py

@@ -37,11 +37,12 @@ Tools: bash, read, write, edit, upload
 from __future__ import annotations
 
 import io
+import json
 import sys
 import time
 import uuid
 from dataclasses import dataclass, field
-from typing import Any, Callable
+from typing import Any, Callable, Generator
 
 import httpx
 from huggingface_hub import CommitOperationAdd, HfApi
@@ -97,8 +98,9 @@ CMD ["python", "sandbox_server.py"]
 
 _SANDBOX_SERVER = '''\
 """Minimal FastAPI server for sandbox operations."""
-import os, subprocess, pathlib
+import os, subprocess, pathlib, asyncio, json
 from fastapi import FastAPI
+from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
 from typing import Optional
 import uvicorn
@@ -148,6 +150,41 @@ def bash(req: BashReq):
     except Exception as e:
         return {"success": False, "output": "", "error": str(e)}
 
+@app.post("/api/bash/stream")
+async def bash_stream(req: BashReq):
+    """Stream bash output line-by-line as NDJSON."""
+    async def generate():
+        try:
+            proc = await asyncio.create_subprocess_shell(
+                req.command,
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.STDOUT,
+                cwd=req.work_dir,
+            )
+            total_chars = 0
+            try:
+                line_bytes = await asyncio.wait_for(proc.stdout.readline(), timeout=req.timeout)
+                while line_bytes:
+                    line = line_bytes.decode(errors="replace")
+                    total_chars += len(line)
+                    if total_chars <= 30000:
+                        yield json.dumps({"type": "output", "data": line}) + "\\n"
+                    elif total_chars - len(line) <= 30000:
+                        yield json.dumps({"type": "output", "data": "... (truncated)\\n"}) + "\\n"
+                    line_bytes = await asyncio.wait_for(proc.stdout.readline(), timeout=req.timeout)
+            except asyncio.TimeoutError:
+                try:
+                    proc.kill()
+                except ProcessLookupError:
+                    pass
+                yield json.dumps({"type": "error", "data": f"Timeout after {req.timeout}s"}) + "\\n"
+                return
+            rc = await proc.wait()
+            yield json.dumps({"type": "exit", "code": rc}) + "\\n"
+        except Exception as e:
+            yield json.dumps({"type": "error", "data": str(e)}) + "\\n"
+    return StreamingResponse(generate(), media_type="application/x-ndjson")
+
 @app.post("/api/read")
 def read(req: ReadReq):
     try:
@@ -492,6 +529,58 @@ class Sandbox:
             timeout=timeout,
         )
 
+    def bash_stream(
+        self,
+        command: str,
+        *,
+        work_dir: str | None = None,
+        timeout: int | None = None,
+    ) -> Generator[tuple[str, str | int], None, None]:
+        """Stream bash output line-by-line from the sandbox.
+
+        Yields:
+            (event_type, data) tuples:
+              - ("output", line_text)
+              - ("error", error_message)
+              - ("exit", return_code)
+        """
+        effective_timeout = min(timeout or self.timeout, MAX_TIMEOUT)
+        payload = {
+            "command": command,
+            "work_dir": work_dir or self.work_dir,
+            "timeout": effective_timeout,
+        }
+        try:
+            with self._client.stream(
+                "POST",
+                "bash/stream",
+                json=payload,
+                timeout=httpx.Timeout(effective_timeout + 30, connect=30),
+            ) as resp:
+                if resp.status_code != 200:
+                    yield ("error", f"HTTP {resp.status_code}")
+                    return
+                for raw_line in resp.iter_lines():
+                    if not raw_line:
+                        continue
+                    try:
+                        msg = json.loads(raw_line)
+                    except json.JSONDecodeError:
+                        continue
+                    msg_type = msg.get("type", "")
+                    if msg_type == "output":
+                        yield ("output", msg.get("data", ""))
+                    elif msg_type == "exit":
+                        yield ("exit", msg.get("code", -1))
+                    elif msg_type == "error":
+                        yield ("error", msg.get("data", "unknown error"))
+        except httpx.TimeoutException:
+            yield ("error", f"Timeout after {effective_timeout}s")
+        except httpx.ConnectError:
+            yield ("error", f"Cannot connect to sandbox. Is {self.space_id} running?")
+        except Exception as e:
+            yield ("error", str(e))
+
     def read(
         self, path: str, *, offset: int | None = None, limit: int | None = None
     ) -> ToolResult:

agent/tools/sandbox_tool.py

@@ -190,8 +190,77 @@ async def sandbox_create_handler(
     ), True
 
 
+def _make_bash_streaming_handler():
+    """Create a bash handler that streams output via tool_log events."""
+
+    async def handler(args: dict[str, Any], session: Any = None) -> tuple[str, bool]:
+        if not session or not getattr(session, "sandbox", None):
+            return "No sandbox running. Call sandbox_create first to start one.", False
+
+        sb = session.sandbox
+        queue: asyncio.Queue = asyncio.Queue()
+        loop = asyncio.get_running_loop()
+
+        def producer():
+            try:
+                for event_type, data in sb.bash_stream(
+                    args["command"],
+                    work_dir=args.get("work_dir"),
+                    timeout=args.get("timeout"),
+                ):
+                    loop.call_soon_threadsafe(queue.put_nowait, (event_type, data))
+                loop.call_soon_threadsafe(queue.put_nowait, None)
+            except Exception as e:
+                loop.call_soon_threadsafe(queue.put_nowait, ("error", str(e)))
+                loop.call_soon_threadsafe(queue.put_nowait, None)
+
+        fut = loop.run_in_executor(None, producer)
+
+        all_output: list[str] = []
+        exit_code = 0
+        error_msg = ""
+
+        while True:
+            item = await queue.get()
+            if item is None:
+                break
+
+            event_type, data = item
+            if event_type == "output":
+                all_output.append(data)
+                if session:
+                    await session.send_event(
+                        Event(
+                            event_type="tool_log",
+                            data={"tool": "bash", "log": data.rstrip("\n")},
+                        )
+                    )
+            elif event_type == "exit":
+                exit_code = data
+            elif event_type == "error":
+                error_msg = data
+
+        await fut
+
+        output = "".join(all_output)
+        if not output and not error_msg:
+            output = "(no output)"
+
+        if error_msg:
+            if output:
+                return f"{output}\nERROR: {error_msg}", False
+            return f"ERROR: {error_msg}", False
+
+        if exit_code != 0:
+            return f"{output}\nExit code {exit_code}", False
+
+        return output, True
+
+    return handler
+
+
 def _make_tool_handler(sandbox_tool_name: str):
-    """Factory: create a handler for a sandbox operation tool."""
+    """Factory: create a handler for a sandbox operation tool (non-bash)."""
 
     async def handler(args: dict[str, Any], session: Any = None) -> tuple[str, bool]:
         # Require sandbox to exist — user must approve sandbox_create first
@@ -235,12 +304,17 @@ def get_sandbox_tools():
     # Operation tools (auto-execute, no approval needed)
     for name in Sandbox.TOOLS.keys():
         spec = Sandbox.TOOLS[name]
+        # Bash uses streaming handler; other tools use standard handler
+        if name == "bash":
+            handler = _make_bash_streaming_handler()
+        else:
+            handler = _make_tool_handler(name)
         tools.append(
             ToolSpec(
                 name=name,
                 description=spec["description"],
                 parameters=spec["parameters"],
-                handler=_make_tool_handler(name),
+                handler=handler,
             )
         )
 

tests/integration/tools/test_sandbox_streaming.py

@@ -0,0 +1,267 @@
+#!/usr/bin/env python3
+"""
+Integration test for sandbox bash streaming.
+
+Creates a real sandbox on HF, runs commands, verifies that:
+  1. bash_stream() yields lines incrementally
+  2. The streaming tool handler emits tool_log events
+  3. The final output is collected correctly
+  4. Error/exit codes propagate
+
+Requires HF_TOKEN in environment.
+"""
+import asyncio
+import os
+import sys
+
+sys.path.insert(0, ".")
+
+from agent.tools.sandbox_client import Sandbox
+
+GREEN = "\033[92m"
+RED = "\033[91m"
+BLUE = "\033[94m"
+YELLOW = "\033[93m"
+RESET = "\033[0m"
+
+
+def ok(msg):
+    print(f"{GREEN}  OK{RESET} {msg}")
+
+
+def fail(msg):
+    print(f"{RED}FAIL{RESET} {msg}")
+
+
+def info(msg):
+    print(f"{BLUE}INFO{RESET} {msg}")
+
+
+async def test_bash_stream_basic(sb: Sandbox):
+    """Test that bash_stream yields output lines and an exit event."""
+    info("bash_stream: echo test")
+    lines = []
+    exit_code = None
+    for event_type, data in sb.bash_stream("echo hello && echo world"):
+        if event_type == "output":
+            lines.append(data.rstrip("\n"))
+        elif event_type == "exit":
+            exit_code = data
+
+    if lines == ["hello", "world"]:
+        ok(f"Got expected lines: {lines}")
+    else:
+        fail(f"Unexpected lines: {lines}")
+        return False
+
+    if exit_code == 0:
+        ok(f"Exit code: {exit_code}")
+    else:
+        fail(f"Unexpected exit code: {exit_code}")
+        return False
+
+    return True
+
+
+async def test_bash_stream_multiline(sb: Sandbox):
+    """Test streaming a script that prints multiple lines with delays."""
+    info("bash_stream: multi-line with sleep")
+    lines = []
+    for event_type, data in sb.bash_stream(
+        'for i in 1 2 3; do echo "line $i"; sleep 0.2; done'
+    ):
+        if event_type == "output":
+            line = data.rstrip("\n")
+            lines.append(line)
+            info(f"  streamed: {line}")
+
+    if len(lines) == 3 and lines == ["line 1", "line 2", "line 3"]:
+        ok(f"Got all 3 lines incrementally")
+    else:
+        fail(f"Unexpected lines: {lines}")
+        return False
+
+    return True
+
+
+async def test_bash_stream_stderr(sb: Sandbox):
+    """Test that stderr is also streamed (merged into stdout)."""
+    info("bash_stream: stderr output")
+    lines = []
+    for event_type, data in sb.bash_stream("echo out && echo err >&2"):
+        if event_type == "output":
+            lines.append(data.rstrip("\n"))
+
+    if "out" in lines and "err" in lines:
+        ok(f"Both stdout and stderr captured: {lines}")
+    else:
+        fail(f"Missing output: {lines}")
+        return False
+
+    return True
+
+
+async def test_bash_stream_exit_code(sb: Sandbox):
+    """Test that non-zero exit codes are reported."""
+    info("bash_stream: non-zero exit code")
+    exit_code = None
+    for event_type, data in sb.bash_stream("exit 42"):
+        if event_type == "exit":
+            exit_code = data
+
+    if exit_code == 42:
+        ok(f"Got expected exit code: {exit_code}")
+    else:
+        fail(f"Unexpected exit code: {exit_code}")
+        return False
+
+    return True
+
+
+async def test_bash_stream_empty(sb: Sandbox):
+    """Test command with no output."""
+    info("bash_stream: no output command")
+    lines = []
+    exit_code = None
+    for event_type, data in sb.bash_stream("true"):
+        if event_type == "output":
+            lines.append(data)
+        elif event_type == "exit":
+            exit_code = data
+
+    if exit_code == 0:
+        ok(f"Exit code 0, output lines: {len(lines)}")
+    else:
+        fail(f"Unexpected exit code: {exit_code}")
+        return False
+
+    return True
+
+
+async def test_bash_stream_tool_handler(sb: Sandbox):
+    """Test the full tool handler path with mock session that captures events."""
+    info("Tool handler: streaming bash with event capture")
+
+    from agent.core.session import Event
+    from agent.tools.sandbox_tool import _make_bash_streaming_handler
+
+    # Mock session with event capture
+    class MockSession:
+        def __init__(self, sandbox):
+            self.sandbox = sandbox
+            self.event_queue = asyncio.Queue()
+            self.events = []
+
+        async def send_event(self, event: Event):
+            self.events.append(event)
+            await self.event_queue.put(event)
+
+    session = MockSession(sb)
+    handler = _make_bash_streaming_handler()
+
+    output, success = await handler(
+        {"command": 'for i in a b c; do echo "$i"; sleep 0.1; done'},
+        session=session,
+    )
+
+    # Check tool_log events were emitted
+    log_events = [e for e in session.events if e.event_type == "tool_log"]
+    log_lines = [e.data["log"] for e in log_events]
+
+    if log_lines == ["a", "b", "c"]:
+        ok(f"Tool handler emitted {len(log_events)} tool_log events: {log_lines}")
+    else:
+        fail(f"Unexpected log events: {log_lines}")
+        return False
+
+    if success and "a" in output and "b" in output and "c" in output:
+        ok(f"Handler returned success with full output")
+    else:
+        fail(f"Handler returned success={success}, output={output!r}")
+        return False
+
+    return True
+
+
+async def test_original_bash_still_works(sb: Sandbox):
+    """Verify the non-streaming bash endpoint still works."""
+    info("Original bash endpoint (non-streaming)")
+    result = await asyncio.to_thread(sb.bash, "echo legacy_test")
+    if result.success and "legacy_test" in result.output:
+        ok(f"Original bash works: {result.output.strip()}")
+    else:
+        fail(f"Original bash failed: {result}")
+        return False
+
+    return True
+
+
+async def main():
+    print("=" * 60)
+    print(f"{BLUE}Sandbox Bash Streaming — Integration Tests{RESET}")
+    print("=" * 60)
+
+    token = os.environ.get("HF_TOKEN")
+    if not token:
+        fail("HF_TOKEN not set")
+        sys.exit(1)
+
+    from huggingface_hub import HfApi
+
+    api = HfApi(token=token)
+    owner = api.whoami().get("name", "")
+    info(f"HF user: {owner}")
+
+    sb = None
+    try:
+        info("Creating sandbox (this takes ~2-4 minutes)...")
+        sb = Sandbox.create(owner=owner, token=token, hardware="cpu-basic")
+        ok(f"Sandbox ready: {sb.space_id}")
+        print()
+
+        tests = [
+            test_bash_stream_basic,
+            test_bash_stream_multiline,
+            test_bash_stream_stderr,
+            test_bash_stream_exit_code,
+            test_bash_stream_empty,
+            test_bash_stream_tool_handler,
+            test_original_bash_still_works,
+        ]
+
+        passed = 0
+        failed = 0
+        for test in tests:
+            print()
+            try:
+                result = await test(sb)
+                if result:
+                    passed += 1
+                else:
+                    failed += 1
+            except Exception as e:
+                fail(f"{test.__name__}: {e}")
+                import traceback
+                traceback.print_exc()
+                failed += 1
+
+        print()
+        print("=" * 60)
+        if failed == 0:
+            print(f"{GREEN}All {passed} tests passed!{RESET}")
+        else:
+            print(f"{RED}{failed} failed{RESET}, {GREEN}{passed} passed{RESET}")
+        print("=" * 60)
+
+    finally:
+        if sb and sb._owns_space:
+            info(f"Cleaning up sandbox: {sb.space_id}")
+            try:
+                sb.delete()
+                ok("Sandbox deleted")
+            except Exception as e:
+                fail(f"Cleanup failed: {e}")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())