feat(remote): keep slow remote session requests alive with a heartbeat

Titles, tracks and licence stream a newline every 30s until the body is ready, so a slow service outlasts the client timeout and proxy idle limits; add a per-server timeout.

tests/remote/unit/test_extract_pssh_strict.py

@@ -5,6 +5,7 @@ loop licensed a PlayReady-only track as Widevine by accident. A Widevine device
 license PlayReady-only content, but only as a deliberate step in the batch loop.
 """
 
+import asyncio
 import json
 from types import SimpleNamespace
 
@@ -44,6 +45,7 @@ def test_extractor_is_type_strict():
 
 async def _run_batch(monkeypatch, tracks, cdm_type, calls):
     session = SimpleNamespace(
+        lock=asyncio.Lock(),
         service_tag="EXAMPLE",
         service_instance=SimpleNamespace(),
         served_keys={},

tests/remote/unit/test_input_bridge.py

@@ -107,3 +107,12 @@ def test_answered_only_after_an_accepted_response() -> None:
     assert bridge.submit_response("CODE") is True
     worker.join(timeout=2)
     assert bridge.answered is True
+
+
+def test_request_input_after_authentication_fails_fast() -> None:
+    """The client stops polling for prompts once authenticated, so a later prompt must not wait for an answer."""
+    bridge = InputBridge()
+    bridge.status = AuthStatus.AUTHENTICATED
+    with pytest.raises(RuntimeError, match="after authentication"):
+        bridge.request_input("Pick a profile", timeout=5)
+    assert bridge.status is AuthStatus.AUTHENTICATED

tests/remote/unit/test_remote_client.py

@@ -168,3 +168,43 @@ def test_timeout_raises_systemexit(client: RemoteClient) -> None:
     )
     with pytest.raises(SystemExit):
         client.get("/api/health")
+
+
+@responses.activate
+def test_late_error_after_heartbeat_is_fatal(client: RemoteClient) -> None:
+    """A heartbeat response carries a late failure as a 200 error body after leading newlines."""
+    body = "\n\n" + json.dumps({"status": "error", "error_code": "SERVICE_ERROR", "message": "late"})
+    responses.add(responses.GET, "http://srv:8786/api/session/s/titles", body=body, status=200)
+    with pytest.raises(SystemExit):
+        client.get("/api/session/s/titles")
+    assert client.request("get", "/api/session/s/titles", optional=True) == {}
+
+
+@responses.activate
+def test_heartbeat_body_parses_past_leading_newlines(client: RemoteClient) -> None:
+    responses.add(responses.GET, "http://srv:8786/api/session/s/titles", body='\n\n{"titles": []}', status=200)
+    assert client.get("/api/session/s/titles") == {"titles": []}
+
+
+def test_read_timeout_in_body_reports_timeout(client: RemoteClient, caplog: pytest.LogCaptureFixture) -> None:
+    """requests re-raises a read timeout in the body as a ConnectionError; the client still names the timeout."""
+    import requests
+    import urllib3
+
+    def stalled(*args: object, **kwargs: object) -> None:
+        raise requests.ConnectionError(urllib3.exceptions.ReadTimeoutError(None, "", "Read timed out."))
+
+    client.session.get = stalled  # type: ignore[method-assign]
+    with pytest.raises(SystemExit):
+        client.get("/api/session/s/titles")
+    assert "timed out after 120" in caplog.text
+
+
+@responses.activate
+def test_proxy_error_page_logs_its_title(client: RemoteClient, caplog: pytest.LogCaptureFixture) -> None:
+    page = "<!DOCTYPE html><html><head><title>example.com | 524: A timeout occurred</title></head><body>x</body></html>"
+    responses.add(responses.GET, "http://srv:8786/api/session/s/titles", body=page, status=524)
+    with pytest.raises(SystemExit):
+        client.get("/api/session/s/titles")
+    assert "HTTP 524: example.com | 524: A timeout occurred" in caplog.text
+    assert "DOCTYPE" not in caplog.text

tests/remote/unit/test_routes.py

@@ -7,12 +7,16 @@ stubbing the route table for selected paths.
 
 from __future__ import annotations
 
+import asyncio
+import json
+
 import pytest
 from aiohttp import web
 
+from unshackle.core.api import routes
 from unshackle.core.api.compression import compression_middleware
 from unshackle.core.api.errors import APIError, APIErrorCode
-from unshackle.core.api.routes import api_handler, cors_middleware, setup_routes
+from unshackle.core.api.routes import api_handler, cors_middleware, heartbeat, setup_routes
 
 pytestmark = pytest.mark.unit
 
@@ -131,6 +135,93 @@ async def test_api_handler_maps_apierror(aiohttp_client) -> None:
     assert body["details"] == {"x": 1}
 
 
+async def test_heartbeat_keeps_slow_request_alive(aiohttp_client, monkeypatch: pytest.MonkeyPatch) -> None:
+    """A fast route answers unchanged; a slow one streams newlines, then its body, and a late error arrives as 200."""
+    monkeypatch.setattr(routes, "HEARTBEAT_INTERVAL", 0.05)
+
+    @heartbeat
+    async def fast(request: web.Request) -> web.Response:
+        return web.json_response({"x": 1}, status=201)
+
+    @heartbeat
+    async def slow(request: web.Request) -> web.Response:
+        await asyncio.sleep(0.2)
+        return web.json_response({"titles": [1]})
+
+    @heartbeat
+    async def slow_fail(request: web.Request) -> web.Response:
+        await asyncio.sleep(0.2)
+        raise APIError(APIErrorCode.SERVICE_ERROR, "late")
+
+    app = web.Application()
+    app["debug_api"] = False
+    app.router.add_get("/fast", fast)
+    app.router.add_get("/slow", slow)
+    app.router.add_get("/slow_fail", slow_fail)
+    client = await aiohttp_client(app)
+
+    resp = await client.get("/fast")
+    assert resp.status == 201 and await resp.json() == {"x": 1}
+
+    resp = await client.get("/slow")
+    text = await resp.text()
+    assert resp.status == 200 and text.startswith("\n") and json.loads(text) == {"titles": [1]}
+    assert resp.headers["X-Accel-Buffering"] == "no"
+
+    resp = await client.get("/slow_fail")
+    body = json.loads(await resp.text())
+    assert resp.status == 200 and body["status"] == "error" and body["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
+
+    monkeypatch.setattr(routes, "HEARTBEAT_INTERVAL", 0.05)
+    monkeypatch.setattr(stats, "keys", {})
+
+    @heartbeat
+    async def slow_fail(request: web.Request) -> web.Response:
+        await asyncio.sleep(0.2)
+        raise APIError(APIErrorCode.SERVICE_ERROR, "late")
+
+    app = web.Application(middlewares=[stats_middleware])
+    app["debug_api"] = False
+    app.router.add_get("/slow_fail", slow_fail)
+    client = await aiohttp_client(app)
+
+    resp = await client.get("/slow_fail")
+    sent = len(await resp.read())
+    (entry,) = stats.keys.values()
+    assert resp.status == 200 and entry.rejected == 1 and entry.bytes_out >= sent > 0
+
+
+async def test_heartbeat_client_disconnect_is_quiet(
+    aiohttp_client, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
+) -> None:
+    """A client that leaves during the heartbeat logs no traceback, and the work runs to the end."""
+    monkeypatch.setattr(routes, "HEARTBEAT_INTERVAL", 0.02)
+    finished = asyncio.Event()
+
+    @heartbeat
+    async def slow(request: web.Request) -> web.Response:
+        await asyncio.sleep(0.3)
+        finished.set()
+        return web.json_response({})
+
+    app = web.Application()
+    app["debug_api"] = False
+    app.router.add_get("/slow", slow)
+    client = await aiohttp_client(app)
+
+    resp = await client.get("/slow")
+    await resp.content.read(1)
+    resp.close()
+    await asyncio.wait_for(finished.wait(), 2)
+    await asyncio.sleep(0.1)
+    assert "Error handling request" not in caplog.text
+
+
 async def test_cors_preflight_returns_headers(make_app, aiohttp_client) -> None:
     app = make_app(remote_only=True)
     client = await aiohttp_client(app)

tests/remote/unit/test_segment_filter.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import asyncio
 import json
 from types import SimpleNamespace
 
@@ -29,6 +30,7 @@ BASE = "https://cdn.test/hls/media.m3u8"
 
 def make_session(monkeypatch, track):
     session = SimpleNamespace(
+        lock=asyncio.Lock(),
         tracks={"t1": track},
         service_instance=SimpleNamespace(session=SimpleNamespace(get=lambda url: SimpleNamespace(text=PLAYLIST))),
     )

tests/remote/unit/test_server_cdm_absorb_routes.py

@@ -6,6 +6,7 @@ the first PSSH alone and so handed one track another track's narrower key set, a
 with no readable PSSH told the client nothing.
 """
 
