feat(environment): route BaseEnvironmentServerWithTools tool calls to the worker that runs the episode

.github/workflows/unit-tests.yml

@@ -69,6 +69,7 @@ jobs:
             nemo_gym/client/**
             nemo_gym/agents/base_responses_api_agent.py
             nemo_gym/server_runtime/base_resources_server.py
+            nemo_gym/server_runtime/base_environment_server.py
             nemo_gym/server_runtime/mcp_auto_exposure.py
             tests/e2e/session_routing/**
 

nemo_gym/server_runtime/base_environment_server.py

@@ -17,7 +17,16 @@
 from anyio import CancelScope
 from fastapi import Body, FastAPI, Request
 from fastapi.responses import JSONResponse
-from pydantic import ConfigDict, Field, JsonValue, PositiveFloat, PositiveInt, field_validator, model_validator
+from pydantic import (
+    ConfigDict,
+    Field,
+    JsonValue,
+    PositiveFloat,
+    PositiveInt,
+    ValidationError,
+    field_validator,
+    model_validator,
+)
 from typing_extensions import Self
 
 from nemo_gym.client.http_client import format_exception_with_response_body
@@ -38,6 +47,7 @@
 from nemo_gym.server_runtime.episode_tokens import EPISODE_TOKEN_HEADER, EpisodeTokenClaims, EpisodeTokens
 from nemo_gym.server_runtime.episode_tools import EpisodeToolset, ToolArgumentsError, ToolCallRecord, ToolSpec
 from nemo_gym.server_runtime.request_handling.rollout_context import RolloutContextMiddleware
+from nemo_gym.server_runtime.request_handling.session_routing import SessionKind, SessionOwners
 from nemo_gym.server_runtime.tool_access import (
     DirectHTTPToolAccess,
     MCPStreamableHTTPConnection,
@@ -540,6 +550,15 @@ async def aggregate_metrics(self, body: AggregateMetricsRequest = Body()) -> Agg
         """Aggregate per-rollout scores into task-level metrics."""
 
 
+# An episode's tool calls name it only by the episode token header; no cookie, seed or close names it.
+EPISODE_TOOL_SESSIONS = SessionKind(
+    cookie_key=None,
+    id_field=None,
+    seed_path=None,
+    close_path=None,
+    cookie_session=False,
+)
+
 ToolTransport = Literal["direct_http", "mcp"]
 # The path ``nemo_gym.server_runtime.mcp_auto_exposure.mount_mcp_endpoint`` serves. Kept here so that a server without MCP never
 # imports the MCP SDK.
@@ -615,17 +634,23 @@ class BaseEnvironmentServerWithTools(BaseEnvironmentServer[EpisodeRequestT, Epis
     Both transports check the token the same way. Calls to one episode run one at a time, whatever their
     transport, and each is recorded in ``ServedTools.calls``.
 
+    With several workers, the worker that runs the episode owns it (see ``session_routing``), so a tool call that
+    reaches another worker, over either transport, is forwarded there, by the episode its token names.
+
     When the server shuts down, it finishes the cleanups of every episode still running, such as stopping its
     sandbox.
     """
 
     config: BaseEnvironmentServerWithToolsConfig
 
+    routes_sessions_to_owner = True
     _episode_tokens: Optional[EpisodeTokens] = None
     _open_episodes: Optional[dict[EpisodeId, _OpenEpisode]] = None
     # Keyed by id(): two attempts of one episode can run at once without checkpointing.
     _running_cleanups: Optional[dict[int, CleanupContext]] = None
     _reserved_tool_names: frozenset[str] = frozenset()
+    # With several workers, which worker runs each episode (see nemo_gym.server_runtime.request_handling.session_routing).
+    _session_owners: Optional[SessionOwners] = None
 
     def model_post_init(self, context: Any, /) -> None:
         super().model_post_init(context)
@@ -682,8 +707,27 @@ async def finish_running_cleanups(self) -> None:
         """
         await asyncio.gather(*(cleanup.aclose() for cleanup in list(self._running_cleanups.values())))
 
-    def multiple_workers_unsupported_reason(self) -> Optional[str]:
-        return "a tool call reaches whichever worker accepts it, and only the worker that runs the episode holds it"
+    def owned_session_routing(self) -> dict[str, Any]:
+        return {
+            "kind": EPISODE_TOOL_SESSIONS,
+            "session_header": EPISODE_TOKEN_HEADER,
+            "session_from_header": self._episode_session_id,
+            "session_holder": self,
+        }
+
+    def _episode_session_id(self, token: str) -> Optional[str]:
+        # Routing needs only the episode the token names; the worker that runs the episode checks its nonce.
+        claims = self._episode_tokens.read(token)
+        return claims.episode_id.capture_key if claims is not None else None
+
+    def holds_session(self, session_id: str) -> bool:
+        try:
+            return EpisodeId.from_capture_key(session_id) in self._open_episodes
+        except (ValueError, ValidationError):
+            return False
+
+    def use_session_owners(self, owners: SessionOwners) -> None:
+        self._session_owners = owners
 
     async def serve_tools(
         self, request: EpisodeRequestT, toolset: EpisodeToolset, cleanup: CleanupContext
@@ -709,13 +753,19 @@ async def serve_tools(
             if unnamable:
                 raise ValueError(f"MCP tool names may use only letters, digits, _ and -: {', '.join(unnamable)}")
 
+        session_owners = self._session_owners
+        if session_owners is not None and session_owners.acquire(episode_id.capture_key) != session_owners.routing_id:
+            raise ValueError(f"Episode {episode_id.capture_key} already serves its tools on another worker")
+
         token, claims = self._episode_tokens.mint(episode_id)
         episode = _OpenEpisode(claims=claims, toolset=toolset, specs=specs, calls=[])
         self._open_episodes[episode_id] = episode
 
         async def stop_serving() -> None:
             if self._open_episodes.get(episode_id) is episode:
                 del self._open_episodes[episode_id]
+                if session_owners is not None:
+                    session_owners.release(episode_id.capture_key)
 
         cleanup.register_cleanup("episode tools", stop_serving)
         base_url = self.server_client._resolve_base_url(self.config.name).rstrip("/")

nemo_gym/server_runtime/request_handling/session_routing.py

@@ -39,6 +39,7 @@
 | Any request to an agent server | The session cookie's ``agent_session_id`` |
 | An agent session seed or close | The body's ``agent_session_id`` |
 | An agent rollout stop without a cookie | The path's ``/rollouts/<rollout id>``, which the session's worker also owns |
+| A tool call to a ``BaseEnvironmentServerWithTools`` | The episode named in its episode token header |
 
 The body is read only on the seed and close routes, where it names the session before any cookie exists.
 If the cookie and the body name different sessions, the reply is 409.
@@ -64,7 +65,9 @@
 - a resources session's ownership lasts while the session is open,
   from a successful seed until its close or ``/verify``,
   and expires ``session_close_retry_window_seconds`` after the last request for it otherwise.
-  A ``/verify`` that says a close follows, as an environment server's does, leaves the session open for that close.
+  A ``/verify`` that says a close follows, as an environment server's does, leaves the session open for that close;
+- an episode of a ``BaseEnvironmentServerWithTools`` is owned by the worker that runs it, from when it starts serving
+  the episode's tools until it stops (see ``SessionHolder``).
 
 An agent worker also owns ``/rollouts/<rollout id>`` from the start of a session's seed until the session closes,
 so a rollout stop, which names only the rollout, reaches the worker that runs it.
@@ -163,15 +166,19 @@ def _paths(paths: tuple[str, ...]) -> re.Pattern[str]:
 
 @dataclass(frozen=True)
 class SessionKind:
-    """Where one kind of server's requests name their session."""
+    """Where one kind of server's requests name their session.
+
+    A kind whose requests name their session only in a header, such as the episode tool calls of a
+    ``BaseEnvironmentServerWithTools``, has no cookie key, body field or seed and close routes.
+    """
 
     #: The key of the session ID in the session cookie.
-    cookie_key: str
+    cookie_key: Optional[str]
     #: The body field that names the session on the seed and close routes.
-    id_field: str
+    id_field: Optional[str]
     #: The seed and close routes, whose body names the session.
-    seed_path: str
-    close_path: str
+    seed_path: Optional[str]
+    close_path: Optional[str]
     #: Whether the session ID is the cookie session's own ID, as on a resources server:
     #: an episode seed's ID becomes it, and a request without one gets a new ID,
     #: owned by this worker if the request starts a session or its handler uses the session.
@@ -182,8 +189,9 @@ class SessionKind:
     path_key: Optional[re.Pattern[str]] = None
 
     @property
-    def body_paths(self) -> re.Pattern[str]:
-        return _paths((self.seed_path, self.close_path))
+    def body_paths(self) -> Optional[re.Pattern[str]]:
+        paths = tuple(path for path in (self.seed_path, self.close_path) if path is not None)
+        return _paths(paths) if paths else None
 
 
 def create_session_socket_dir() -> str:
@@ -545,6 +553,42 @@ def __len__(self) -> int:
         return len(self._until.keys() | self._retained)
 
 
+class SessionHolder(Protocol):
+    """A server that acquires each of its sessions itself, and owns it for as long as it holds the session.
+
+    A ``BaseEnvironmentServerWithTools`` is one: it acquires an episode when it starts serving the episode's tools, and
+    releases it when it stops.
+    """
+
+    def holds_session(self, session_id: str) -> bool:
+        """Whether this worker still holds the session."""
+        ...
+
+    def use_session_owners(self, owners: SessionOwners) -> None:
+        """Keep the session owners, to acquire a session when it starts and release it when it ends."""
+        ...
+
+
+class _HeldSessionLifetimes:
+    """A holder's own state decides what it owns: it acquires and releases each session itself."""
+
+    def __init__(self, holder: SessionHolder) -> None:
+        self.holder = holder
+
+    def holds(self, session_id: str) -> bool:
+        return self.holder.holds_session(session_id)
+
+    def served(self, session_id: str, scope: Scope, status: Optional[int]) -> None:
+        pass
+
+    def starts(self, path: str) -> bool:
+        # The holder acquires a session when it starts it, not on a request.
+        return False
+
+    def sweep(self) -> list[str]:
+        return []
+
+
 class AgentSessionHolder(Protocol):
     """An agent server, whose sessions are owned by ``agent_session_id`` across workers."""
 
@@ -610,6 +654,8 @@ def __init__(
         mcp_token_header: Optional[str] = None,
         mcp_token_serializer: Optional[itsdangerous.URLSafeSerializer] = None,
         mcp_path: str = "/mcp",
+        session_header: Optional[str] = None,
+        session_from_header: Optional[Callable[[str], Optional[str]]] = None,
     ) -> None:
         self.app = app
         self.kind = kind
@@ -626,13 +672,20 @@ def __init__(
         self.mcp_token_header = mcp_token_header
         self.mcp_token_serializer = mcp_token_serializer
         self.mcp_path = mcp_path
+        # A header that names the session, and the server's reader of its value; ahead of the cookie and MCP token.
+        self.session_header = session_header
+        self.session_from_header = session_from_header
 
     async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
         if scope["type"] != "http":
             await self.app(scope, receive, send)
             return
         owners = self.owners
-        names_by_body = scope["method"] == "POST" and self.body_paths.fullmatch(scope["path"]) is not None
+        names_by_body = (
+            scope["method"] == "POST"
+            and self.body_paths is not None
+            and self.body_paths.fullmatch(scope["path"]) is not None
+        )
         if owners is None and not names_by_body:
             await self.app(scope, receive, send)
             return
@@ -722,7 +775,8 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
         elif forwarded:
             # The owner released the ID while the request was forwarded here, and another worker took ownership since.
             await _conflict(
-                f"This {self.kind.id_field} moved to another server worker while the request was forwarded. Retry it.",
+                f"This {self.kind.id_field or 'session'} moved to another server worker while the request was "
+                "forwarded. Retry it.",
                 scope,
                 receive,
                 send,
@@ -749,7 +803,13 @@ async def send_and_record(message: Message) -> None:
                     self.lifetimes.served(session_id, scope, status[0] if status else None)
 
     def _session_id_from_headers(self, connection: HTTPConnection) -> Optional[str]:
-        cookie = connection.cookies.get(self.session_cookie)
+        if self.session_header is not None:
+            value = connection.headers.get(self.session_header)
+            if value is not None:
+                # A value the reader rejects names no session; the handler refuses the request.
+                return self.session_from_header(value)
+        # A kind without a cookie key never names its session in the cookie.
+        cookie = connection.cookies.get(self.session_cookie) if self.kind.cookie_key is not None else None
         token = connection.headers.get(self.mcp_token_header) if self.mcp_token_serializer is not None else None
         path = connection.scope["path"]
         is_mcp = path == self.mcp_path or path.startswith(self.mcp_path + "/")
@@ -964,7 +1024,10 @@ def install_session_routing(
     secret_key: str,
     mcp_token_header: Optional[str] = None,
     mcp_token_serializer: Optional[itsdangerous.URLSafeSerializer] = None,
+    session_header: Optional[str] = None,
+    session_from_header: Optional[Callable[[str], Optional[str]]] = None,
     agent_sessions: Optional[AgentSessionHolder] = None,
+    session_holder: Optional[SessionHolder] = None,
     close_retry_window_seconds: Optional[float] = None,
     session_starts: tuple[str, ...] = (),
     session_ends: tuple[str, ...] = (),
@@ -978,16 +1041,22 @@ def install_session_routing(
     With one, it records session owners in it and returns this worker's ID,
     also stored in ``app.state.nemo_gym_routing_id``.
     Pass the agent for an agent server, whose own records decide how long it owns its sessions,
+    the holder for a server that acquires and releases each session itself (see ``SessionHolder``),
     or the close retry window and the session start and end routes for a resources server
     (see ``ResourcesSessionLifetimes``).
+    A server whose requests name their session in a header passes the header and its reader of the header's value.
     """
     route_by = dict(
         kind=kind,
         session_cookie=session_cookie,
         secret_key=secret_key,
         mcp_token_header=mcp_token_header,
         mcp_token_serializer=mcp_token_serializer,
+        session_header=session_header,
+        session_from_header=session_from_header,
     )
+    if session_header is not None and session_from_header is None:
+        raise ValueError("install_session_routing needs session_from_header with a session_header")
     if socket_dir is None:
         # An agent's own seed handler already refuses a cookie of another session, so it needs nothing here.
         if kind.cookie_session:
@@ -1000,6 +1069,9 @@ def install_session_routing(
     if agent_sessions is not None:
         lifetimes: SessionLifetimes = _AgentSessionLifetimes(agent_sessions)
         window = None
+    elif session_holder is not None:
+        lifetimes = _HeldSessionLifetimes(session_holder)
+        window = None
     elif close_retry_window_seconds is not None:
         resources_lifetimes = ResourcesSessionLifetimes(
             close_retry_window_seconds,
@@ -1012,13 +1084,16 @@ def install_session_routing(
         window = close_retry_window_seconds
     else:
         raise ValueError(
-            "install_session_routing needs agent_sessions or close_retry_window_seconds with a socket_dir"
+            "install_session_routing needs agent_sessions, session_holder or close_retry_window_seconds "
+            "with a socket_dir"
         )
     owners = SessionOwners(socket_dir, routing_id, lifetimes.holds)
     app.state.nemo_gym_session_owners = owners
     if agent_sessions is not None:
         agent_sessions.use_agent_session_owners(owners)
         window = getattr(getattr(agent_sessions, "config", None), "session_close_retry_window_seconds", None)
+    if session_holder is not None:
+        session_holder.use_session_owners(owners)
     sweep_seconds = min(_MAX_SWEEP_SECONDS, max(_MIN_SWEEP_SECONDS, (window or _MAX_SWEEP_SECONDS) / 4))
     app.add_middleware(
         SessionRoutingMiddleware,