Upload folder using huggingface_hub

server/app.py

@@ -4,6 +4,11 @@ Mirrors the ``e2b_desktop`` pattern: pass a ``gradio_builder`` to
 ``create_app`` and let OpenEnv handle the Gradio mount (including the
 HF-Space-compatible ``/web`` path). No manual ``gr.mount_gradio_app``.
 
+Also mounts a bespoke SSE endpoint at ``GET /rollouts/{rollout_id}/events``
+that multiplexes opencode serve's ``/event`` stream with our proxy's
+per-turn frames. MCP tools don't support streaming; this gives the UI
+and interactive clients a live feed.
+
 Usage::
 
     # Development:
@@ -69,6 +74,165 @@ app = create_app(
 )
 
 
+def _find_active_environment(request):
+    """Locate the currently-active OpenCodeEnvironment instance for a request.
+
+    ``create_app`` keeps per-session envs behind the scenes; for the SSE
+    endpoint we just grab the most recent one (single-worker Space), so
+    we poke at ``app.state.env_cache`` and fall back to ``web_manager``.
+    """
+    cache = getattr(app.state, "env_cache", None)
+    if cache:
+        try:
+            return next(iter(cache.values()))
+        except StopIteration:
+            pass
+    try:
+        return _web_manager.get_environment()  # type: ignore[name-defined]
+    except Exception:
+        return None
+
+
+@app.get("/rollouts/{rollout_id}/events")
+async def rollout_events(rollout_id: str):
+    """Server-Sent Events feed for a rollout started via ``start_rollout``.
+
+    Merges two streams:
+
+    1. opencode serve's ``GET /event`` (session-level events: message
+       parts, tool calls, idle/abort markers) — forwarded as-is.
+    2. our proxy's ``proxy_trace.jsonl`` in the sandbox (per-turn
+       LLM turns + logprobs) — tailed and emitted as
+       ``{type: "proxy.turn", turn, tokens, logprobs, finish_reason, ...}``.
+
+    Terminates on a final ``{"type": "rollout.done", ...}`` frame once the
+    session has idled or erred.
+    """
+    import asyncio
+    import json as _json
+
+    from starlette.responses import StreamingResponse
+
+    env = _find_active_environment(None)
+    if env is None:
+        return StreamingResponse(
+            iter([f"data: {_json.dumps({'type': 'error', 'reason': 'env not found'})}\n\n"]),
+            media_type="text/event-stream",
+        )
+
+    registry = getattr(env, "_registry", None)
+    handle = registry.get(rollout_id) if registry else None
+    if handle is None:
+        async def _single_error():
+            yield (
+                "data: "
+                + _json.dumps({"type": "error", "rollout_id": rollout_id, "reason": "unknown rollout"})
+                + "\n\n"
+            )
+        return StreamingResponse(_single_error(), media_type="text/event-stream")
+
+    async def _gen():
+        # Wait briefly for the serve client to be wired by the worker.
+        for _ in range(60):
+            if handle.session is not None and getattr(handle.session, "serve_client", None):
+                break
+            if handle.is_done():
+                break
+            await asyncio.sleep(0.25)
+        session = handle.session
+        if session is None:
+            yield (
+                "data: "
+                + _json.dumps({
+                    "type": "error",
+                    "rollout_id": rollout_id,
+                    "reason": "session never created",
+                    "detail": handle.error,
+                })
+                + "\n\n"
+            )
+            return
+
+        sandbox = session.sandbox
+        proxy_trace_path = session._proxy_trace_path
+        serve_client = getattr(session, "serve_client", None)
+
+        # Task A: forward opencode serve events.
+        serve_q: asyncio.Queue = asyncio.Queue()
+
+        async def forward_serve():
+            if serve_client is None:
+                return
+            try:
+                async for ev in serve_client.astream_events():
+                    await serve_q.put({"source": "serve", **ev})
+                    if handle.is_done():
+                        break
+            except Exception as exc:  # noqa: BLE001
+                await serve_q.put({"source": "serve", "type": "error", "reason": str(exc)})
+            finally:
+                await serve_q.put(None)
+
+        # Task B: tail proxy trace file (incremental) from the sandbox.
+        async def tail_proxy():
+            last_len = 0
+            while not handle.is_done():
+                try:
+                    if proxy_trace_path:
+                        content = sandbox.read_text(proxy_trace_path) or ""
+                        if len(content) > last_len:
+                            new = content[last_len:]
+                            last_len = len(content)
+                            for line in new.splitlines():
+                                line = line.strip()
+                                if not line:
+                                    continue
+                                try:
+                                    turn = _json.loads(line)
+                                except Exception:
+                                    continue
+                                await serve_q.put({
+                                    "source": "proxy",
+                                    "type": "proxy.turn",
+                                    "turn": turn.get("turn"),
+                                    "finish_reason": turn.get("finish_reason"),
+                                    "n_tokens": len(turn.get("completion_tokens") or []),
+                                    "first_tokens": (turn.get("completion_tokens") or [])[:6],
+                                    "first_logps": (turn.get("per_token_logps") or [])[:6],
+                                    "latency_s": turn.get("latency_s"),
+                                })
+                except Exception:
+                    pass
+                await asyncio.sleep(1.0)
+
+        t_serve = asyncio.create_task(forward_serve())
+        t_proxy = asyncio.create_task(tail_proxy())
+
+        try:
+            while True:
+                try:
+                    ev = await asyncio.wait_for(serve_q.get(), timeout=1.0)
+                except asyncio.TimeoutError:
+                    ev = None
+                if ev is None:
+                    if handle.is_done():
+                        break
+                    continue
+                yield "data: " + _json.dumps(ev) + "\n\n"
+        finally:
+            t_serve.cancel()
+            t_proxy.cancel()
+
+        yield "data: " + _json.dumps({
+            "source": "server",
+            "type": "rollout.done",
+            "rollout_id": rollout_id,
+            "error": handle.error,
+        }) + "\n\n"
+
+    return StreamingResponse(_gen(), media_type="text/event-stream")
+
+
 def main(host: str = "0.0.0.0", port: int = 8000) -> None:
     import uvicorn
 

