nemo_gym.session_routing

Route each session’s requests to the uvicorn worker that holds its state.

uvicorn’s workers share one listening socket, so a request reaches whichever worker accepts the connection, while a session’s state lives in the memory of the worker that created it. With more than one worker, each worker also serves the same app on a private Unix socket, in a directory the main process creates for the server, and requests are routed by one rule:

Find the request’s session ID.
Look up the session’s claim, which names the worker that created the session.
Forward the request to that worker’s private socket, unchanged, and stream its reply back unchanged. If this worker holds the claim, or the session has no claim, the request is handled here.

A request that arrives over a private socket is always handled where it arrives, so it is forwarded at most once. If the claimant’s socket is gone, because its worker exited, the request gets a 410 instead of silently running against empty state.

Where the session ID comes from:

Request	Session ID
Any request to a resources server	The session cookie’s `session_id`
A resources tool call over MCP, or a request with an MCP session token and no cookie	The token’s `sid`
A typed resources `/seed_session` or `/close_session`	The body’s `resources_session_id`
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`

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. A typed resources session uses its resources_session_id as its cookie’s session_id too, so the seed, its close and every request in between name the session by the same ID. A resources request with no session ID gets a new ID, which this router mints. A session start, such as a seed, claims it before its handler runs, since its reply’s cookie names the session. Any other request claims it only if the handler used the session, before the reply leaves the worker. A request that keeps nothing, such as a health check or /aggregate_metrics, leaves no claim and gets no cookie, so no cookie names a session nobody claims. A request whose cookie names a session whose claim ended is claimed the same way by the worker that serves it.

A claim is a file in the socket directory, named by a hash of the session ID and holding the claimant’s worker ID. The worker that creates the session creates the claim, atomically, so two workers can never claim one ID. Seeds and closes create claims, so a retried seed, a cleanup close, or a delayed seed after a close all reach the worker that saw the ID first. A claim ends once its session’s last record on the claimant is gone:

an agent session’s claim ends when the agent drops the session’s last record, which is its close receipt or failed seed, session_close_retry_window_seconds after it was written;
a resources session’s claim lasts while the session is open, from a successful seed until its close or /verify, and ends 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.

The claims of a worker that exited stay until the server exits, so its sessions keep getting a 410. The whole directory goes with the server’s runtime directory when the server exits. With one worker there is no private socket and no claim.

Module Contents

Classes

Name	Description
`AgentSessionHolder`	An agent server, whose sessions are claimed by `agent_session_id` across workers.
`ResourcesSessionLifetimes`	How long a resources server’s worker keeps the claim on a session, judged from the requests it served.
`SessionClaims`	Which worker created each session, shared by the workers of one server.
`SessionKind`	Where one kind of server’s requests name their session.
`SessionLifetimes`	How long a worker keeps the claims on the sessions it created.
`SessionRoutingMiddleware`	Find a request’s session ID, and forward the request to the worker that claims the session.
`_AgentSessionLifetimes`	An agent’s own records decide its claims: it releases a claim itself when it drops a session’s last record.
`_PrivateSocketProtocol`	uvicorn’s httptools protocol without access logs, which it would otherwise write whenever the logger has handlers.

Functions

Name	Description
`_conflict`	-
`_mark_forwarded`	-
`_paths`	Match any of paths, bare or under a rollout prefix.
`_read_body`	Read the request body, or None if the client disconnected.
`_replay_body`	A receive that gives the app the body already read, then passes later messages, such as a disconnect, on.
`_session_id`	The session ID in `field` of a seed or close body, or None for a body that has no non-empty string one.
`_sweep_claims`	Release the claims whose session has ended, so a claim ends even when no more requests arrive.
`_valid_routing_id`	-
`cookie_session`	Return the session stored in a session cookie, or None if Starlette’s SessionMiddleware would reject it.
`create_session_socket_dir`	The directory for one server’s private worker sockets and claims: this process’s runtime directory.
`install_session_routing`	Route this worker’s requests by session ID, and serve forwarded requests on a private socket.
`mcp_token_session_id`	Return the session ID in an MCP session token, or None for a token with a bad signature.
`serve_private_socket`	Serve app on a Unix socket for requests other workers forward here, on the running event loop.
`worker_socket_path`	-

Data

LOG

SESSION_CLAIMS_DIR

SESSION_CLAIM_SCOPE_KEY

SESSION_ID_SCOPE_KEY

SESSION_SOCKET_DIR_ENV

_FORWARDED_SCOPE_KEY

_FORWARD_KEEPALIVE_SECONDS

_HOP_BY_HOP_HEADERS

_MAX_SWEEP_SECONDS

_MIN_SWEEP_SECONDS

_PRIVATE_SERVER_KEEPALIVE_SECONDS

_ROLLOUT_PREFIX

_ROUTING_ID_PATTERN

_SESSION_MAX_AGE_SECONDS

API

class nemo_gym.session_routing.AgentSessionHolder()

Protocol

An agent server, whose sessions are claimed by agent_session_id across workers.

nemo_gym.session_routing.AgentSessionHolder.expire_agent_session_records() -> None

Drop the close receipts and failed seeds whose window has passed, releasing their claims.

nemo_gym.session_routing.AgentSessionHolder.holds_agent_session(
    agent_session_id: str
) -> bool

Whether this worker still keeps any record of the session, open, failed or closed.

nemo_gym.session_routing.AgentSessionHolder.use_agent_session_claims(
    claims: nemo_gym.session_routing.SessionClaims
) -> None

Keep the claims, to release one when the worker drops the session’s last record.

class nemo_gym.session_routing.ResourcesSessionLifetimes(
    window_seconds: float,
    starts: tuple[str, ...],
    ends: tuple[str, ...],
    close_follows_header: typing.Optional[str] = None
)

How long a resources server’s worker keeps the claim on a session, judged from the requests it served.

A resources server keeps its per-session state in its own subclass, so the router judges by the requests:

A successful seed, or another declared session start such as a Gymnasium /reset, opens the session: its claim lasts until the session ends.
A successful /close_session or /verify, or SimpleResourcesServer.checkpoint_session_ended, ends it, unless the request says a close follows, as an environment server’s /verify does.
A session that is not open keeps its claim window_seconds after the last request for it. The caller’s retried or cleanup close reaches the claimant within that window. A delayed seed after a close reaches it too, where the server refuses it while it remembers the close.
A failed seed or close of an open session leaves it open.

So a claim outlives the session’s last request by at most the window, except for a session never closed, which keeps its claim as long as the server keeps the session.

_close_follows

_ends

= _paths(ends)

_expiries

list[tuple[float, str]] = []

_queued

set[str] = set()

_starts

= _paths(starts)

_until

dict[str, float] = {}

nemo_gym.session_routing.ResourcesSessionLifetimes.__len__() -> int

nemo_gym.session_routing.ResourcesSessionLifetimes._close_follows_in(
    scope: starlette.types.Scope
) -> bool

nemo_gym.session_routing.ResourcesSessionLifetimes._expire_after_window(
    session_id: str
) -> None

nemo_gym.session_routing.ResourcesSessionLifetimes.ended(
    session_id: str
) -> None

End a session this worker claims, from a route that ends a session under another name.

nemo_gym.session_routing.ResourcesSessionLifetimes.holds(
    session_id: str
) -> bool

nemo_gym.session_routing.ResourcesSessionLifetimes.served(
    session_id: str,
    scope: starlette.types.Scope,
    status: typing.Optional[int]
) -> None

nemo_gym.session_routing.ResourcesSessionLifetimes.starts(
    path: str
) -> bool

nemo_gym.session_routing.ResourcesSessionLifetimes.sweep() -> list[str]

class nemo_gym.session_routing.SessionClaims(
    socket_dir: str,
    routing_id: str,
    holds: collections.abc.Callable[[str], bool]
)

Which worker created each session, shared by the workers of one server.

A claim is a file in the claims directory, named by a hash of the session ID and holding the claimant’s worker ID. It is created with a hard link from a complete temporary file, which fails if the claim exists, so two workers cannot both claim one ID and a reader never sees a partly written claim.

A claim lives as long as holds says its claimant keeps anything of the session. It is not released while a request for the session is being served here, since that request may be about to create a record. Only the claimant removes its claim, so a removal cannot delete another worker’s claim made since.

_mine

set[str] = set()

_serving

dict[str, int] = {}

directory

= os.path.join(socket_dir, SESSION_CLAIMS_DIR)

nemo_gym.session_routing.SessionClaims.__len__() -> int

nemo_gym.session_routing.SessionClaims._path(
    session_id: str
) -> str

nemo_gym.session_routing.SessionClaims.claim(
    session_id: str
) -> typing.Optional[str]

Claim the ID for this worker unless another worker has; return the claimant.

Returns None for a claim file that names no valid worker, which nothing here writes.

nemo_gym.session_routing.SessionClaims.claimant(
    session_id: str
) -> typing.Optional[str]

The worker that claims the ID, or None if nobody does.

nemo_gym.session_routing.SessionClaims.mine(
    session_id: str
) -> bool

nemo_gym.session_routing.SessionClaims.release(
    session_id: str
) -> None

Remove this worker’s claim on the ID, unless a request for it is in progress or the session remains.

nemo_gym.session_routing.SessionClaims.serving(
    session_id: str
) -> collections.abc.Iterator[None]

Hold the claim while a request for the ID is served here, and release it after if nothing is left.

class nemo_gym.session_routing.SessionKind(
    cookie_key: str,
    id_field: str,
    seed_path: str,
    close_path: str,
    cookie_session: bool
)

Dataclass

Where one kind of server’s requests name their session.

body_paths

Pattern[str]

close_path

str

cookie_key

str

cookie_session

bool

id_field

str

seed_path

str

class nemo_gym.session_routing.SessionLifetimes()

Protocol

How long a worker keeps the claims on the sessions it created.

nemo_gym.session_routing.SessionLifetimes.holds(
    session_id: str
) -> bool

Whether this worker still keeps any record of the session.

nemo_gym.session_routing.SessionLifetimes.served(
    session_id: str,
    scope: starlette.types.Scope,
    status: typing.Optional[int]
) -> None

Record a request for a session this worker claims, once it is served; status is None if it failed.

nemo_gym.session_routing.SessionLifetimes.starts(
    path: str
) -> bool

Whether a request on path starts a session, so it claims its session whatever its handler does.

nemo_gym.session_routing.SessionLifetimes.sweep() -> list[str]

Drop the records whose time has passed, and return the sessions whose claims can be released.

class nemo_gym.session_routing.SessionRoutingMiddleware(
    app: starlette.types.ASGIApp,
    kind: nemo_gym.session_routing.SessionKind,
    session_cookie: str,
    secret_key: str,
    routing_id: typing.Optional[str] = None,
    socket_dir: typing.Optional[str] = None,
    clients: typing.Optional[dict[str, aiohttp.ClientSession]] = None,
    claims: typing.Optional[nemo_gym.session_routing.SessionClaims] = None,
    lifetimes: typing.Optional[nemo_gym.session_routing.SessionLifetimes] = None,
    mcp_token_header: typing.Optional[str] = None,
    mcp_token_serializer: typing.Optional[itsdangerous.URLSafeSerializer] = None,
    mcp_path: str = '/mcp'
)

Find a request’s session ID, and forward the request to the worker that claims the session.

It sits outside SessionMiddleware and reads the cookie by decoding it with the same secret. A forwarded request therefore never reaches this worker’s SessionMiddleware, so the reply carries only the claimant’s Set-Cookie, built from the session as the claimant left it. Sitting inside SessionMiddleware instead would add this worker’s stale copy of the session as a second Set-Cookie.

With one worker there are no claims, and it only reads the session ID of seeds and closes: it refuses a request whose cookie and body name different sessions, and gives a typed resources session its ID in the cookie.

body_paths

= kind.body_paths

clients

= clients if clients is not None else {}

signer

= itsdangerous.TimestampSigner(secret_key)

nemo_gym.session_routing.SessionRoutingMiddleware.__call__(
    scope: starlette.types.Scope,
    receive: starlette.types.Receive,
    send: starlette.types.Send
) -> None

async

nemo_gym.session_routing.SessionRoutingMiddleware._client(
    owner: str
) -> aiohttp.ClientSession

nemo_gym.session_routing.SessionRoutingMiddleware._discard_client(
    owner: str
) -> None

async

nemo_gym.session_routing.SessionRoutingMiddleware._forward(
    owner: str,
    scope: starlette.types.Scope,
    receive: starlette.types.Receive,
    send: starlette.types.Send,
    body: typing.Optional[bytes] = None
) -> None

async

nemo_gym.session_routing.SessionRoutingMiddleware._serve_claimed(
    claims: nemo_gym.session_routing.SessionClaims,
    session_id: str,
    scope: starlette.types.Scope,
    receive: starlette.types.Receive,
    send: starlette.types.Send
) -> None

async

Serve a request for a session this worker claims, or may claim, and record it for the claim’s lifetime.

nemo_gym.session_routing.SessionRoutingMiddleware._session_id_from_headers(
    connection: starlette.requests.HTTPConnection
) -> typing.Optional[str]

class nemo_gym.session_routing._AgentSessionLifetimes(
    agent: nemo_gym.session_routing.AgentSessionHolder
)

An agent’s own records decide its claims: it releases a claim itself when it drops a session’s last record.

nemo_gym.session_routing._AgentSessionLifetimes.holds(
    session_id: str
) -> bool

nemo_gym.session_routing._AgentSessionLifetimes.served(
    session_id: str,
    scope: starlette.types.Scope,
    status: typing.Optional[int]
) -> None

nemo_gym.session_routing._AgentSessionLifetimes.starts(
    path: str
) -> bool

nemo_gym.session_routing._AgentSessionLifetimes.sweep() -> list[str]

class nemo_gym.session_routing._PrivateSocketProtocol(
    args: typing.Any = (),
    kwargs: typing.Any = {}
)

Protocol

Bases: HttpToolsProtocol

uvicorn’s httptools protocol without access logs, which it would otherwise write whenever the logger has handlers.

nemo_gym.session_routing._conflict(
    detail: str,
    scope: starlette.types.Scope,
    receive: starlette.types.Receive,
    send: starlette.types.Send
) -> None

async

nemo_gym.session_routing._mark_forwarded(
    app: starlette.types.ASGIApp
) -> starlette.types.ASGIApp

nemo_gym.session_routing._paths(
    paths: tuple[str, ...]
) -> re.Pattern[str]

Match any of paths, bare or under a rollout prefix.

nemo_gym.session_routing._read_body(
    receive: starlette.types.Receive
) -> typing.Optional[bytes]

async

Read the request body, or None if the client disconnected.

nemo_gym.session_routing._replay_body(
    body: bytes,
    receive: starlette.types.Receive
) -> starlette.types.Receive

A receive that gives the app the body already read, then passes later messages, such as a disconnect, on.

nemo_gym.session_routing._session_id(
    body: bytes,
    field: str
) -> typing.Optional[str]

The session ID in field of a seed or close body, or None for a body that has no non-empty string one.

nemo_gym.session_routing._sweep_claims(
    claims: nemo_gym.session_routing.SessionClaims,
    lifetimes: nemo_gym.session_routing.SessionLifetimes,
    interval_seconds: float
) -> None

async

Release the claims whose session has ended, so a claim ends even when no more requests arrive.

nemo_gym.session_routing._valid_routing_id(
    owner: typing.Any
) -> typing.Optional[str]

nemo_gym.session_routing.cookie_session(
    cookie: str,
    signer: itsdangerous.TimestampSigner
) -> typing.Optional[dict[str, typing.Any]]

Return the session stored in a session cookie, or None if Starlette’s SessionMiddleware would reject it.

Decoded exactly as SessionMiddleware decodes it. A cookie it rejects (badly signed or expired) starts a fresh session.

nemo_gym.session_routing.create_session_socket_dir() -> str

The directory for one server’s private worker sockets and claims: this process’s runtime directory.

nemo_gym.session_routing.install_session_routing(
    app: fastapi.FastAPI,
    kind: nemo_gym.session_routing.SessionKind,
    socket_dir: typing.Optional[str],
    session_cookie: str,
    secret_key: str,
    mcp_token_header: typing.Optional[str] = None,
    mcp_token_serializer: typing.Optional[itsdangerous.URLSafeSerializer] = None,
    agent_sessions: typing.Optional[nemo_gym.session_routing.AgentSessionHolder] = None,
    close_window_seconds: typing.Optional[float] = None,
    session_starts: tuple[str, ...] = (),
    session_ends: tuple[str, ...] = (),
    close_follows_header: typing.Optional[str] = None
) -> typing.Optional[str]

Route this worker’s requests by session ID, and serve forwarded requests on a private socket.

Call once the app’s middleware is otherwise in place: the router must sit outside SessionMiddleware. Without a socket_dir (one worker), it only reads the session ID of seeds and closes, and claims nothing. With one, it claims sessions 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 its claims’ lifetimes, or the close window and the session start and end routes for a resources server (see ResourcesSessionLifetimes).

nemo_gym.session_routing.mcp_token_session_id(
    token: str,
    serializer: itsdangerous.URLSafeSerializer
) -> typing.Optional[str]

Return the session ID in an MCP session token, or None for a token with a bad signature.

A rejected token is handled locally, where the MCP endpoint’s own token check refuses it.

nemo_gym.session_routing.serve_private_socket(
    app: starlette.types.ASGIApp,
    path: str
) -> collections.abc.AsyncIterator[None]

async

Serve app on a Unix socket for requests other workers forward here, on the running event loop.

nemo_gym.session_routing.worker_socket_path(
    socket_dir: str,
    routing_id: str
) -> str

nemo_gym.session_routing.LOG = logging.getLogger(__name__)

nemo_gym.session_routing.SESSION_CLAIMS_DIR = 'session-claims'

nemo_gym.session_routing.SESSION_CLAIM_SCOPE_KEY = 'nemo_gym.claim_session'

nemo_gym.session_routing.SESSION_ID_SCOPE_KEY = 'nemo_gym.session_id'

nemo_gym.session_routing.SESSION_SOCKET_DIR_ENV = 'NEMO_GYM_SESSION_SOCKET_DIR'

nemo_gym.session_routing._FORWARDED_SCOPE_KEY = 'nemo_gym.session_forwarded'

nemo_gym.session_routing._FORWARD_KEEPALIVE_SECONDS = 15.0

nemo_gym.session_routing._HOP_BY_HOP_HEADERS = frozenset({'connection', 'keep-alive', 'transfer-encoding', 'content-length', 'u...

nemo_gym.session_routing._MAX_SWEEP_SECONDS = 30.0

nemo_gym.session_routing._MIN_SWEEP_SECONDS = 0.05

nemo_gym.session_routing._PRIVATE_SERVER_KEEPALIVE_SECONDS = 30

nemo_gym.session_routing._ROLLOUT_PREFIX = f'(?:/{ROLLOUT_PATH_PREFIX}/[^/]+(?:/[^/]+)?)?'

nemo_gym.session_routing._ROUTING_ID_PATTERN = re.compile('[0-9a-f]{32}')

nemo_gym.session_routing._SESSION_MAX_AGE_SECONDS = 14 * 24 * 60 * 60