fix(nemo-gym): survive a real multi-server NeMo-RL GRPO run

docs/engineering/status.md

@@ -55,7 +55,7 @@ drivers not yet live-run and the post-hardening CRIU path are called out explici
 | CU provider adapters (Anthropic, OpenAI, **Kimi-VL**) → canonical ACI | done | fixture-tested, no live API; #75/#76 + Kimi-VL |
 | **String-form XML tool calls as first-class action input** (`shinken.dialect.parse_xml_actions` + `parse_actions(text, format="auto"\|"xml"\|"dialect")`: JSON-in-XML `<tool_call>`, `<invoke>`/`<parameter>` blocks, `<function=…>` element XML, attribute/element XML; tolerant recovery, typed `DialectError` on unknown/unsupported — never a silent drop; adapters expose `.from_text`) | done | corpus-driven `tests/test_dialect_xml.py` (39 real-shaped samples) · [aci-spec §3.1](../design/aci-spec.md) |
 | **RL rollout-server runtime plugin** ([ProRL-Agent-Server](https://github.com/NVIDIA-NeMo/ProRL-Agent-Server) `BaseRuntime` contract — `integrations/prorl_agent_server.py`, INIT = resume-from-golden) | done | fixture-tested, no hard dep; live Docker roundtrip env-gated (`SHINKEN_DOCKER_LIVE=1`) · [agent-runtime.md](../design/agent-runtime.md) |
-| **NeMo Gym computer-use resources engine** (`integrations/nemo_gym.py`): fork-from-golden seed, generation-fenced tool/verify/end/reseed, same-session linearization, active abandoned-rollout reaping, leased TTL/LRU-bounded golden ownership, terminal shutdown barrier; backend/OS-neutral lifecycle with Linux/X11 demo evidence | done (single-process instance) | deterministic fixture tests cover concurrent seed/tool/verify/close/reap and deletion retry; real two-demo Docker loop + recorded NeMo Gym rollout collection · [example](../../examples/nemo_gym/README.md). Multi-worker startup fails closed; horizontal session routing/control plane is not built |
+| **NeMo Gym computer-use resources engine** (`integrations/nemo_gym.py`): fork-from-golden seed, generation-fenced tool/verify/end/reseed (with a `session_id`-keyed generation fallback + per-response session re-assert so a multi-server agent's chained cookies keep threading our namespaced session), same-session linearization, active abandoned-rollout reaping, leased TTL/LRU-bounded golden ownership, terminal shutdown barrier (Starlette lifespan-compat); backend/OS-neutral lifecycle with Linux/X11 demo evidence | done (single-process instance); **real NeMo-RL GRPO step closed end-to-end 2026-07-15** | deterministic fixture tests cover concurrent seed/tool/verify/close/reap, deletion retry, the lifespan shim, and the session/cookie recovery; real two-demo Docker loop + recorded NeMo Gym rollout collection + a **live NeMo-RL GRPO training step** (Qwen3-4B, 1 GPU: validation + `num_prompts×num_generations` training group forked from the golden checkpoint → rewards→advantages→logprobs) · [example](../../examples/nemo_gym/README.md). Multi-worker startup fails closed; horizontal session routing/control plane is not built |
 | **Agent-runtime narrow waist** (`shinken.runtime`: Session/rollout/Trajectory, zero scorer/reward) + **Workload registry** | done | #220/#221/#227 · [agent-runtime.md](../design/agent-runtime.md) |
 | **Provider registry** + `DockerLocalProvider` + out-of-tree plugin loaders | done | #219/#226 |
 | **Pluggable `shinkend` injector** (`shinken.inject`: `docker`/`ssh`/`osworld-exec`) | done | #230/#233 — chunked upload, shell-wrapped start, configurable remote path, surfaced errors |

examples/nemo_gym/README.md

@@ -116,6 +116,21 @@ exactly the shape Shinken's fleet numbers are built for: `num_generations_per_pr
 maps to N forks of one golden checkpoint at ~0.5 s each, where re-provisioning
 environments per generation is the cost the trainer otherwise eats.
 
+**Verified end-to-end (2026-07-15).** A real NeMo RL GRPO run (not a bespoke trainer)
+drove rollouts through this resources server on a single GPU node — Qwen3-4B policy in
+vLLM, the `run_grpo_nemo_gym.py` entrypoint over the container's Gym submodule — and closed
+one full training step: `Collecting rollouts` completed the validation set and the
+`num_prompts_per_step × num_generations_per_prompt` training group **as forks of the golden
+desktop checkpoint**, then `Processing rewards → Computing advantages → Computing logprobs`.
+Bringing that up surfaced three integration fixes in `shinken.integrations.nemo_gym` (all
+covered by `tests/test_nemo_gym_integration.py`): a Starlette lifespan-compat shim for the
+engine-shutdown binding (newer Starlette dropped `add_event_handler`), a `session_id`-keyed
+generation fallback, and re-asserting the rollout generation into the session on every
+tool/verify response so the multi-server agent's chained `resources_server_cookies` keeps
+carrying our namespaced session cookie end to end (otherwise a later call arrives
+session-less and the fence rejects it). Set `SHINKEN_NG_DEBUG=1` to trace inbound
+cookies/session per request when adapting to a different agent.
+
 ## 5. Local optimizer-step smoke (not GRPO equivalence)
 
 NeMo RL is the production trainer, but it needs CUDA. [`local_grpo.py`](local_grpo.py)

examples/nemo_gym/app.py

@@ -55,6 +55,10 @@ def engine_factory(_config: object) -> ShinkenComputerEngine:
             lifecycle[argument] = float(os.environ[env_name])
     if "SHINKEN_MAX_GOLDENS" in os.environ:
         lifecycle["max_goldens"] = int(os.environ["SHINKEN_MAX_GOLDENS"])
+    if "SHINKEN_SCORER_ERROR_REWARD" in os.environ:
+        # RL-tolerant scoring: a corpus task whose reward.py crashes on the unsolved state
+        # scores this value instead of aborting the collection batch (default: strict raise).
+        lifecycle["scorer_error_reward"] = float(os.environ["SHINKEN_SCORER_ERROR_REWARD"])
     if "SHINKEN_MAX_PENDING_CLEANUP" in os.environ:
         lifecycle["max_pending_cleanup"] = int(
             os.environ["SHINKEN_MAX_PENDING_CLEANUP"]

sdk/python/src/shinken/integrations/nemo_gym.py

@@ -44,6 +44,7 @@
 import json
 import logging
 import math
+import os
 import secrets
 import threading
 import time
@@ -287,6 +288,7 @@ def __init__(
         reap_interval_s: float = 30.0,
         max_pending_cleanup: int = 64,
         cleanup_retry_batch: int = 16,
+        scorer_error_reward: float | None = None,
         env_factory: Any = ShinkenCuaGymEnv,
     ) -> None:
         if type(max_goldens) is not int or max_goldens < 0:
@@ -318,6 +320,13 @@ def __init__(
         self.reap_interval_s = lifetimes["reap_interval_s"]
         self.max_pending_cleanup = max_pending_cleanup
         self.cleanup_retry_batch = cleanup_retry_batch
+        # None keeps the strict eval contract (a broken scorer is a typed fault). A float
+        # makes verify tolerant for RL collection over messy corpora: a reward.py that
+        # crashes/emits no reward (e.g. CUA-Gym self-tests that `assert reward == 1.0` on
+        # the unsolved state) scores this value with a warning instead of aborting the batch.
+        self.scorer_error_reward = (
+            None if scorer_error_reward is None else float(scorer_error_reward)
+        )
         self.env_factory = env_factory
         self._rollouts: dict[str, _Rollout] = {}
         self._goldens: dict[str, _Golden] = {}
@@ -496,18 +505,36 @@ def _golden_lease(self, task: CuaGymTask, env: ShinkenCuaGymEnv) -> Iterator[Any
             self._evict_goldens()
 
     def verify(self, session_id: str, *, generation: int) -> float:
-        """Score the rollout with the bundle's ``reward.py`` and tear the replica down."""
+        """Score the rollout with the bundle's ``reward.py`` and tear the replica down.
+
+        A scorer fault is a typed error by default; with ``scorer_error_reward`` set it is
+        logged and scored as that value so one badly-authored corpus task cannot abort an
+        RL collection batch."""
         with self._operation(session_id=session_id):
             with self._session_lock(session_id):
                 rollout = self._current_rollout(session_id, generation)
                 with self._cleanup_admission():
                     try:
-                        return rollout.env.evaluate()
+                        return self._score(rollout)
                     finally:
                         detached = self._detach_rollout(session_id, rollout)
                         if detached is not None:
                             self._close_env(detached.env)
 
+    def _score(self, rollout: _Rollout) -> float:
+        if self.scorer_error_reward is None:
+            return rollout.env.evaluate()
+        try:
+            return rollout.env.evaluate()
+        except CuaGymError as exc:
+            _LOG.warning(
+                "reward.py fault on task %r scored as %s (RL-tolerant): %s",
+                getattr(rollout.task, "task_id", "?"),
+                self.scorer_error_reward,
+                exc,
+            )
+            return self.scorer_error_reward
+
     def end(self, session_id: str, *, generation: int) -> None:
         """Tear down a rollout's replica without scoring (idempotent)."""
         with self._operation(fail_if_closed=False, session_id=session_id) as admitted:
@@ -694,6 +721,17 @@ def _session_lock(self, session_id: str) -> threading.Lock:
     def _golden_lock(self, task_id: str) -> threading.Lock:
         return self._golden_locks[hash(task_id) % len(self._golden_locks)]
 
+    def current_generation(self, session_id: str) -> int | None:
+        """The generation of the session's live rollout, or ``None`` when none is active.
+
+        Lets a transport that cannot reliably thread the per-rollout generation (e.g. a
+        multi-server SessionMiddleware whose ``session`` cookie collides across servers,
+        while the ``session_id`` still threads) recover the fence value from the
+        reliably-keyed ``session_id``. Single rollout per session, so this is unambiguous."""
+        with self._lock:
+            rollout = self._rollouts.get(session_id)
+            return rollout.generation if rollout is not None else None
+
     def _current_rollout(self, session_id: str, generation: int) -> _Rollout:
         with self._lock:
             rollout = self._rollouts.get(session_id)
@@ -1051,18 +1089,64 @@ def extract_task_id(payload: Mapping[str, Any]) -> str | None:
     return None
 
 
+def _install_request_tracer(app: Any, session_id_key: str) -> None:
+    """SHINKEN_NG_DEBUG diagnostic: log every inbound request's path, the cookies it
+    carried, and the session_id/generation the server resolved — to pin down where a
+    multi-server session cookie drops the per-rollout state. Off unless the env is set."""
+    import sys
+
+    @app.middleware("http")
+    async def _trace(request: Any, call_next: Any) -> Any:
+        cookie_names = sorted(request.cookies.keys())
+        try:
+            session = request.session
+            sid = session.get(session_id_key)
+            gen = session.get(_SESSION_GENERATION_KEY)
+        except Exception as exc:  # noqa: BLE001 — diagnostic must never break the request
+            sid, gen = f"<no-session:{exc}>", None
+        print(
+            f"[NG_TRACE] {request.method} {request.url.path} cookies={cookie_names} "
+            f"session_id={sid!r} generation={gen!r}",
+            file=sys.stderr,
+            flush=True,
+        )
+        return await call_next(request)
+
+
 def _install_engine_shutdown(app: Any, engine: ShinkenComputerEngine) -> None:
-    """Bind active maintenance and engine-owned resources to the web app lifetime."""
+    """Bind active maintenance and engine-owned resources to the web app lifetime.
+
+    Starlette's classic ``add_event_handler``/``on_event`` surface was removed in newer
+    releases (lifespan-only apps), so when the app doesn't expose it, wrap the router's
+    lifespan context instead — same lifetime contract on both API generations."""
     import asyncio
+    import contextlib
 
     def start_engine() -> None:
         engine.start_maintenance()
 
     async def close_engine() -> None:
         await asyncio.to_thread(engine.close)
 
-    app.add_event_handler("startup", start_engine)
-    app.add_event_handler("shutdown", close_engine)
+    add_handler = getattr(app, "add_event_handler", None)
+    if callable(add_handler):
+        add_handler("startup", start_engine)
+        add_handler("shutdown", close_engine)
+        return
+
+    router = app.router
+    inner = router.lifespan_context
+
+    @contextlib.asynccontextmanager
+    async def lifespan_with_engine(app_: Any) -> Any:
+        start_engine()
+        try:
+            async with inner(app_) as state:
+                yield state
+        finally:
+            await close_engine()
+
+    router.lifespan_context = lifespan_with_engine
 
 
 def _request_generation(request: Any) -> int:
@@ -1130,12 +1214,14 @@ def setup_webserver(self) -> FastAPI:
             # live replicas and golden snapshots. Tie those lifetimes together so normal
             # server shutdown cannot orphan provider resources.
             _install_engine_shutdown(app, self._engine)
+            if os.environ.get("SHINKEN_NG_DEBUG"):
+                _install_request_tracer(app, SESSION_ID_KEY)
             return app
 
         def _make_tool_route(self, name: str):
             async def route(request: Request) -> PlainTextResponse:
                 session_id = request.session[SESSION_ID_KEY]
-                generation = _request_generation(request)
+                generation = self._resolve_generation(request, session_id)
                 args = await request.json()
                 out = await asyncio.to_thread(
                     self._engine.tool,
@@ -1164,12 +1250,43 @@ async def seed_session(  # type: ignore[override]
                 self._engine.seed, session_id, task_id, generation=generation
             )
             request.session[_SESSION_GENERATION_KEY] = seeded["generation"]
+            if os.environ.get("SHINKEN_NG_DEBUG"):
+                import sys
+
+                print(
+                    f"[NG_TRACE] SEED session_id={session_id!r} task={task_id!r} "
+                    f"generation={seeded['generation']!r}",
+                    file=sys.stderr,
+                    flush=True,
+                )
             return BaseSeedSessionResponse()
 
         async def verify(self, request: Request, body: BaseVerifyRequest) -> BaseVerifyResponse:
             session_id = request.session[SESSION_ID_KEY]
-            generation = _request_generation(request)
+            generation = self._resolve_generation(request, session_id)
             reward = await asyncio.to_thread(self._engine.verify, session_id, generation=generation)
             return BaseVerifyResponse(**body.model_dump(), reward=reward)
 
+        def _resolve_generation(self, request: Request, session_id: str) -> int:
+            """The rollout generation for a fenced tool/verify call, re-asserted into the
+            session so Starlette re-emits the Set-Cookie.
+
+            Value precedence: the client session's threaded value, else the engine's
+            active generation for this ``session_id`` (recovering from a transport that
+            drops the per-rollout cookie but keeps ``session_id``). The re-assert is the
+            load-bearing part over a multi-server agent: each server namespaces its own
+            session cookie and the agent chains ``resources_server_cookies`` from one tool
+            response to the next, so a response that does NOT re-emit our cookie breaks
+            the chain and the next call arrives session-less. Writing the generation back
+            marks the session modified on every call, keeping the cookie alive end to end."""
+            generation = request.session.get(_SESSION_GENERATION_KEY)
+            if type(generation) is not int:  # missing / non-int (bool) — recover it
+                generation = self._engine.current_generation(session_id)
+                if generation is None:
+                    raise CuaGymError(
+                        "session carries no Shinken rollout generation — seed_session first"
+                    )
+            request.session[_SESSION_GENERATION_KEY] = generation
+            return generation
+
     return ShinkenComputerResourcesServer