server/opencode_environment.py

@@ -21,6 +21,7 @@ from __future__ import annotations
 
 import json
 import os
+import threading
 import time
 from typing import Any, Optional
 from uuid import uuid4
@@ -50,17 +51,87 @@ WORKDIR_PATH = "/home/user/workdir"
 VERIFIER_TIMEOUT_S = 120
 
 
+class _RolloutRegistry:
+    """Per-environment bookkeeping for in-flight non-blocking rollouts.
+
+    Every ``start_rollout`` call spawns a background thread that owns a
+    session through its full lifecycle (create sandbox → install opencode
+    → start proxy + opencode serve → fire instruction → wait for agent →
+    finalize on demand). The registry keeps one :class:`_RolloutHandle`
+    per ``rollout_id`` so subsequent ``subscribe_events`` / ``abort_rollout``
+    / ``finalize`` tool calls can find it.
+    """
+
+    def __init__(self) -> None:
+        self._lock = threading.Lock()
+        self._handles: dict[str, "_RolloutHandle"] = {}
+
+    def add(self, handle: "_RolloutHandle") -> None:
+        with self._lock:
+            self._handles[handle.rollout_id] = handle
+
+    def get(self, rollout_id: str) -> "_RolloutHandle | None":
+        with self._lock:
+            return self._handles.get(rollout_id)
+
+    def drop(self, rollout_id: str) -> None:
+        with self._lock:
+            self._handles.pop(rollout_id, None)
+
+
+class _RolloutHandle:
+    """Everything the server needs to finalize / observe / abort a rollout."""
+
+    def __init__(
+        self,
+        rollout_id: str,
+        task_id: str,
+        session_factory_kwargs: dict[str, Any],
+        task: Any,
+    ) -> None:
+        self.rollout_id = rollout_id
+        self.task_id = task_id
+        self._kwargs = session_factory_kwargs
+        self._task = task
+
+        # Set by the worker thread as the session progresses.
+        self.session: Any | None = None
+        self.error: str | None = None
+        self.started_at = time.time()
+        self.finished_at: float | None = None
+        self._done = threading.Event()
+        self._worker: threading.Thread | None = None
+        self._verifier_stdout = ""
+        self._verifier_stderr = ""
+        self._test_exit_code: int | None = None
+        self._test_script: str = ""
+
+    def is_done(self) -> bool:
+        return self._done.is_set()
+
+    def wait(self, timeout: float | None = None) -> bool:
+        return self._done.wait(timeout)
+
+
 class OpenCodeEnvironment(MCPEnvironment):
-    """One-tool MCP environment.
+    """Two-tier MCP environment for OpenCode rollouts.
 
