diff --git a/CHANGELOG.md b/CHANGELOG.md index c519661..e1added 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/). ### Added +- `LiveRunnerRegistration.stop_session()` to stop reserved sessions +- `LiveRunnerSessionEvent.session_token` field for session management - `call_runner(on_payment_session=...)` callback for each new payment session - `LivePaymentSession.manifest_id` for payment attribution - Flexible request body types on call_runner and selectors. diff --git a/src/livepeer_gateway/live_runner.py b/src/livepeer_gateway/live_runner.py index 82c9e77..48e77ed 100644 --- a/src/livepeer_gateway/live_runner.py +++ b/src/livepeer_gateway/live_runner.py @@ -85,6 +85,7 @@ class LiveRunnerSessionEvent: event: Literal["reserved", "released"] timestamp: str | None raw: dict[str, Any] + session_token: str = "" LiveRunnerSessionCallback = Callable[[LiveRunnerSessionEvent], None | Awaitable[None]] @@ -318,7 +319,7 @@ def __init__( self._unregister_on_close = unregister_on_close self._on_session_reserve = on_session_reserve self._on_session_release = on_session_release - self._active_session_ids: list[str] = [] + self._active_sessions: dict[str, str] = {} self.o2r_channel: LiveRunnerTrickleChannel | None = None self._o2r_reader: ChannelReader | None = None self._closed = False @@ -332,8 +333,7 @@ async def start(self) -> LiveRunnerRegistration: @property def active_session_ids(self) -> tuple[str, ...]: - # Return an immutable snapshot; internal storage stays list-backed to preserve reservation order. - return tuple(self._active_session_ids) + return tuple(self._active_sessions) async def close(self) -> None: self._closed = True @@ -379,6 +379,23 @@ async def __aenter__(self) -> LiveRunnerRegistration: async def __aexit__(self, exc_type: object, exc: object, tb: object) -> None: await self.close() + async def stop_session(self, session_id: str) -> None: + """Stop a reserved session using its cached token.""" + session_id = session_id.strip() + token = self._active_sessions.get(session_id) + if not token: + raise LivepeerGatewayError(f"Live runner session {session_id!r} has no session token") + await _post_empty( + _join_endpoint( + self.orchestrator_url, + f"/runner/{quote(self.runner_id, safe='')}/session/{quote(session_id, safe='')}/stop", + ), + headers={"Livepeer-Session-Token": token}, + timeout=self._timeout, + ) + if self._active_sessions.get(session_id) == token: + self._release_session_id(session_id) + async def create_trickle_channels( self, session: str | LiveRunnerSessionRequest, @@ -570,8 +587,12 @@ async def _handle_o2r_message(self, message: dict[str, Any]) -> None: session_id = session_id.strip() typed_event = cast(Literal["reserved", "released"], event) + session_token = "" if typed_event == "reserved": - self._reserve_session_id(session_id) + token = message.get("session_token") + session_token = token.strip() if isinstance(token, str) else "" + if session_token or session_id not in self._active_sessions: + self._active_sessions[session_id] = session_token else: self._release_session_id(session_id) @@ -585,6 +606,7 @@ async def _handle_o2r_message(self, message: dict[str, Any]) -> None: event=typed_event, timestamp=timestamp if isinstance(timestamp, str) else None, raw=message, + session_token=session_token, ) try: result = callback(session_event) @@ -593,14 +615,8 @@ async def _handle_o2r_message(self, message: dict[str, Any]) -> None: except Exception: _LOG.exception("Live runner %s callback failed for session %s", typed_event, session_id) - def _reserve_session_id(self, session_id: str) -> None: - if session_id not in self._active_session_ids: - self._active_session_ids.append(session_id) - def _release_session_id(self, session_id: str) -> None: - try: - self._active_session_ids.remove(session_id) - except ValueError: + if self._active_sessions.pop(session_id, None) is None: _LOG.debug("Live runner session %s was already released", session_id) diff --git a/tests/test_live_runner.py b/tests/test_live_runner.py index 9492757..32a58d6 100644 --- a/tests/test_live_runner.py +++ b/tests/test_live_runner.py @@ -857,7 +857,97 @@ async def close(self, **_kwargs: object) -> None: await asyncio.gather(self._task, return_exceptions=True) +def _registration(*, runner_id: str = "runner-1", **kwargs) -> LiveRunnerRegistration: + return LiveRunnerRegistration( + orchestrator_url="http://localhost:8935", + secret="secret", + runner_id=runner_id, + runner_url="http://localhost:9000", + app="test", + price_info=LiveRunnerPriceInfo(0), + **kwargs, + ) + + class TestLiveRunnerRegistration: + async def test_reserve_callback_can_stop_before_application_request(self) -> None: + released_events: list[LiveRunnerSessionEvent] = [] + + async def on_reserve(event: LiveRunnerSessionEvent) -> None: + assert event.session_token == "session-token" + await reg.stop_session(event.session_id) + + reg = _registration( + runner_id="runner/1", + timeout=12.0, + on_session_reserve=on_reserve, + on_session_release=released_events.append, + ) + reg.orchestrator_url = "https://service.example.com/api" + stop = mock.AsyncMock() + with mock.patch.object(live_runner, "_post_empty", stop): + await reg._handle_o2r_message( + {"event": "reserved", "session": " session/1 ", "session_token": " session-token "} + ) + stop.assert_awaited_once_with( + "https://service.example.com/api/runner/runner%2F1/session/session%2F1/stop", + headers={"Livepeer-Session-Token": "session-token"}, + timeout=12.0, + ) + assert reg.active_session_ids == () + assert released_events == [] + await reg._handle_o2r_message({"event": "released", "session": "session/1"}) + + assert [event.session_id for event in released_events] == ["session/1"] + assert released_events[0].session_token == "" + + async def test_stop_session_failure_retains_credentials(self) -> None: + reg = _registration() + await reg._handle_o2r_message( + {"event": "reserved", "session": "session-1", "session_token": "token"} + ) + with ( + mock.patch.object(live_runner, "_post_empty", side_effect=LivepeerHTTPError(503, "url", "unavailable")), + pytest.raises(LivepeerHTTPError), + ): + await reg.stop_session("session-1") + assert reg._active_sessions == {"session-1": "token"} + + async def test_stop_session_preserves_replacement_reservation(self) -> None: + reg = _registration() + await reg._handle_o2r_message( + {"event": "reserved", "session": "session-1", "session_token": "old-token"} + ) + + async def replace_before_response(*args, **kwargs): + await reg._handle_o2r_message({"event": "released", "session": "session-1"}) + await reg._handle_o2r_message( + {"event": "reserved", "session": "session-1", "session_token": "new-token"} + ) + + with mock.patch.object(live_runner, "_post_empty", side_effect=replace_before_response): + await reg.stop_session("session-1") + assert reg._active_sessions == {"session-1": "new-token"} + + @pytest.mark.parametrize( + "token_fields", [None, {}, {"session_token": " "}], + ids=["unknown-session", "missing-token", "blank-token"], + ) + async def test_stop_session_requires_token( + self, token_fields: dict[str, object] | None + ) -> None: + reg = _registration() + if token_fields is not None: + await reg._handle_o2r_message( + {"event": "reserved", "session": "session-1", **token_fields} + ) + with ( + mock.patch.object(live_runner, "_post_empty") as stop, + pytest.raises(LivepeerGatewayError, match="has no session token"), + ): + await reg.stop_session("session-1") + stop.assert_not_awaited() + async def test_initial_heartbeat_starts_o2r_channel_reader_callback_once( self, ) -> None: @@ -984,12 +1074,17 @@ def _post_json( { "event": "reserved", "session": "session-b", + "session_token": "original-token", "timestamp": "2026-05-20T17:00:00Z", } ) await on_event({"event": "reserved", "session": "session-a"}) + await on_event( + {"event": "reserved", "session": "session-b", "session_token": "refreshed-token"} + ) await on_event({"event": "reserved", "session": "session-b"}) assert reg.active_session_ids == ("session-b", "session-a") + assert reg._active_sessions["session-b"] == "refreshed-token" await on_event( { @@ -1014,6 +1109,7 @@ def _post_json( "session-b", "session-a", "session-b", + "session-b", ] assert sync_events[0].event == "reserved" assert sync_events[0].timestamp == "2026-05-20T17:00:00Z"