openenv/harbor/serving.py

"""Serve Harbor tasks: Task API, MCP rollouts, and a human UI.

Locally this is two ports:

    :port          env server    Task API + MCP + UI   (faces the trainer / a browser)
    :capture_port  capture proxy                       (faces the sandbox, published)

Two ports on purpose. The sandbox is off-cluster and must reach the capture proxy over a public URL;
the env server has no business being publicly reachable, and sharing one port would expose it as
soon as the capture proxy became reachable.

On a hosted platform that inverts. A Space gets exactly one public URL and exposes one port, so
there is no second port to publish and nothing to forward. The capture app is mounted onto the env
server's own app at `CAPTURE_MOUNT` instead, and the sandbox reaches it at `<space-url>/capture`.
The proxy still refuses unregistered callers, which is what keeps a public mount from becoming an
open relay.
"""

from __future__ import annotations

import os
import threading
from typing import Any

from .runner import CaptureServer

# Where the capture app is mounted when the env server hosts it directly.
CAPTURE_MOUNT = "/capture"


def space_public_url() -> str:
    """The public URL of the Space this process is running in, or `""` when it is not on one.

    Returns:
        `str`: e.g. `https://owner-name.hf.space`, with no trailing slash.
    """
    host = os.environ.get("SPACE_HOST", "").strip()
    if host:
        return "https://" + host.rstrip("/").removeprefix("https://").removeprefix(
            "http://"
        )
    # SPACE_HOST is the direct answer, but SPACE_ID is the variable that is always set, so derive
    # the hostname the same way `auto.auto_env` does.
    space_id = os.environ.get("SPACE_ID", "").strip()
    if space_id and "/" in space_id:
        slug = space_id.replace("/", "-").replace("_", "-").replace(".", "-").lower()
        return f"https://{slug}.hf.space"
    return ""


class HarborService:
    """Long-lived state for a serving process: capture proxy, forwarding, datasets.

    Held at module scope by `serve_harbor` so that the environment instances OpenEnv builds
    per-request can reach it. They must not own it: `/metadata` and `/schema` construct a throwaway
    environment on every call, so anything expensive on `__init__` would be paid per docs hit.
    """

    _instance: "HarborService | None" = None

    def __init__(
        self,
        *,
        llm_url: str,
        model: str,
        datasets: list[str],
        capture_port: int = 8100,
        expose: str = "gradio",
    ) -> None:
        self.llm_url = llm_url
        self.model = model
        self.datasets = datasets
        self.capture = CaptureServer(llm_url=llm_url, model=model, port=capture_port)
        self._expose_kind = expose
        self.public_url = ""
        self.mounted = False
        self._forwarder: Any = None
        self._lock = threading.Lock()

    def start(self) -> str:
        """Make the capture proxy reachable from the sandbox and return its public URL.

        Two different situations, and conflating them is what makes the hosted case awkward:

        - **Hosted** (a Space). The platform already gives this process one public URL and exposes
          exactly one port. So the capture app is mounted onto the env server's own app under
          `CAPTURE_MOUNT` and reached at `<space>/capture`. No second port, no forwarding, nothing
          for the platform to object to.
        - **Local.** The sandbox runs off-cluster and cannot reach `127.0.0.1`, so the capture port
          is published by whichever forwarder was selected.
        """
        public = space_public_url()
        if public:
            # The env server's app serves it; `build_app` performs the mount.
            self.mounted = True
            self.public_url = f"{public}{CAPTURE_MOUNT}"
            return self.public_url

        from openenv.core.harness.capture.forwarding import make_forwarder

        self.capture.start()
        # A half-started service is worse than a failed one: the capture server owns a port and a
        # background thread, so leaving it up after the forwarder fails makes the next attempt fail
        # too, on a port conflict that has nothing to do with the real error.
        try:
            self._forwarder = make_forwarder(self._expose_kind)
            self.public_url = self._forwarder.start(self.capture.port)
        except BaseException:
            self._forwarder = None
            self.capture.stop()
            raise
        return self.public_url

    def stop(self) -> None:
        if self._forwarder is not None:
            self._forwarder.stop()
        self.capture.stop()

    @classmethod
    def current(cls) -> "HarborService | None":
        return cls._instance

    @classmethod
    def set_current(cls, service: "HarborService") -> None:
        cls._instance = service


def serve_harbor(
    *,
    llm_url: str,
    datasets: list[str],
    model: str | None = None,
    host: str = "0.0.0.0",
    port: int = 8000,
    capture_port: int = 8100,
    expose: str = "gradio",
    env_file: str | None = None,
) -> None:
    """Boot the capture proxy, then serve the env server with the UI mounted.

    Args:
        llm_url (`str`):
            OpenAI-spec inference endpoint.
        datasets (`list[str]`):
            Dataset specs to serve as splits.
        port (`int`, *optional*, defaults to `8000`):
            Env server port.
        capture_port (`int`, *optional*, defaults to `8100`):
            Capture proxy port. This is the one published to the sandbox. Ignored on a hosted
            platform, where the proxy is mounted on the env server's own app instead.
        expose (`str`, *optional*, defaults to `"gradio"`):
            How the sandbox reaches the capture proxy locally: `gradio`, `cloudflare` or `direct`.
    """
    import uvicorn

    from .startup import prepare

    caps = prepare(
        llm_url=llm_url,
        model=model,
        datasets=datasets,
        env_file=env_file,
        require_llm=True,
        quiet=False,
    )
    model = caps.llm.get("model") or model or ""

    service = HarborService(
        llm_url=llm_url,
        model=model,
        datasets=datasets,
        capture_port=capture_port,
        expose=expose,
    )
    public = service.start()
    HarborService.set_current(service)

    where = "mounted on this app" if service.mounted else f":{capture_port}"
    print(f"\ncapture   {where} -> {public}")
    print(
        f"server    http://{host}:{port}    (UI at /web, Task API at /{{env}}/splits)"
    )
    print("Ctrl-C to stop\n")

    # The UI is the whole point of this entry point, so turn it on rather than making the operator
    # discover an env var.
    os.environ.setdefault("ENABLE_WEB_INTERFACE", "true")

    app = build_app(datasets=datasets, llm_url=llm_url, model=model, llm=caps.llm)
    try:
        uvicorn.run(app, host=host, port=port, log_level="info")
    finally:
        service.stop()


def build_app(
    *,
    datasets: list[str],
    llm_url: str = "",
    model: str = "",
    llm: dict[str, Any] | None = None,
) -> Any:
    """The FastAPI app: Task API + MCP + the Gradio UI."""
    from openenv.core.env_server.http_server import create_app
    from openenv.core.env_server.mcp_environment import (
        CallToolAction,
        CallToolObservation,
    )

    from .environment import HarborEnvironment
    from .ui import harbor_gradio_builder

    HarborEnvironment.configure(
        datasets=datasets, llm_url=llm_url, model=model, llm=llm
    )

    def gradio_builder(
        _web_manager: Any = None,
        _action_fields: Any = None,
        _metadata: Any = None,
        _is_chat: Any = None,
        display_title: str = "",
        _quick_start: Any = None,
    ) -> Any:
        """OpenEnv calls this positionally with six web-interface arguments.

        Only the title is useful here: the Harbor UI drives rollouts through its own handlers rather
        than the generic action-field form, because a rollout is one long tool call, not a step.
        """
        return harbor_gradio_builder(datasets=datasets, title=display_title or "Harbor")

    app = create_app(
        HarborEnvironment,
        CallToolAction,
        CallToolObservation,
        env_name="harbor_env",
        max_concurrent_envs=int(os.getenv("MAX_CONCURRENT_ENVS", "4")),
        gradio_builder=gradio_builder,
        custom_tab_name="Harbor",
        custom_tab_primary=True,
        show_default_tab=False,
    )

    # When the platform gives us a single port, the capture proxy rides on this app instead of
    # being published separately. Mounting strips the prefix, so the proxy's own catch-all still
    # sees `/v1/chat/completions` and every dialect keeps working unchanged.
    service = HarborService.current()
    if service is not None and service.mounted:
        app.mount(CAPTURE_MOUNT, service.capture.app)

    return app