-    The single tool ``run_rollout`` is synchronous and returns a JSON string
-    — one call = one complete agent rollout. Each call creates and destroys
-    its own E2B sandbox, so the environment is stateless across calls.
+    The **macro tool** ``run_rollout`` runs a blocking end-to-end rollout
+    and returns the final :class:`RolloutResult` (kept for training, which
+    doesn't benefit from per-turn streaming).
+
+    The **fine-grained tools** ``start_rollout`` / ``subscribe_events``
+    (server-side SSE via a helper endpoint) / ``abort_rollout`` / ``finalize``
+    are meant for the Gradio UI and interactive inspection — they expose
+    the PR #471-style session primitives directly.
     """
 
     SUPPORTS_CONCURRENT_SESSIONS = True
 
     def __init__(self) -> None:
+        # Per-environment rollout registry (one per HTTP session/worker).
+        self._registry = _RolloutRegistry()
+
         # Import inside __init__ to keep module import cheap and to allow
         # patching for tests. Dual-import pattern: package when installed,
         # flat when run directly out of the repo via ``server.app:app``.
@@ -151,6 +222,143 @@ class OpenCodeEnvironment(MCPEnvironment):
                 agent_timeout_s=agent_timeout_s,
             )
 
+        # ── Fine-grained tools (non-blocking) — Phase 2b ────────────────────
+        #
+        # These wrap the primitive's ``driver="serve"`` path so the UI and
+        # interactive clients can observe a rollout in flight. Training keeps
+        # using the blocking ``run_rollout`` tool above.
+
+        @mcp.tool
+        def start_rollout(
+            vllm_url: str,
+            model: str,
+            instruction: str,
+            test_script: str = "",
+            task_id: str = "",
+            setup_shell: str = "",
+            upload_files: Optional[dict[str, str]] = None,
+            provider: str = "openai_compatible",
+            api_key: str = "intercepted",
+            mode: str = "transparent_proxy",
+            disable_thinking: bool = False,
+            max_tokens_cap: int = 4096,
+            agent_timeout_s: float = 600.0,
+        ) -> str:
+            """Start a rollout asynchronously; return a ``rollout_id`` immediately.
+
+            Spawns a background worker that creates the sandbox, installs
+            opencode, boots ``opencode serve``, and fires the instruction.
+            The caller then uses ``subscribe_events`` / ``get_state`` /
+            ``abort_rollout`` / ``finalize`` with the returned id.
+            """
+            rid = uuid4().hex[:12]
+            handle = self._spawn_async_rollout(
+                rollout_id=rid,
+                vllm_url=vllm_url,
+                model=model,
+                instruction=instruction,
+                test_script=test_script,
+                task_id=task_id,
+                setup_shell=setup_shell,
+                upload_files=upload_files or {},
+                provider=provider,
+                api_key=api_key,
+                mode=mode,
+                disable_thinking=disable_thinking,
+                max_tokens_cap=max_tokens_cap,
+                agent_timeout_s=agent_timeout_s,
+            )
+            return json.dumps({
+                "rollout_id": handle.rollout_id,
+                "task_id": handle.task_id,
+                "status": "starting",
+                "started_at": handle.started_at,
+            })
+
+        @mcp.tool
+        def get_state(rollout_id: str) -> str:
+            """Return a small JSON snapshot of the rollout's current state.
+
+            Useful for UI polling between SSE reconnects. The payload mirrors
+            what an ``agent_done`` event would carry plus a ``status`` field.
+            """
+            handle = self._registry.get(rollout_id)
+            if handle is None:
+                return json.dumps({"rollout_id": rollout_id, "status": "unknown"})
+            session = handle.session
+            serve_session_id = getattr(session, "serve_session_id", None)
+            turns_n = 0
+            if session is not None and session._proxy_trace_path:
+                try:
+                    content = session.sandbox.read_text(session._proxy_trace_path)
+                    turns_n = sum(1 for line in content.splitlines() if line.strip())
+                except Exception:
+                    pass
+            return json.dumps({
+                "rollout_id": rollout_id,
+                "task_id": handle.task_id,
+                "status": "done" if handle.is_done() else "running",
+                "serve_session_id": serve_session_id,
+                "proxy_turns_so_far": turns_n,
+                "error": handle.error,
+                "started_at": handle.started_at,
+                "finished_at": handle.finished_at,
+            })
+
+        @mcp.tool
+        def abort_rollout(rollout_id: str) -> str:
+            """Cancel an in-flight rollout.
+
+            Calls ``POST /session/:id/abort`` on the sandbox's opencode serve
+            (which stops token generation immediately), then lets the worker
+            unwind and close the sandbox. Returns a small JSON payload.
+            """
+            handle = self._registry.get(rollout_id)
+            if handle is None:
+                return json.dumps({"rollout_id": rollout_id, "aborted": False, "reason": "unknown"})
+            aborted = False
+            if handle.session is not None:
+                try:
+                    aborted = bool(handle.session.abort())
+                except Exception as exc:  # noqa: BLE001
+                    return json.dumps({
+                        "rollout_id": rollout_id, "aborted": False, "reason": str(exc),
+                    })
+            return json.dumps({"rollout_id": rollout_id, "aborted": aborted})
+
+        @mcp.tool
+        def finalize_rollout(rollout_id: str, wait_s: float = 300.0) -> str:
+            """Block until the rollout's worker finishes, then return its
+            full :class:`RolloutResult` as JSON (same shape as ``run_rollout``).
+
+            Must be called exactly once per rollout — this also de-registers
+            the handle and closes the sandbox.
+            """
+            handle = self._registry.get(rollout_id)
+            if handle is None:
+                return json.dumps({
+                    "task_id": "", "sandbox_id": "", "reward": None,
+                    "exit_code": 0, "wall_s": 0.0, "mode": "",
+                    "proxy_turns": [], "workdir_files": {},
+                    "agent_log_tail": "", "verifier_stdout": "",
+                    "verifier_stderr": "", "test_exit_code": None,
+                    "error": f"unknown rollout_id: {rollout_id}",
+                    "proxy_log_tail": "", "install_log_tail": "",
+                })
+            if not handle.wait(timeout=wait_s):
+                return json.dumps({
+                    "task_id": handle.task_id, "sandbox_id": "",
+                    "reward": None, "exit_code": -1, "wall_s": 0.0,
+                    "mode": "", "proxy_turns": [], "workdir_files": {},
+                    "agent_log_tail": "", "verifier_stdout": "",
+                    "verifier_stderr": "", "test_exit_code": None,
+                    "error": f"timeout waiting for rollout_id={rollout_id}",
+                    "proxy_log_tail": "", "install_log_tail": "",
+                })
+            payload = self._finalize_handle(handle)
+            self._registry.drop(rollout_id)
+            return payload
+
         super().__init__(mcp)
 
     # ── OpenEnv lifecycle ──���────────────────────────────────────────────────
@@ -397,6 +605,150 @@ def _clamp_turn(turn: dict[str, Any]) -> dict[str, Any]:
     return out
 
 
+    # ── Async rollout plumbing ─────────────────────────────────────────────
+
+    def _spawn_async_rollout(
+        self,
+        *,
+        rollout_id: str,
+        vllm_url: str,
+        model: str,
+        instruction: str,
+        test_script: str,
+        task_id: str,
+        setup_shell: str,
+        upload_files: dict[str, str],
+        provider: str,
+        api_key: str,
+        mode: str,
+        disable_thinking: bool,
+        max_tokens_cap: int,
+        agent_timeout_s: float,
+    ) -> _RolloutHandle:
+        from opencode_env import OpenCodeTask
+
+        # Build the task payload up-front; staging happens on the worker.
+        merged_uploads = dict(upload_files)
+        if test_script:
+            merged_uploads[REMOTE_TEST_PATH] = test_script
+        task = OpenCodeTask(
+            instruction=instruction,
+            setup_shell=setup_shell or None,
+            upload_files=merged_uploads,
+            metadata={"task_id": task_id},
+        )
+
+        provider_model = _qualify_model(provider, model)
+        config = self._OpenCodeConfig(
+            provider=provider,
+            base_url=vllm_url.rstrip("/"),
+            api_key=api_key,
+            model=provider_model,
+            agent_timeout_s=agent_timeout_s,
+            proxy_disable_thinking=disable_thinking,
+            proxy_max_tokens_cap=max_tokens_cap if max_tokens_cap > 0 else None,
+        )
+
+        handle = _RolloutHandle(
+            rollout_id=rollout_id,
+            task_id=task_id,
+            session_factory_kwargs={
+                "config": config,
+                "mode": mode,
+                "agent_timeout_s": agent_timeout_s,
+            },
+            task=task,
+        )
+        handle._test_script = test_script
+        self._registry.add(handle)
+
+        def worker() -> None:
+            try:
+                factory = self._OpenCodeSessionFactory(
+                    config=config,
+                    sandbox_backend=self._E2BSandboxBackend(),
+                    mode=mode,
+                    verifier=None,
+                    driver="serve",  # Phase 2b path
+                )
+                handle.session = factory.create(task=task)
+                # Block until the agent idles. The caller can abort via
+                # ``abort_rollout`` any time; that triggers the serve
+                # ``/abort`` endpoint and ``wait_for_completion`` returns.
+                try:
+                    handle.session.wait_for_completion(timeout_s=agent_timeout_s)
+                except Exception as exc:  # noqa: BLE001
+                    handle.error = f"wait_for_completion: {type(exc).__name__}: {exc}"
+            except Exception as exc:  # noqa: BLE001
+                handle.error = f"{type(exc).__name__}: {exc}"
+            finally:
+                handle.finished_at = time.time()
+                handle._done.set()
+
+        t = threading.Thread(target=worker, daemon=True, name=f"rollout-{rollout_id}")
+        t.start()
+        handle._worker = t
+        return handle
+
+    def _finalize_handle(self, handle: _RolloutHandle) -> str:
+        """Run the verifier (if test_script present), collect the trace, and
+        return a JSON-serialized :class:`RolloutResult` matching the shape
+        returned by ``run_rollout``. Closes the session + sandbox."""
+        result = self._result_cls(task_id=handle.task_id, mode=handle._kwargs.get("mode", ""))
+        session = handle.session
+        if session is None:
+            result.error = handle.error or "session never created"
+            return result.model_dump_json()
+
+        result.sandbox_id = session.sandbox.sandbox_id
+        result.exit_code = 0  # serve-driver has no exit code; use 0 unless aborted
+        wall_s = (handle.finished_at or time.time()) - handle.started_at
+        result.wall_s = round(wall_s, 3)
+
+        try:
+            if handle._test_script:
+                session.sandbox.exec(
+                    "mkdir -p /home/user/logs/verifier /home/user/tests && "
+                    f"chmod +x {REMOTE_TEST_PATH}",
+                    timeout=15,
+                )
+                v = session.sandbox.exec(
+                    f"bash {REMOTE_TEST_PATH}",
+                    cwd=WORKDIR_PATH,
+                    timeout=VERIFIER_TIMEOUT_S,
+                )
+                result.test_exit_code = int(v.exit_code)
+                result.verifier_stdout = (v.stdout or "")[:4000]
+                result.verifier_stderr = (v.stderr or "")[:2000]
+                result.reward = _read_reward(session.sandbox, REMOTE_REWARD_PATH)
+
+            summary = self._collect_rollout_summary(session)
+            result.agent_log_tail = _tail(summary.opencode_events, 20)
+            result.workdir_files = {
+                path: (contents or "")[:8000]
+                for path, contents in (summary.workdir_contents or {}).items()
+            }
+            for raw in summary.proxy_turns:
+                result.proxy_turns.append(self._turn_cls(**_clamp_turn(raw)))
+            result.proxy_log_tail = _read_safe(
+                session.sandbox, "/home/user/logs/agent/proxy.log", tail_chars=4000
+            )
+        except Exception as exc:  # noqa: BLE001
+            if not handle.error:
+                result.error = f"finalize: {type(exc).__name__}: {exc}"
+            else:
+                result.error = handle.error
+        finally:
+            try:
+                session.close()
+            except Exception:
+                pass
+
+        if handle.error and not result.error:
+            result.error = handle.error
+        return result.model_dump_json()
+
+
 def _tail(events: list[dict[str, Any]], n: int) -> str:
     """Return the last ``n`` opencode event lines as a newline-joined string."""
     if not events:

uv.lock

@@ -1710,7 +1710,7 @@ dev = [
 [[package]]
 name = "openenv-core"
 version = "0.2.3"
-source = { git = "https://github.com/adithya-s-k/OpenEnv.git?rev=opencode-harness#4f8bc23dc0e860ac990d8fd0ebdac7a5498e56a3" }
+source = { git = "https://github.com/adithya-s-k/OpenEnv.git?rev=opencode-harness#5d760274516f8c917212506240983e781b467c1d" }
 dependencies = [
     { name = "fastapi" },
     { name = "fastmcp" },
@@ -1741,7 +1741,7 @@ core = [
 [[package]]
 name = "openenv-opencode-env"
 version = "0.1.0"
-source = { git = "https://github.com/adithya-s-k/OpenEnv.git?subdirectory=envs%2Fopencode_env&rev=opencode-harness#4f8bc23dc0e860ac990d8fd0ebdac7a5498e56a3" }
+source = { git = "https://github.com/adithya-s-k/OpenEnv.git?subdirectory=envs%2Fopencode_env&rev=opencode-harness#5d760274516f8c917212506240983e781b467c1d" }
 dependencies = [
     { name = "e2b" },
     { name = "fastapi" },