+import asyncio
 import base64
 import json
 from types import SimpleNamespace
@@ -77,6 +78,7 @@ async def test_single_track_path_keeps_siblings(env, monkeypatch):
         drm=[_FakePlayReady(pssh_b64=PSSH, kid=KID_A), _FakePlayReady(pssh_b64=PSSH, kid=KID_B)],
     )
     session = SimpleNamespace(
+        lock=asyncio.Lock(),
         service_tag="EX",
         service_instance=SimpleNamespace(),
         served_keys={},
@@ -111,6 +113,7 @@ async def test_batch_cache_does_not_starve_second_track(env, monkeypatch):
     )
     tracks = {"vid": vid, "aud": aud}
     session = SimpleNamespace(
+        lock=asyncio.Lock(),
         service_tag="EX",
         service_instance=SimpleNamespace(),
         served_keys={},
@@ -142,6 +145,7 @@ async def test_no_pssh_tells_the_client_why(env, monkeypatch):
     buf = SessionLogBuffer()
     track = SimpleNamespace(id="vid", drm=[object()])
     session = SimpleNamespace(
+        lock=asyncio.Lock(),
         service_tag="EX",
         service_instance=SimpleNamespace(),
         served_keys={},

tests/remote/unit/test_server_cdm_manifest_kid.py

@@ -8,6 +8,7 @@ named different KIDs got one licence; and a KID the licence did not cover was lo
 server alone.
 """
 
+import asyncio
 import base64
 import json
 from types import SimpleNamespace
@@ -100,6 +101,7 @@ def fake_playready(monkeypatch):
 
 def _session(tracks: dict, buf=None):
     session = SimpleNamespace(
+        lock=asyncio.Lock(),
         service_tag="EX",
         service_instance=SimpleNamespace(),
         served_keys={},

tests/remote/unit/test_server_cdm_multikey.py

@@ -6,6 +6,7 @@ instead of short-circuiting on a vault hit for the PSSH KID alone, and the batch
 path must give each track the full bundle without borrowing another PSSH's keys.
 """
 
+import asyncio
 import base64
 import json
 from types import SimpleNamespace
@@ -118,6 +119,7 @@ async def test_batch_shares_full_bundle_per_pssh(monkeypatch):
 
     tracks = {tid: SimpleNamespace(id=tid, drm=[object()]) for tid in track_pssh}
     session = SimpleNamespace(
+        lock=asyncio.Lock(),
         service_tag="EXAMPLE",
         service_instance=SimpleNamespace(),
         served_keys={},

tests/remote/unit/test_server_cdm_type_override.py

@@ -8,6 +8,7 @@ client's choice decides, and the type the routes plan with must be the type of t
 they load.
 """
 
+import asyncio
 import json
 from types import SimpleNamespace
 from uuid import UUID
@@ -74,6 +75,7 @@ def _session(track):
         get_playready_license=lambda **k: b"",
     )
     return SimpleNamespace(
+        lock=asyncio.Lock(),
         service_tag="PCOK",
         service_instance=service,
         served_keys={},

tests/remote/unit/test_session_store.py

@@ -76,6 +76,17 @@ async def test_cleanup_expired_drops_old_authenticated(store: SessionStore, monk
     assert store.session_count == 0
 
 
+async def test_session_with_work_in_progress_does_not_expire(store: SessionStore) -> None:
+    """A request that holds the lock for longer than the TTL keeps its session and cache directory."""
+    from datetime import datetime, timedelta, timezone
+
+    entry = await store.create("EXAMPLE", _FakeService(), session_id="busy")
+    entry.last_accessed = datetime.now(timezone.utc) - timedelta(seconds=store.ttl + 100)
+    async with entry.lock:
+        assert await store.cleanup_expired() == 0
+        assert await store.get("busy") is entry
+
+
 async def test_cleanup_expired_keeps_pending_input_under_grace(store: SessionStore) -> None:
     """Sessions awaiting user input get a longer grace period (10 min) than authenticated TTL."""
     entry = await store.create("EXAMPLE", _FakeService())