fix(serve): apply staged service reloads when sessions end, with a max staging age

A staged tag was only retried on the refresh tick or a job completion, so a service with steady remote sessions never swapped in. Session end, expiry and eviction now trigger a single-flight apply, the tick applies after the pull, and serve.services_staged_max_age swaps an overdue tag in even while busy.

docs/reference/configuration/services.md

@@ -125,7 +125,8 @@ is the [REST API](../../dev/rest-api/index.md) section. These are the config key
 | `max_sessions` | int or null | `100` | Maximum concurrent remote sessions; at the cap the server evicts the oldest. Set it to `null` or `0` for no limit, which is also what `/api/dashboard/status` then reports, so a dashboard can tell "no cap" from a cap that happens to sit at the default. |
 | `history_limit` | int | `100` | How many finished jobs to retain in history. |
 | `compression_level` | int | `1` | gzip level for responses. |
-| `services_refresh_interval` | int (s) | `0` | How often the server pulls the git-backed service repositories in `directories.services` and hot-reloads the services that changed. `0` turns it off. A service with a running or queued job, or a live remote session, swaps to the new code as soon as its last job and session finish. |
+| `services_refresh_interval` | int (s) | `0` | How often the server pulls the git-backed service repositories in `directories.services` and hot-reloads the services that changed. `0` turns it off. A service with a running or queued job, or a live remote session, is staged and swaps to the new code as soon as its last job and session finish. |
+| `services_staged_max_age` | int (s) | `0` | How long a staged service waits for its jobs and sessions before the server swaps it in anyway. `0` waits forever. Set it on a server whose popular services never have a quiet moment. A session that is live at the swap keeps the old code until it ends; a running job keeps its old code too. |
 | `global_speed_limit` | str | *(unlimited)* | Server-wide download speed cap, e.g. `10M`, `1.5G` or plain bytes/sec (same format as `speed_limit`). One shared budget across all concurrent jobs; the server ignores per-job speed limits while it is set. |
 | `cdm_overrides` | list or bool | *(unset)* | Allowed per-request CDM overrides: a list of permitted device names, or `true` for any. Unset rejects every override. |
 | `allow_job_credentials` | bool | `false` | Permit clients to supply credentials per job. |

tests/core/test_services_reload.py

@@ -105,3 +105,19 @@ def test_apply_pending_never_reports_a_failed_reimport_as_applied(
     services.PENDING.add("FOO")
     assert services.apply_pending(busy=set()) == ["FOO"]
     assert services.LOADED_COMMITS["FOO"] == "new-head"
+
+
+def test_overdue_staged_service_applies_while_busy(service_dir: Path, monkeypatch: pytest.MonkeyPatch):
+    write_service(service_dir, 2)
+    monkeypatch.setattr(services, "repo_specs", lambda: ["example/repo"])
+    monkeypatch.setattr(services, "refresh_repo", lambda spec: (service_dir, ["~FOO"]))
+    monkeypatch.setattr(services, "staged_max_age", lambda: 60.0)
+
+    services.refresh_and_reload(busy={"FOO"})
+    assert services.apply_pending(busy={"FOO"}) == []
+    assert services.MODULES["FOO"].VERSION == 1
+
+    services.PENDING_SINCE["FOO"] -= 61
+    assert services.apply_pending(busy={"FOO"}) == ["FOO"]
+    assert services.MODULES["FOO"].VERSION == 2
+    assert services.PENDING == set() and "FOO" not in services.PENDING_SINCE

tests/remote/unit/test_serve_service_refresh.py

@@ -0,0 +1,57 @@
+"""Pins the serve refresh loop: staged reloads apply after the pull, not only on the next tick."""
+
+from __future__ import annotations
+
+import asyncio
+from typing import Any
+
+import pytest
+from aiohttp import web
+
+from unshackle.commands.serve import _install_service_refresh
+from unshackle.core import services
+from unshackle.core.api import download_manager
+from unshackle.core.config import config
+
+pytestmark = pytest.mark.unit
+
+
+async def test_refresh_tick_applies_tags_staged_during_the_pull(monkeypatch: pytest.MonkeyPatch) -> None:
+    """A tag whose last session ended while the pull ran is staged although idle; the same tick swaps it in."""
+    pulled = asyncio.Event()
+    seen: list[set[str]] = []
+
+    def fake_refresh(busy: set[str]) -> list[dict[str, Any]]:
+        # The pull read the busy set before the session ended, so the tag lands in PENDING while idle.
+        services.PENDING.add("EXAMPLE")
+        pulled.set()
+        return []
+
+    def fake_apply(busy: set[str]) -> list[str]:
+        seen.append(set(services.PENDING))
+        ready = sorted(services.PENDING - busy)
+        services.PENDING.difference_update(ready)
+        return ready
+
+    monkeypatch.setattr(config, "serve", {"services_refresh_interval": 1})
+    monkeypatch.setattr(services, "PENDING", set())
+    monkeypatch.setattr(services, "repo_specs", lambda: ["x"])
+    monkeypatch.setattr(services, "log_load_issues", lambda: None)
+    monkeypatch.setattr(services, "record_loaded_commits", lambda: None)
+    monkeypatch.setattr(services, "refresh_and_reload", fake_refresh)
+    monkeypatch.setattr(services, "apply_pending", fake_apply)
+    monkeypatch.setattr(download_manager, "busy_services", set)
+    monkeypatch.setattr(download_manager, "publish_service_event", lambda *a, **k: None)
+    monkeypatch.setattr(download_manager, "_reload_task", None)
+
+    app = web.Application()
+    _install_service_refresh(app)
+    await app.on_startup[-1](app)
+    try:
+        await asyncio.wait_for(pulled.wait(), timeout=3)
+        await asyncio.sleep(0.2)
+    finally:
+        await app.on_cleanup[-1](app)
+
+    assert {"EXAMPLE"} in seen
+    assert services.PENDING == set()

tests/remote/unit/test_session_store.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import asyncio
+import time
 
 import pytest
 
@@ -142,3 +143,123 @@ async def test_get_expiry_removes_cache_dir(store: SessionStore, monkeypatch: py
     entry.last_accessed = datetime.now(timezone.utc) - timedelta(seconds=store.ttl + 100)
     assert await store.get(entry.session_id) is None
     assert removed == ["_sessions/k/s/EXAMPLE"]
+
+
+async def test_session_end_applies_staged_service_reload(store: SessionStore, monkeypatch: pytest.MonkeyPatch) -> None:
+    """A staged reload swaps in as soon as the last session on that tag ends, not on the next refresh tick."""
+    from unshackle.core import services
+    from unshackle.core.api import download_manager
+
+    applied: list[set[str]] = []
+
+    def fake_apply(busy: set[str]) -> list[str]:
+        applied.append(busy)
+        return sorted(services.PENDING - busy)
+
+    monkeypatch.setattr(services, "PENDING", {"EXAMPLE"})
+    monkeypatch.setattr(services, "apply_pending", fake_apply)
+    monkeypatch.setattr(download_manager, "publish_service_event", lambda *a, **k: None)
+
+    entry = await store.create("EXAMPLE", _FakeService())
+    assert await store.delete(entry.session_id)
+    await asyncio.sleep(0.05)
+    assert applied == [set()]
+
+
+def _stage_reload(monkeypatch: pytest.MonkeyPatch, pending: set[str], delay: float = 0.0) -> list[set[str]]:
+    """Stage ``pending`` for reload and record every apply_pending call; returns the call log."""
+    from unshackle.core import services
+    from unshackle.core.api import download_manager
+
+    applied: list[set[str]] = []
+
+    def fake_apply(busy: set[str]) -> list[str]:
+        if delay:
+            time.sleep(delay)
+        applied.append(busy)
+        return sorted(services.PENDING - busy)
+
+    monkeypatch.setattr(services, "PENDING", pending)
+    monkeypatch.setattr(services, "apply_pending", fake_apply)
+    monkeypatch.setattr(download_manager, "publish_service_event", lambda *a, **k: None)
+    monkeypatch.setattr(download_manager, "_reload_task", None)
+    return applied
+
+
+async def test_session_end_on_unstaged_tag_skips_reload(store: SessionStore, monkeypatch: pytest.MonkeyPatch) -> None:
+    applied = _stage_reload(monkeypatch, {"OTHER"})
+
+    entry = await store.create("EXAMPLE", _FakeService())
+    assert await store.delete(entry.session_id)
+    await asyncio.sleep(0.05)
+    assert applied == []
+
+
+async def test_concurrent_session_ends_reload_once(store: SessionStore, monkeypatch: pytest.MonkeyPatch) -> None:
+    """Session ends that overlap an in-flight reload share it instead of each re-importing the tag."""
+    applied = _stage_reload(monkeypatch, {"EXAMPLE"}, delay=0.05)
+
+    ids = [(await store.create("EXAMPLE", _FakeService())).session_id for _ in range(10)]
+    await asyncio.gather(*(store.delete(sid) for sid in ids))
+    await asyncio.sleep(0.3)
+    assert len(applied) == 1
+
+
+async def test_eviction_applies_staged_service_reload(store: SessionStore, monkeypatch: pytest.MonkeyPatch) -> None:
+    applied = _stage_reload(monkeypatch, {"EXAMPLE"})
+    monkeypatch.setattr(type(store), "max_sessions", property(lambda _: 1))
+
+    await store.create("EXAMPLE", _FakeService(), session_id="a")
+    await store.create("OTHER", _FakeService(), session_id="b")
+    await asyncio.sleep(0.05)
+    assert len(applied) == 1
+
+
+async def test_inline_expiry_applies_staged_service_reload(
+    store: SessionStore, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    from datetime import datetime, timedelta, timezone
+
+    applied = _stage_reload(monkeypatch, {"EXAMPLE"})
+
+    entry = await store.create("EXAMPLE", _FakeService())
+    entry.last_accessed = datetime.now(timezone.utc) - timedelta(seconds=store.ttl + 100)
+    assert await store.get(entry.session_id) is None
+    await asyncio.sleep(0.05)
+    assert applied == [set()]
+
+
+async def test_cleanup_expired_applies_staged_service_reload_once(
+    store: SessionStore, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """A cleanup sweep that drops several sessions on one tag schedules one reload for that tag."""
+    from datetime import datetime, timedelta, timezone
+
+    applied = _stage_reload(monkeypatch, {"EXAMPLE"})
+    stale = datetime.now(timezone.utc) - timedelta(seconds=store.ttl + 100)
+    for sid in ("a", "b"):
+        (await store.create("EXAMPLE", _FakeService(), session_id=sid)).last_accessed = stale
+
+    assert await store.cleanup_expired() == 2
+    await asyncio.sleep(0.05)
+    assert applied == [set()]
+
+
+async def test_live_session_on_same_tag_keeps_it_busy(store: SessionStore, monkeypatch: pytest.MonkeyPatch) -> None:
+    """Ending one of two sessions on a staged tag reports the tag busy; ending the last one frees it."""
+    from unshackle.core.api import download_manager
+    from unshackle.core.api import session_store as session_store_module
+
+    applied = _stage_reload(monkeypatch, {"EXAMPLE"})
+    monkeypatch.setattr(session_store_module, "session_store", store)
+    monkeypatch.setattr(download_manager, "download_manager", None)
+
+    first = await store.create("EXAMPLE", _FakeService())
+    second = await store.create("EXAMPLE", _FakeService())
+    assert await store.delete(first.session_id)
+    await asyncio.sleep(0.05)
+    assert applied == [{"EXAMPLE"}]
+
+    assert await store.delete(second.session_id)
+    await asyncio.sleep(0.05)
+    assert applied == [{"EXAMPLE"}, set()]

unshackle/commands/serve.py

@@ -52,7 +52,7 @@ def _install_service_refresh(app: web.Application) -> None:
         return
 
     async def loop() -> None:
-        from unshackle.core.api.download_manager import busy_services
+        from unshackle.core.api.download_manager import busy_services, schedule_pending_reload
 
         while True:
             await asyncio.sleep(interval)
@@ -70,6 +70,9 @@ def _install_service_refresh(app: web.Application) -> None:
                     for err in r["load_errors"]:
                         log.error(f"Service reload failed: {err}")
                 publish_refresh_events(repos)
+                # The busy set was read before the pull. A tag whose last session ended during the pull
+                # is staged although idle, and its session-end trigger ran before it was staged.
+                schedule_pending_reload()
             except Exception:
                 log.exception("Service refresh failed")
 

unshackle/core/api/download_manager.py

@@ -747,21 +747,7 @@ class DownloadQueueManager:
 
     def schedule_pending_reload(self) -> None:
         """Swap in staged service updates whose last job has finished, off the event loop."""
-        from unshackle.core import services
-
-        if not services.PENDING:
-            return
-
-        async def run() -> None:
-            try:
-                applied = await asyncio.to_thread(services.apply_pending, busy_services())
-                if applied:
-                    log.info(f"Services reloaded after job completion: {', '.join(applied)}")
-                    publish_service_event("applied", applied)
-            except Exception:
-                log.exception("Applying pending service reloads failed")
-
-        asyncio.get_running_loop().create_task(run())
+        schedule_pending_reload()
 
     @staticmethod
     def put(queue: asyncio.Queue, item: Optional[Dict[str, Any]]) -> None:
@@ -1210,3 +1196,41 @@ def busy_services() -> set[str]:
     if stats.mode != "remote_only":
         busy |= get_download_manager().busy_services()
     return busy
+
+
+_reload_task: Optional[asyncio.Task[None]] = None
+
+
+def schedule_pending_reload(tag: Optional[str] = None) -> None:
+    """Swap in staged service updates whose last job or remote session has ended, off the event loop.
+
+    Runs on every job completion and remote session removal, because a busy service on a popular
+    tag can otherwise stay staged across every periodic refresh tick. With ``tag``, the reload only
+    runs when ``PENDING`` holds that tag.
+    """
+    global _reload_task
+    from unshackle.core import services
+
+    if not services.PENDING or (tag is not None and tag not in services.PENDING):
+        return
+    # apply_pending computes its ready set before it takes RELOAD_LOCK, so concurrent
+    # session ends would each re-import the same tag; one in-flight task covers them all.
+    if _reload_task is not None and not _reload_task.done():
+        return
+
+    async def run() -> None:
+        try:
+            applied = await asyncio.to_thread(services.apply_pending, busy_services())
+            if applied:
+                log.info(f"Services reloaded after their jobs and sessions ended: {', '.join(applied)}")
+                publish_service_event("applied", applied)
+        except Exception:
+            log.exception("Applying pending service reloads failed")
+
+    _reload_task = asyncio.get_running_loop().create_task(run())
+
+    def clear(_: asyncio.Task[None]) -> None:
+        global _reload_task
+        _reload_task = None
+
+    _reload_task.add_done_callback(clear)

unshackle/core/api/session_store.py

@@ -95,6 +95,13 @@ class SessionEntry:
         }
 
 
+def _schedule_pending_reload(service_tag: str) -> None:
+    """Apply staged service reloads now that a remote session ended; the task runs after the store lock releases."""
+    from unshackle.core.api.download_manager import schedule_pending_reload
+
+    schedule_pending_reload(service_tag)
+
+
 def _publish(action: str, entry: SessionEntry, reason: Optional[str] = None) -> None:
     data = {"action": action, **entry.summary()}
     if reason:
@@ -145,6 +152,7 @@ class SessionStore:
                     evicted.input_bridge.cancel()
                 self.cleanup_cache_dir(evicted.cache_tag)
                 _publish("delete", evicted, "evicted")
+                _schedule_pending_reload(evicted.service_tag)
 
             session_id = session_id or str(uuid.uuid4())
             entry = SessionEntry(
@@ -173,6 +181,7 @@ class SessionStore:
                     log.info(f"Session {sanitize_log(session_id)} expired (elapsed={elapsed:.0f}s, ttl={self.ttl}s)")
                     self.cleanup_cache_dir(entry.cache_tag)
                     _publish("delete", self._sessions.pop(session_id), "expired")
+                    _schedule_pending_reload(entry.service_tag)
                     return None
 
             entry.touch()
@@ -196,6 +205,7 @@ class SessionStore:
                 self.cleanup_cache_dir(entry.cache_tag)
                 _publish("delete", entry, "closed")
                 log.info(f"Deleted session {sanitize_log(session_id)}")
+                _schedule_pending_reload(entry.service_tag)
                 return True
             return False
 
@@ -211,12 +221,16 @@ class SessionStore:
                         expired.append(sid)
                 elif elapsed > self.ttl:
                     expired.append(sid)
+            ended_tags: set[str] = set()
             for sid in expired:
                 entry = self._sessions.pop(sid)
                 if entry.input_bridge:
                     entry.input_bridge.cancel()
                 self.cleanup_cache_dir(entry.cache_tag)
                 _publish("delete", entry, "expired")
+                ended_tags.add(entry.service_tag)
+            for tag in ended_tags:
+                _schedule_pending_reload(tag)
             if expired:
                 log.info(f"Cleaned up {len(expired)} expired sessions")
             return len(expired)

unshackle/core/services.py

@@ -272,13 +272,24 @@ def failed_tags(errors: Iterable[str]) -> set[str]:
     return {err.split(":", 1)[0] for err in errors if ":" in err}
 
 
+def staged_max_age() -> float:
+    """Seconds a staged tag may wait for its jobs and sessions before it swaps in anyway; 0 waits forever."""
+    return float((config.serve or {}).get("services_staged_max_age", 0) or 0)
+
+
 def apply_pending(busy: set[str] | None = None) -> list[str]:
     """Reload every staged tag whose service is no longer busy; returns the tags now running.
 
-    A tag whose re-import failed is not one of them: the old module keeps serving, so calling
-    it applied would tell the dashboard that code is live when it is not.
+    A tag staged for longer than ``serve.services_staged_max_age`` reloads even while busy:
+    its live sessions keep the old class objects they already hold, and a popular tag would
+    otherwise never see a quiet moment. A tag whose re-import failed is not counted as
+    running: the old module keeps serving, so calling it applied would tell the dashboard
+    that code is live when it is not.
     """
-    ready = sorted(PENDING - (busy or set()))
+    max_age = staged_max_age()
+    now = time.time()
+    overdue = {tag for tag, since in PENDING_SINCE.items() if max_age and now - since >= max_age}
+    ready = sorted((PENDING - (busy or set())) | (PENDING & overdue))
     if not ready:
         return []
     errors = reload_services(ready)