fix(runs): reach the platform's internal RFT router for hosted runs

packages/prime-runs/README.md

@@ -65,9 +65,14 @@ run.finish()
 
 - The platform enables external runs per team; a team outside the allowlist
   gets a `ForbiddenError` from `init()`.
-- `init(kind="train", id=os.environ["RUN_ID"])` attaches to an external run a
-  launcher already created: nothing is registered, the platform keeps the run's
-  failure marking, and a clean `finish()` still completes it.
+- `init(kind="train", id=os.environ["RUN_ID"])` attaches to a run a launcher
+  already created: nothing is registered, the platform keeps the run's failure
+  marking, and a clean `finish()` still completes it. A *hosted* run (one the
+  platform launched) is reachable only through the platform's internal RFT
+  root: pass the `PRIME_API_BASE` its launcher injects (`…/api/internal/rft`) as
+  `base_url=` and the SDK addresses that router, sending the run's token as
+  `x-api-key` too. Registering a run or setting its status is not available
+  there; the public API answers 400 for a hosted run's id.
 - Metrics are one row per `log_metrics` call, on their own uploader. The sample
   table gets one Parquet object per upload, every 10th step, keyed by the step
   an episode was dispatched at; a step logged in several calls gets several
@@ -109,7 +114,10 @@ table today's viewer reads. `log_*()` are queue puts, safe inside a coroutine;
 | `~/.prime/config.json` | Shared prime CLI config (`api_key`, `team_id`, `base_url`)             |
 
 Precedence is `init()` argument → environment variable → config file. A missing
-API key disables the run with a warning.
+API key disables the run with a warning. `base_url` is normally the platform
+origin; the internal RFT root a hosted training run is given
+(`…/api/internal`, with or without `/rft`) is accepted too and switches the
+client to that router (attached runs only).
 
 ## Not yet available
 

packages/prime-runs/src/prime_runs/_http.py

@@ -1,4 +1,12 @@
-"""Authenticated client for ``{base_url}/api/v1`` with a per-call retry policy.
+"""Authenticated client for the platform API with a per-call retry policy.
+
+Requests are addressed under ``{base_url}/api/v1``. When ``base_url`` is the
+platform's *internal* RFT root instead (``…/api/internal/rft``, what a hosted
+run's launcher injects as ``$PRIME_API_BASE``), they go under
+``{base_url}/api/internal`` and the token is repeated in ``x-api-key``, the
+header that router authenticates a run's own token on. A trailing ``/rft`` is
+accepted on either root; only attached runs are reachable through the internal
+one (it registers nothing and has no status endpoint).
 
 A failure is *ambiguous* when the server may already have processed the request
 (a 502/504, a read timeout). Replaying one is safe for a GET or PUT and not for
@@ -10,7 +18,8 @@
 import json
 import sys
 import time
-from typing import Any, Dict, Mapping, Optional, Union
+from typing import Any, Dict, Mapping, Optional, Tuple, Union
+from urllib.parse import urlsplit, urlunsplit
 
 import httpx
 from prime_traces.core.client import AMBIGUOUS_TRANSPORT_ERRORS, raise_for_response, retry_delay
@@ -25,6 +34,34 @@
 #: intermediary after forwarding.
 UNAMBIGUOUS_RETRY_STATUS = frozenset({429})
 DEFAULT_MAX_ATTEMPTS = 5
+PUBLIC_API_PATH = "/api/v1"
+#: The platform's internal RFT root, served to hosted runs; see the module docstring.
+INTERNAL_API_PATH = "/api/internal"
+
+
+#: API roots a base URL may end in, longest first; ``/rft`` is only recognised
+#: after an API path, so a host that happens to be named ``rft`` or ``api`` stays.
+_API_ROOTS = (
+    ("/api/internal/rft", INTERNAL_API_PATH),
+    ("/api/internal", INTERNAL_API_PATH),
+    ("/api/v1/rft", PUBLIC_API_PATH),
+    ("/api/v1", PUBLIC_API_PATH),
+)
+
+
+def split_api_root(base_url: str) -> Tuple[str, str]:
+    """``(origin, api_path)`` for a base URL that is a platform origin or one of
+    its API roots, with or without the ``/rft`` the RFT client historically got.
+    Only the URL's path is inspected, never its host."""
+    parts = urlsplit(base_url)
+    path = parts.path.rstrip("/")
+    api_path = PUBLIC_API_PATH
+    for suffix, candidate in _API_ROOTS:
+        if path.endswith(suffix):
+            path, api_path = path[: -len(suffix)], candidate
+            break
+    origin = urlunsplit((parts.scheme, parts.netloc, path, "", ""))
+    return origin.rstrip("/"), api_path
 
 
 def _user_agent() -> str:
@@ -51,11 +88,16 @@ def __init__(
         max_attempts: int = DEFAULT_MAX_ATTEMPTS,
         client: Optional[httpx.Client] = None,
     ) -> None:
-        self.base_url = base_url.rstrip("/").removesuffix("/api/v1")
-        self.api_prefix = f"{self.base_url}/api/v1"
+        self.base_url, api_path = split_api_root(base_url)
+        self.api_prefix = f"{self.base_url}{api_path}"
+        #: Addressing the internal RFT router rather than the public API.
+        self.internal = api_path == INTERNAL_API_PATH
         self.max_attempts = max(1, max_attempts)
         self._owns_client = client is None
         self._headers = {"Authorization": f"Bearer {api_key}", "User-Agent": _user_agent()}
+        if self.internal:
+            # The internal router reads a hosted run's own token from this header.
+            self._headers["x-api-key"] = api_key
         self._timeout = timeout
         self._client = client or self._new_client()
         if self._owns_client:  # only a pool we opened is ours to rebuild after a fork
@@ -86,9 +128,14 @@ def request(
         ``idempotent`` defaults to ``method != "POST"``."""
         url = f"{self.api_prefix}{path}"
         body = content if content is not None else (encode_json(json_body) if json_body else None)
+        # Sent per request as well as set on the pool, so an injected client is
+        # authenticated the same way as one we opened.
+        headers = dict(self._headers)
+        if body is not None:
+            headers["Content-Type"] = "application/json"
         request_kwargs: Dict[str, Any] = {
             "content": body,
-            "headers": {"Content-Type": "application/json"} if body is not None else None,
+            "headers": headers,
             "params": dict(params) if params else None,
         }
         if timeout is not None:  # None would disable httpx timeouts, not restore the default

packages/prime-runs/src/prime_runs/backend.py

@@ -289,8 +289,11 @@ def create(self, spec: RunSpec) -> RunHandle:
 
     def attach(self, run_id: str) -> RunHandle:
         """A handle for a run a launcher already created. No request: the first
-        write proves access. The run must be an external one; the monitoring
-        endpoints answer 400 for a hosted run's id."""
+        write proves access. An external run is reached through the public API,
+        whose monitoring endpoints answer 400 for a hosted run's id; a hosted run
+        is reached through the platform's internal RFT root — pass the
+        ``$PRIME_API_BASE`` its launcher injects as ``base_url`` and the client
+        authenticates with the run's token (see :mod:`prime_runs._http`)."""
         return RunHandle(id=run_id, url=self.url_for(run_id))
 
     def url_for(self, run_id: str) -> str:

packages/prime-runs/src/prime_runs/run.py

@@ -423,7 +423,10 @@ def init(
     model, ``environments`` the hub ids, ``training`` the display fields, and a
     team is required. ``id`` attaches to an external run a launcher already
     created (``$RUN_ID``): nothing is registered, the platform keeps the run's
-    failure marking, and a clean finish still completes it.
+    failure marking, and a clean finish still completes it. ``base_url`` is the
+    platform origin; a hosted run passes the internal RFT root its launcher
+    injects (``$PRIME_API_BASE``, ``…/api/internal/rft``), which serves attached
+    runs only.
     """
     settings = Config()
     api_key = api_key if api_key is not None else settings.api_key

packages/prime-runs/tests/test_http.py

@@ -206,3 +206,90 @@ def test_idempotent_methods_still_replay_ambiguous_failures(no_sleep):
     client = client_for(lambda request: responses.pop(0))
 
     assert client.put("/evaluations/x", json_body={"metrics": {}}) == {"ok": True}
+
+
+def _recording(seen):
+    def handler(request):
+        seen.append(request)
+        return httpx.Response(200, json={"data": {"status": "success"}})
+
+    return handler
+
+
+@pytest.mark.parametrize(
+    "base_url",
+    [
+        "http://testserver/api/internal/rft",
+        "http://testserver/api/internal/rft/",
+        "http://testserver/api/internal",
+    ],
+)
+def test_the_internal_rft_root_addresses_that_router_with_the_run_token(base_url):
+    """A hosted run's ``$PRIME_API_BASE``: paths go under ``/api/internal`` and the
+    token travels in ``x-api-key`` as well, which is what that router checks."""
+    seen = []
+    client = client_for(_recording(seen), base_url=base_url)
+
+    client.post("/rft/metrics", json_body={"run_id": "run-1", "metrics": {"loss": 1.0}})
+
+    assert client.internal
+    assert client.api_prefix == "http://testserver/api/internal"
+    assert str(seen[0].url) == "http://testserver/api/internal/rft/metrics"
+    assert seen[0].headers["x-api-key"] == "test-key"
+    assert seen[0].headers["Authorization"] == "Bearer test-key"
+
+
+@pytest.mark.parametrize(
+    "base_url",
+    [
+        "http://testserver",
+        "http://testserver/",
+        "http://testserver/api/v1",
+        "http://testserver/api/v1/rft",
+    ],
+)
+def test_the_public_root_stays_under_api_v1_without_the_run_token_header(base_url):
+    seen = []
+    client = client_for(_recording(seen), base_url=base_url)
+
+    client.post("/rft/metrics", json_body={"run_id": "run-1", "metrics": {"loss": 1.0}})
+
+    assert not client.internal
+    assert client.api_prefix == "http://testserver/api/v1"
+    assert str(seen[0].url) == "http://testserver/api/v1/rft/metrics"
+    assert "x-api-key" not in seen[0].headers
+    assert seen[0].headers["Authorization"] == "Bearer test-key"
+
+
+def test_an_injected_client_is_authenticated_per_request():
+    """The pool's default headers are also sent per call, so a caller-supplied
+    ``httpx.Client`` without them still authenticates."""
+    seen = []
+    client = client_for(_recording(seen))
+
+    client.get("/evaluations/x")
+
+    assert seen[0].headers["Authorization"] == "Bearer test-key"
+    assert seen[0].headers["User-Agent"].startswith("prime-runs/")
+
+
+@pytest.mark.parametrize(
+    "base_url,expected_prefix",
+    [
+        ("http://rft", "http://rft/api/v1"),
+        ("http://rft/", "http://rft/api/v1"),
+        ("https://rft.internal:8080", "https://rft.internal:8080/api/v1"),
+        ("http://api", "http://api/api/v1"),
+        ("http://api/v1", "http://api/v1/api/v1"),
+        (
+            "https://platform.test/proxy/api/internal/rft",
+            "https://platform.test/proxy/api/internal",
+        ),
+    ],
+)
+def test_only_the_path_is_normalised_never_a_host_named_like_a_suffix(base_url, expected_prefix):
+    """A service host called ``rft`` or ``api`` (Docker, Kubernetes) is not a path
+    suffix; the API path is recognised on the URL's path only."""
+    client = client_for(lambda request: httpx.Response(200, json={}), base_url=base_url)
+
+    assert client.api_prefix == expected_prefix

packages/prime-runs/tests/test_rft_backend.py

@@ -4,6 +4,7 @@
 import pytest
 from conftest import RecordingHandler
 
+from prime_runs._http import PlatformClient
 from prime_runs.backend import RftBackend
 from prime_runs.exceptions import ConfigurationError, ForbiddenError, RetryableAPIError
 from prime_runs.models import EnvironmentRef, RunSpec, RunStatus, TrainingSpec
@@ -212,3 +213,27 @@ def test_marking_an_already_closed_run_failed_is_not_an_error(make_platform_clie
     backend.finalize("run-abc", status=RunStatus.FAILED, error="boom")
 
     assert handler.paths() == ["PUT /api/v1/rft/external-runs/run-abc/status"]
+
+
+def test_a_hosted_run_is_finalized_through_the_internal_router():
+    """Attached with the launcher's ``$PRIME_API_BASE``: the same call reaches
+    ``/api/internal/rft/finalize`` carrying the run's token in ``x-api-key``."""
+    seen = []
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        seen.append(request)
+        return httpx.Response(200, json={"data": {"status": "success"}})
+
+    client = PlatformClient(
+        api_key="run-token",
+        base_url="http://testserver/api/internal/rft",
+        client=httpx.Client(transport=httpx.MockTransport(handler)),
+    )
+    backend = RftBackend(client, frontend_url="https://app.example")
+
+    handle = backend.attach("run-managed")
+    backend.finalize(handle.id, status=RunStatus.COMPLETED)
+
+    assert [str(r.url) for r in seen] == ["http://testserver/api/internal/rft/finalize"]
+    assert seen[0].headers["x-api-key"] == "run-token"
+    assert seen[0].headers["Authorization"] == "Bearer run-token"