[framework, gateway] fix: settle cancelled session creation

tests/uni_agent/framework/test_generate_sequences_on_cpu.py

@@ -2126,3 +2126,35 @@ async def scored_runner(**kwargs):
         },
     ]
     assert "runner_reward_info" not in framework._trajectory_meta(trajectories[0])
+
+
+@pytest.mark.cpu
+@pytest.mark.level0
+@pytest.mark.asyncio
+async def test_runner_error_survives_gateway_abort_failure():
+    async def failing_runner(**kwargs):
+        raise ValueError("runner exploded")
+
+    class _GatewayManager(_FakeGatewayManager):
+        async def abort_session(self, session_id: str) -> None:
+            await super().abort_session(session_id)
+            raise RuntimeError("abort failed")
+
+    runtime = _GatewayManager({"session-sample-0-rollout-0": [_trajectory()]})
+    framework = await _build_framework_with_agent_runners(
+        agent_runners={"runner": _inline_runner_config(failing_runner)},
+        gateway_manager=runtime,
+    )
+
+    # The abort failure is logged; the runner error is what callers see.
+    with pytest.raises(ValueError, match="runner exploded"):
+        await framework._run_agent_episode(
+            sample_fields={"raw_prompt": [], "uid": "uid-0"},
+            sample_index=0,
+            session_index=0,
+            global_steps=7,
+            runner_name="runner",
+            runner_config=framework.runner_registry["runner"],
+            sampling_params={},
+        )
+    assert len(runtime.aborted_sessions) == 1

tests/uni_agent/gateway/test_gateway_manager_on_cpu.py

@@ -325,3 +325,201 @@ async def generate(
         ]
     finally:
         await manager.shutdown()
+
+
+class _LifecycleGateway:
+    """Controllable remote-method stub for close failures."""
+
+    def __init__(self, *, finalize_error=None, abort_error=None):
+        self.finalize_error = finalize_error
+        self.abort_error = abort_error
+        self.abort_calls = []
+        self.finalize_calls = []
+
+    async def _abort(self, session_id):
+        self.abort_calls.append(session_id)
+        if self.abort_error is not None:
+            raise self.abort_error
+
+    async def _finalize(self, session_id):
+        self.finalize_calls.append(session_id)
+        if self.finalize_error is not None:
+            raise self.finalize_error
+        return ["trajectory"]
+
+    abort_session = property(lambda self: _FakeRemoteMethod(self._abort))
+    finalize_session = property(lambda self: _FakeRemoteMethod(self._finalize))
+
+
+def _new_manager(gateway, *, session_ids=()):
+    from uni_agent.gateway.manager import GatewayManager
+
+    manager = GatewayManager.__new__(GatewayManager)
+    manager.gateways = [gateway]
+    manager.active_sessions_per_gateway = [len(session_ids)]
+    manager._session_to_gateway_index = dict.fromkeys(session_ids, 0)
+    return manager
+
+
+class _DetachedRemoteMethod:
+    """Ray-like call: cancelling the caller's await does not stop the actor method."""
+
+    def __init__(self, fn):
+        self._fn = fn
+
+    def remote(self, *args, **kwargs):
+        return asyncio.shield(asyncio.ensure_future(self._fn(*args, **kwargs)))
+
+
+class _SlowCreateGateway(_LifecycleGateway):
+    def __init__(self, *, create_error=None):
+        super().__init__()
+        self.create_error = create_error
+        self.created = []
+        self.create_started = asyncio.Event()
+        self.create_finished = asyncio.Event()
+        self.release_create = asyncio.Event()
+
+    async def _create(self, session_id, **kwargs):
+        self.create_started.set()
+        await self.release_create.wait()
+        self.create_finished.set()
+        if self.create_error is not None:
+            raise self.create_error
+        self.created.append(session_id)
+        return session_id
+
+    create_session = property(lambda self: _DetachedRemoteMethod(self._create))
+    abort_session = property(lambda self: _DetachedRemoteMethod(self._abort))
+
+
+@pytest.mark.cpu
+@pytest.mark.level0
+@pytest.mark.asyncio
+@pytest.mark.parametrize("create_error", [None, RuntimeError("create failed")])
+async def test_gateway_manager_cancelled_create_aborts_what_the_actor_created(create_error):
+    """After the caller cancels, only the manager can still abort the actor-side session."""
+    gateway = _SlowCreateGateway(create_error=create_error)
+    manager = _new_manager(gateway)
+
+    create = asyncio.create_task(manager.create_session("session-cancelled"))
+    await gateway.create_started.wait()
+    create.cancel()
+    await asyncio.sleep(0)
+    gateway.release_create.set()
+    with pytest.raises(asyncio.CancelledError):
+        await create
+    await asyncio.wait_for(gateway.create_finished.wait(), timeout=1)
+
+    assert gateway.abort_calls == ([] if create_error else ["session-cancelled"])
+    assert manager._session_to_gateway_index == {}
+    assert manager.active_sessions_per_gateway == [0]
+
+
+@pytest.mark.cpu
+@pytest.mark.level0
+@pytest.mark.asyncio
+async def test_gateway_manager_cancelled_create_retains_owner_after_abort_failure():
+    gateway = _SlowCreateGateway()
+    gateway.abort_error = RuntimeError("session abort failed before removal")
+    manager = _new_manager(gateway)
+    create = asyncio.create_task(manager.create_session("session-cancelled"))
+    await gateway.create_started.wait()
+    create.cancel()
+    gateway.release_create.set()
+    with pytest.raises(asyncio.CancelledError):
+        await create
+
+    assert manager._session_to_gateway_index == {"session-cancelled": 0}
+    assert manager.active_sessions_per_gateway == [1]
+    gateway.abort_error = None
+    await manager.abort_session("session-cancelled")
+    assert gateway.abort_calls == ["session-cancelled"] * 2
+    assert manager.active_sessions_per_gateway == [0]
+
+
+@pytest.mark.cpu
+@pytest.mark.level0
+@pytest.mark.asyncio
+async def test_gateway_manager_repeated_cancellation_still_settles_remote_create():
+    gateway = _SlowCreateGateway()
+    manager = _new_manager(gateway)
+    create = asyncio.create_task(manager.create_session("session-cancelled"))
+    await gateway.create_started.wait()
+    for _ in range(3):
+        create.cancel()
+        await asyncio.sleep(0)
+    still_waiting = not create.done()
+    gateway.release_create.set()
+    with pytest.raises(asyncio.CancelledError):
+        await create
+    await asyncio.wait_for(gateway.create_finished.wait(), timeout=1)
+
+    assert still_waiting
+    assert gateway.abort_calls == ["session-cancelled"]
+    assert manager._session_to_gateway_index == {}
+    assert manager.active_sessions_per_gateway == [0]
+
+
+@pytest.mark.cpu
+@pytest.mark.level0
+@pytest.mark.asyncio
+async def test_gateway_manager_close_failure_retains_mapping_and_owner():
+    close_error = RuntimeError("actor-close-unknown")
+    gateway = _LifecycleGateway(finalize_error=close_error)
+    manager = _new_manager(gateway, session_ids=("session-close",))
+
+    with pytest.raises(RuntimeError, match="actor-close-unknown") as raised:
+        await manager.finalize_session("session-close")
+
+    assert raised.value is close_error
+    assert gateway.finalize_calls == ["session-close"]
+    assert manager._session_to_gateway_index == {"session-close": 0}
+    assert manager.active_sessions_per_gateway == [1]
+
+    await manager.abort_session("session-close")
+
+    assert gateway.abort_calls == ["session-close"]
+    assert manager._session_to_gateway_index == {}
+    assert manager.active_sessions_per_gateway == [0]
+
+
+@pytest.mark.cpu
+@pytest.mark.level0
+@pytest.mark.asyncio
+async def test_gateway_manager_rejects_duplicate_create_without_touching_owner():
+    gateway = _SlowCreateGateway()
+    manager = _new_manager(gateway)
+    create = asyncio.create_task(manager.create_session("same-id"))
+    await gateway.create_started.wait()
+
+    with pytest.raises(RuntimeError, match="already exists"):
+        await manager.create_session("same-id")
+    assert manager._session_to_gateway_index == {"same-id": 0}
+    assert manager.active_sessions_per_gateway == [1]
+    assert gateway.abort_calls == []
+
+    gateway.release_create.set()
+    assert await create == "same-id"
+    await manager.abort_session("same-id")
+    assert manager.active_sessions_per_gateway == [0]
+
+
+@pytest.mark.cpu
+@pytest.mark.level0
+@pytest.mark.asyncio
+async def test_gateway_manager_concurrent_closes_release_load_slot_once():
+    class _ConcurrentCloseGateway(_LifecycleGateway):
+        async def _finalize(self, session_id):
+            await asyncio.sleep(0)
+            return await super()._finalize(session_id)
+
+        async def _abort(self, session_id):
+            await asyncio.sleep(0)
+            await super()._abort(session_id)
+
+    gateway = _ConcurrentCloseGateway()
+    manager = _new_manager(gateway, session_ids=("session-close",))
+    await asyncio.gather(manager.finalize_session("session-close"), manager.abort_session("session-close"))
+    assert manager._session_to_gateway_index == {}
+    assert manager.active_sessions_per_gateway == [0]

uni_agent/framework/framework.py

@@ -869,7 +869,10 @@ async def _run_agent_episode(
                 raise
             except Exception:
                 logger.exception("session %s failed (runner=%s); aborting session", session_id, runner_name)
-                await self.gateway_manager.abort_session(session_id)
+                try:
+                    await self.gateway_manager.abort_session(session_id)
+                except Exception:
+                    logger.exception("session %s: Gateway abort failed after runner failure", session_id)
                 session_trace.finish(
                     runner_name=runner_name,
                     status="failure",

uni_agent/gateway/manager.py

@@ -8,12 +8,15 @@
 from __future__ import annotations
 
 import asyncio
+import logging
 
 import ray
 
 from uni_agent.gateway.config import GatewayActorConfig
 from verl.workers.rollout.llm_server import LLMServerClient
 
+logger = logging.getLogger(__name__)
+
 
 class GatewayManager:
     """Owns gateway actors and routes sessions to them.
@@ -73,7 +76,9 @@ def _get_gateway(self, session_id: str):
         return self.gateways[gateway_index], gateway_index
 
     async def create_session(self, session_id: str, **kwargs):
-        """Create a session on the least-loaded actor, record the route, and return its handle."""
+        """Create a session; cancellation settles the remote call and aborts only its result."""
+        if session_id in self._session_to_gateway_index:
+            raise RuntimeError(f"Session {session_id} already exists")
         gateway_index = self._select_gateway_index()
         gateway = self.gateways[gateway_index]
         # Reserve the slot synchronously, before the await. Sessions are created
@@ -85,26 +90,53 @@ async def create_session(self, session_id: str, **kwargs):
         self._session_to_gateway_index[session_id] = gateway_index
         self.active_sessions_per_gateway[gateway_index] += 1
         try:
-            return await gateway.create_session.remote(session_id=session_id, **kwargs)
+            create = asyncio.ensure_future(gateway.create_session.remote(session_id=session_id, **kwargs))
+            return await asyncio.shield(create)
+        except asyncio.CancelledError:
+            cleanup = asyncio.ensure_future(self._abort_cancelled_create(session_id, gateway_index, create))
+            # Ray create keeps running after caller cancellation. Every wait must
+            # stay shielded so repeated cancellation cannot orphan its session.
+            while not cleanup.done():
+                try:
+                    await asyncio.shield(cleanup)
+                except asyncio.CancelledError:
+                    continue
+            cleanup.result()
+            raise
         except BaseException:
-            self.active_sessions_per_gateway[gateway_index] -= 1
-            self._session_to_gateway_index.pop(session_id, None)
+            self._drop_route(session_id, gateway_index)
             raise
 
+    async def _abort_cancelled_create(self, session_id: str, gateway_index: int, create: asyncio.Future) -> None:
+        try:
+            await create
+        except (Exception, asyncio.CancelledError):
+            # A failed create owns no actor session; never abort a pre-existing one.
+            self._drop_route(session_id, gateway_index)
+            return
+        try:
+            await self.abort_session(session_id)
+        except (Exception, asyncio.CancelledError):
+            # A failed abort leaves this owner available for an explicit retry.
+            logger.exception("session %s: abort failed after cancelled create", session_id)
+
     async def finalize_session(self, session_id: str):
         """Finalize a session on its owning actor, release the route, and return its trajectories."""
         gateway, gateway_index = self._get_gateway(session_id)
         trajectories = await gateway.finalize_session.remote(session_id=session_id)
-        self._session_to_gateway_index.pop(session_id, None)
-        self.active_sessions_per_gateway[gateway_index] -= 1
+        self._drop_route(session_id, gateway_index)
         return trajectories
 
     async def abort_session(self, session_id: str) -> None:
         """Abort a routed session on its owning actor and release the route."""
         gateway, gateway_index = self._get_gateway(session_id)
         await gateway.abort_session.remote(session_id=session_id)
-        self._session_to_gateway_index.pop(session_id, None)
-        self.active_sessions_per_gateway[gateway_index] -= 1
+        self._drop_route(session_id, gateway_index)
+
+    def _drop_route(self, session_id: str, gateway_index: int) -> None:
+        # Concurrent closes must release the load slot at most once.
+        if self._session_to_gateway_index.pop(session_id, None) is not None:
+            self.active_sessions_per_gateway[gateway_index] -= 1
 
     async def shutdown(self) -> None:
         """Stop owned gateway actors and clear routing state."""