fix(api): stop a service's sys.exit from killing the serve process

Catch SystemExit at service import, in the heartbeat task, the 5xx middleware, search, samples and chapters, and answer it as SERVICE_ERROR.

docs/dev/rest-api/errors.md

@@ -127,7 +127,7 @@ while the server restarts.
 | `TRACK_NOT_FOUND` | 404 | A referenced track ID is not present in the remote session. |
 | `CONFLICT` | 409 | The target resource is in a state that disallows the action (retrying a running job, prioritising a non-queued job, clearing cache during an active download). |
 | `SERVER_DRAINING` | 409 | The server is about to restart and starts no new remote session or download job. Open work continues. Retryable after the restart. A `dl --remote` client waits for the restart and opens a new remote session. |
-| `SERVICE_ERROR` | 422 | The streaming service's API returned an error, or the service does not support the requested operation. |
+| `SERVICE_ERROR` | 422 | The streaming service's API returned an error, or the service does not support the requested operation, or the service code exited (`sys.exit`). |
 | `DRM_ERROR` | 422 | DRM or license acquisition failed. |
 | `RATE_LIMITED` | 429 | The streaming service is rate-limiting requests. Retryable. |
 

tests/core/test_services_reload.py

@@ -121,3 +121,11 @@ def test_overdue_staged_service_applies_while_busy(service_dir: Path, monkeypatc
     assert services.apply_pending(busy={"FOO"}) == ["FOO"]
     assert services.MODULES["FOO"].VERSION == 2
     assert services.PENDING == set() and "FOO" not in services.PENDING_SINCE
+
+
+def test_service_that_exits_on_import_fails_its_reload_only(service_dir: Path):
+    """A service calling sys.exit at import time is a load error, not a stop of the process."""
+    (service_dir / "FOO" / "__init__.py").write_text("import sys\nsys.exit(1)\n")
+    os.utime(service_dir / "FOO" / "__init__.py", (time.time() + 9, time.time() + 9))
+    services.reload_services(["FOO"])
+    assert services.LOAD_ERRORS and services.LOAD_ERRORS[-1].startswith("FOO: failed to import - SystemExit")

tests/remote/unit/test_errors.py

@@ -216,3 +216,9 @@ async def test_outer_middleware_never_answers_5xx(aiohttp_client) -> None:
     resp = await client.get("/subapp")
     assert resp.status == 422 and (await resp.json())["message"] == "device error"
     assert (await client.get("/missing")).status == 404
+
+
+def test_categorize_system_exit_is_a_service_error() -> None:
+    err = categorize_exception(SystemExit(1), {"service": "EXAMPLE"})
+    assert err.error_code == APIErrorCode.SERVICE_ERROR
+    assert err.message == "Service EXAMPLE exited"

tests/remote/unit/test_request_cache.py

@@ -4,8 +4,10 @@ The directory holds only what the client sent, goes back to the client when it l
 itself, and is always removed when the request ends.
 """
 
+import asyncio
 import base64
 import json
+import types
 import zlib
 
 import click
@@ -123,6 +125,17 @@ async def test_cache_dir_removed_on_error(seen, tmp_path):
     assert not request_dirs(tmp_path)
 
 
+@pytest.mark.asyncio
+async def test_search_answers_a_service_sys_exit_with_an_error(seen, tmp_path, monkeypatch):
+    def search(self):
+        raise SystemExit(1)
+
+    monkeypatch.setattr(_CacheService, "search", search)
+    response = await handlers.search_handler({"service": "EXAMPLE", "query": "q"})
+    assert json.loads(response.body)["error_code"] == "SERVICE_ERROR"
+    assert not request_dirs(tmp_path)
+
+
 @pytest.mark.asyncio
 async def test_invalid_client_cache_still_removes_the_dir(seen, tmp_path):
     data = {"service": "EXAMPLE", "query": "q", "cache": {"tokens": 1}}
@@ -192,3 +205,24 @@ async def test_list_tracks_service_sees_dl_params(seen):
 async def test_search_service_gets_no_dl_params(seen):
     await handlers.search_handler({"service": "EXAMPLE", "query": "q", "dl_params": {"a_lang": ["es-419"]}})
     assert "a_lang" not in seen["parent_params"]
+
+
+@pytest.mark.asyncio
+async def test_samples_answers_a_service_exit_in_decrypts_samples(monkeypatch):
+    class Service:
+        def decrypts_samples(self, track):
+            raise SystemExit(1)
+
+    session = types.SimpleNamespace(
+        tracks={"t": object()}, service_instance=Service(), service_tag="EXAMPLE", lock=asyncio.Lock()
+    )
+
+    async def get_session(session_id, request=None):
+        return session
+
+    monkeypatch.setattr(handlers, "get_validated_session", get_session)
+    monkeypatch.setattr(handlers, "require_authenticated", lambda s: None)
+    monkeypatch.setattr(handlers, "server_cdm_allowed", lambda request, tag: True)
+    with pytest.raises(handlers.APIError) as e:
+        await handlers.session_samples_handler("sid", "t", "1", b"")
+    assert e.value.error_code.value == "SERVICE_ERROR"

tests/remote/unit/test_routes.py

@@ -175,6 +175,41 @@ async def test_heartbeat_keeps_slow_request_alive(aiohttp_client, monkeypatch: p
     assert resp.status == 200 and body["status"] == "error" and body["error_code"] == "SERVICE_ERROR"
 
 
+@pytest.mark.parametrize("delay", [0, 0.2])
+async def test_heartbeat_answers_a_service_exit_and_keeps_serving(
+    aiohttp_client, monkeypatch: pytest.MonkeyPatch, delay: float
+) -> None:
+    """A SystemExit inside the heartbeat task becomes a SERVICE_ERROR instead of stopping the event loop."""
+    monkeypatch.setattr(routes, "HEARTBEAT_INTERVAL", 0.05)
+
+    @heartbeat
+    async def exits(request: web.Request) -> web.Response:
+        await asyncio.sleep(delay)
+        raise SystemExit(1)
+
+    app = web.Application()
+    app["debug_api"] = False
+    app.router.add_get("/exits", api_handler(exits))
+    client = await aiohttp_client(app)
+
+    for _ in range(2):
+        resp = await client.get("/exits")
+        assert json.loads(await resp.text())["error_code"] == "SERVICE_ERROR"
+
+
+async def test_unguarded_route_answers_a_service_exit(aiohttp_client) -> None:
+    async def exits(request: web.Request) -> web.Response:
+        raise SystemExit(1)
+
+    app = web.Application(middlewares=[routes.no_5xx_middleware])
+    app.router.add_get("/exits", exits)
+    client = await aiohttp_client(app)
+
+    for _ in range(2):
+        resp = await client.get("/exits")
+        assert resp.status == 422 and (await resp.json())["error_code"] == "SERVICE_ERROR"
+
+
 async def test_heartbeat_late_error_counts_as_rejected(aiohttp_client, monkeypatch: pytest.MonkeyPatch) -> None:
     """Stats count a late failure by its real status and the bytes the stream sent (with framing), not as an empty 200."""
     from unshackle.core.api.stats import stats, stats_middleware

unshackle/core/api/errors.py

@@ -179,7 +179,7 @@ def build_error_response(
 
 
 def categorize_exception(
-    exc: Exception,
+    exc: BaseException,
     context: dict[str, Any] | None = None,
 ) -> APIError:
     """
@@ -196,6 +196,14 @@ def categorize_exception(
         APIError with appropriate error code and details
     """
     context = context or {}
+    if isinstance(exc, SystemExit):
+        exited = f"Service {context['service']} exited" if context.get("service") else "Service exited"
+        return APIError(
+            error_code=APIErrorCode.SERVICE_ERROR,
+            message=redact_secrets(getattr(exc, "message", None)) or exited,
+            details={**context, "reason": "service_exit"},
+            retryable=False,
+        )
     root = exc.__cause__ or exc
     exc_str = f"{exc} {root}".lower()
     exc_type = type(root).__name__
@@ -296,7 +304,7 @@ def categorize_exception(
 
 
 def handle_api_exception(
-    exc: Exception,
+    exc: BaseException,
     context: dict[str, Any] | None = None,
     debug_mode: bool = False,
     extra_debug_info: dict[str, Any] | None = None,

unshackle/core/api/handlers.py

@@ -15,7 +15,7 @@ from datetime import date as date_
 from functools import partial
 from http.cookiejar import CookieJar
 from pathlib import Path
-from typing import Any, Callable, Dict, Iterable, List, Optional, cast
+from typing import Any, Callable, Dict, Iterable, List, Optional, Set, cast
 
 import click
 from aiohttp import web
@@ -1088,6 +1088,8 @@ def live_licence_refusal(session: Any, track: Any) -> Optional[APIError]:
 
 
 _account_counters: Dict[str, int] = {}
+# the event loop holds only weak references to tasks, so a running login task needs a strong one
+_auth_tasks: Set["asyncio.Task[None]"] = set()
 
 
 def validate_server_accounts() -> Dict[str, Any]:
@@ -1885,6 +1887,14 @@ async def search_handler(data: Dict[str, Any], request: Optional[web.Request] =
     try:
         results = await asyncio.to_thread(run_service_search, data, normalized_service, query, request, request_cache)
         return request_cache.json_response({"results": results, "count": len(results)})
+    except APIError:
+        raise
+    except (Exception, SystemExit) as e:
+        log.exception("Error searching")
+        debug_mode = request.app.get("debug_api", False) if request else False
+        return handle_api_exception(
+            e, context={"operation": "search", "service": normalized_service}, debug_mode=debug_mode
+        )
     finally:
         request_cache.cleanup()
 
@@ -3987,7 +3997,9 @@ async def session_create_handler(data: Dict[str, Any], request: Optional[web.Req
                 if store.peek(session_id) is not session:
                     SessionStore.cleanup_cache_dir(session_cache_tag)
 
-        asyncio.create_task(run_auth())
+        auth_task = asyncio.create_task(run_auth())
+        _auth_tasks.add(auth_task)
+        auth_task.add_done_callback(_auth_tasks.discard)
 
         return web.json_response(
             {
@@ -4173,7 +4185,7 @@ async def session_tracks_handler(
             try:
                 chapters = await asyncio.to_thread(service_instance.get_chapters, title)
                 session.chapters_by_title[str(title_id)] = chapters if chapters else []
-            except (NotImplementedError, Exception):
+            except (Exception, SystemExit):
                 session.chapters_by_title[str(title_id)] = []
 
             video_tracks = sorted(tracks.videos, key=lambda t: t.bitrate or 0, reverse=True)
@@ -5271,14 +5283,14 @@ async def session_samples_handler(
             details={"track_id": track_id, "session_id": session_id},
         )
     service = session.service_instance
-    if not service.decrypts_samples(track):
-        raise APIError(APIErrorCode.INVALID_INPUT, f"The service does not decrypt the samples of track {track_id}")
     if not server_cdm_allowed(request, session.service_tag):
         raise APIError(
             APIErrorCode.FORBIDDEN, f"This API key may not decrypt {session.service_tag} samples on the server"
         )
 
     def run() -> web.Response:
+        if not service.decrypts_samples(track):
+            raise APIError(APIErrorCode.INVALID_INPUT, f"The service does not decrypt the samples of track {track_id}")
         decrypter = service.get_sample_decrypter(track, index)
         if isinstance(decrypter, (bytes, bytearray)):
             return web.json_response({"key": bytes(decrypter).hex()})

unshackle/core/api/routes.py

@@ -11,7 +11,13 @@ from aiohttp_swagger3 import SwaggerDocs, SwaggerInfo, SwaggerUiSettings
 
 from unshackle.core import __code_hash__, __version__
 from unshackle.core import services as services_module
-from unshackle.core.api.errors import APIError, APIErrorCode, build_error_response, handle_api_exception
+from unshackle.core.api.errors import (
+    APIError,
+    APIErrorCode,
+    build_error_response,
+    categorize_exception,
+    handle_api_exception,
+)
 from unshackle.core.api.handlers import (
     CORS_HEADERS,
     DASHBOARD_EVENTS_ROUTE,
@@ -104,12 +110,13 @@ async def no_5xx_middleware(
     """Answer an uncaught exception or any unsent 5xx (a CDM subapplication's own) as a 422; see errors.py."""
     try:
         response = await handler(request)
-    except Exception as e:
+    except (Exception, SystemExit) as e:
         if isinstance(e, web.HTTPException) and e.status < 500:
             raise
         # the exception text can hold a caller's key, so the answer gets a generic message
         log.exception(f"Unhandled error in {sanitize_log(request.path)}")
-        response = build_error_response(APIError(APIErrorCode.INTERNAL_ERROR, "Internal server error"))
+        error = categorize_exception(e) if isinstance(e, SystemExit) else None
+        response = build_error_response(error or APIError(APIErrorCode.INTERNAL_ERROR, "Internal server error"))
     if response.status >= 500 and not response.prepared:
         response.set_status(APIError.default_http_status(APIErrorCode.INTERNAL_ERROR))
     return response
@@ -147,9 +154,17 @@ def heartbeat(handler: Handler) -> Handler:
     work continues if the client disconnects, so the remote session state stays complete.
     """
 
+    async def guarded(request: web.Request) -> web.StreamResponse:
+        # a SystemExit leaving any task is re-raised into the event loop, so it must stop inside the task
+        try:
+            return await handler(request)
+        except SystemExit as e:
+            log.exception(f"Service exited in {sanitize_log(request.path)}")
+            raise categorize_exception(e) from e
+
     @functools.wraps(handler)
     async def wrapper(request: web.Request) -> web.StreamResponse:
-        task = asyncio.ensure_future(handler(request))
+        task = asyncio.ensure_future(guarded(request))
         done, _ = await asyncio.wait({task}, timeout=HEARTBEAT_INTERVAL)
         if done:
             return task.result()

unshackle/core/services.py

@@ -97,7 +97,7 @@ def load_service(path: Path) -> object:
     tag = path.parent.stem
     try:
         module = import_module_by_path(path)
-    except Exception as e:
+    except (Exception, SystemExit) as e:
         hint = compiled_import_hint(e.name, path) if isinstance(e, ModuleNotFoundError) and e.name else ""
         raise RuntimeError(f"{tag}: failed to import - {type(e).__name__}: {e}{hint} ({path.name})") from e
     try: