fix(watermark): review follow-ups for #1577 — CI red + bot findings

Two CI failures, both understood:

1. test_shutdown_preload_race_1000 pins the production _cancel_and_await
  _tasks call site by regex; the new fifth handle broke the pattern. The
   guard now pins all FIVE handles (its property — every preload handle
   awaited under one generous bound — is unchanged).

2. test_prefetched_model_gets_one_extra_idle_window flaked only in the
   full suite: many tests boot the app lifespan, and any that exits
   without a lifespan shutdown leaves the deferred watermark-preload
   task pending — 35s later it fires mid-suite in another thread and
   re-stamps _last_used under whatever test is running. conftest now
   defaults OMNIVOICE_PRELOAD_WATERMARK=0 for the test session (a test
   can still opt in), and the grace test neutralizes will_mark so a
   leaked warm-up can't touch it.

Bot findings: Greptile P1 + CodeRabbit — cancelling the preload task
doesn't stop a watermark-pool thread already inside the ~42s cold
import, and nothing drained that pool at shutdown (only the GPU pool
was reset). Shutdown now drains the watermark pool's queue
(shutdown(wait=False, cancel_futures=True)) — bounded abandon, same
documented reality that Python can't kill a running thread. CodeRabbit
Major: the warm-up reads its own delay knob
(OMNIVOICE_PRELOAD_WATERMARK_DELAY, default 35s) instead of reusing the
capture-ASR delay, so a capture env override no longer retimes it.
CodeRabbit Minor: the _prefetched_unused claim/clear transitions now
happen under _generator_lock, so the retention grace can't be granted
to a model that has actually been used; the test fixture resets all
lifecycle globals.

Skipped with reason: gating prefetch on local-checkpoint presence — the
warm-up downloads only what the first embed would download anyway;
time-shifting that download is the feature, not a new network call.

backend/main.py

@@ -383,6 +383,19 @@ def _capture_preload_delay_s() -> float:
         pass
     return 30.0
 
+def _watermark_preload_delay_s() -> float:
+    """Seconds after boot before the AudioSeal generator warm-up fires.
+
+    Own knob, NOT ``_capture_preload_delay_s`` + offset: a capture-specific
+    env override must not retime the watermark warm too, and the two cold
+    imports shouldn't fire on the same tick (CodeRabbit, PR #1577). Default
+    35s sits ~5s past the capture-ASR warm for the same reason."""
+    raw = os.environ.get("OMNIVOICE_PRELOAD_WATERMARK_DELAY", "")
+    try:
+        return float(raw) if raw.strip() else 35.0
+    except ValueError:
+        return 35.0
+
 
 def _capture_preload_ram_ok(min_free_bytes: int = 4 * 1024**3) -> bool:
     """RAM guard for the dictation warm-up: skip below 4 GB free so the
@@ -915,11 +928,7 @@ async def _phase_b(app: FastAPI) -> None:
     # the shared default executor.
     if _env_flag("OMNIVOICE_PRELOAD_WATERMARK", default=True):
         async def _preload_watermark():
-            # Own delay (+5s past the capture warm), not _capture_preload_delay_s:
-            # a capture-specific env override shouldn't retime this, and both
-            # warms firing on the same tick is exactly the I/O overlap the
-            # defer exists to avoid.
-            await asyncio.sleep(_capture_preload_delay_s() + 5.0)
+            await asyncio.sleep(_watermark_preload_delay_s())
             loop = asyncio.get_running_loop()
             from services import watermark as _watermark
 
@@ -1123,6 +1132,17 @@ async def lifespan(app: FastAPI):
         getattr(app.state, "watermark_preload_task", None),
         timeout=20.0,
     )
+    # The watermark warm-up runs on its dedicated 1-worker pool; cancelling
+    # the task above detaches the asyncio side but a thread already inside
+    # the ~42s cold import keeps running (Python can't kill it — same
+    # reality as the GPU-pool note above). Drain the pool's QUEUE so nothing
+    # new starts, mirroring _reset_gpu_pool's bounded-abandon approach.
+    try:
+        from services.model_manager import get_watermark_pool as _get_wm_pool
+
+        _get_wm_pool().shutdown(wait=False, cancel_futures=True)
+    except Exception:
+        pass
     # Unload the model and free GPU memory
     try:
         import services.model_manager as mm

backend/services/watermark.py

@@ -134,7 +134,11 @@ def prefetch_generator() -> None:
             return
         global _prefetched_unused
         _get_generator()
-        _prefetched_unused = True
+        with _generator_lock:
+            # Under the lock so the "prefetched" claim can't land after an
+            # embed/detect already cleared it (CodeRabbit, PR #1577): that
+            # would grant the retention grace to a model that HAS been used.
+            _prefetched_unused = True
         logger.info("AudioSeal generator prefetched in the background")
     except Exception:
         logger.warning(
@@ -286,7 +290,8 @@ def embed_watermark(
 
     try:
         global _prefetched_unused
-        _prefetched_unused = False
+        with _generator_lock:
+            _prefetched_unused = False
         generator = _get_generator()
         msg = torch.tensor(message or OMNI_MESSAGE, dtype=torch.int32).unsqueeze(0)
 
@@ -352,7 +357,8 @@ def detect_watermark(
 
     try:
         global _prefetched_unused
-        _prefetched_unused = False
+        with _generator_lock:
+            _prefetched_unused = False
         detector = _get_detector()
 
         # Normalise shape to (batch, channels, samples)

tests/conftest.py

@@ -45,6 +45,14 @@ if not os.environ.get("OMNIVOICE_ENV_FILE"):
 # need a different value monkeypatch it explicitly.
 os.environ["OMNIVOICE_MODEL"] = "test"
 
+# Background warm-ups must not fire mid-suite: many tests boot the app
+# lifespan via TestClient, and any that exits without a lifespan shutdown
+# leaves the deferred preload task pending — 35s later (mid-suite, in
+# another thread) it loads real models and mutates watermark module state
+# under whatever test happens to be running (seen as a CI-only flake in the
+# prefetch cold-start tests). setdefault so a test can still opt in.
+os.environ.setdefault("OMNIVOICE_PRELOAD_WATERMARK", "0")
+
 
 # ── Test fixtures ──────────────────────────────────────────────────────────
 

tests/test_shutdown_preload_race_1000.py

@@ -140,6 +140,7 @@ def test_production_shutdown_wait_is_generous_enough_for_a_cold_import():
         r"getattr\(app\.state, \"worker_task\", None\),\s*"
         r"getattr\(app\.state, \"preload_task\", None\),\s*"
         r"getattr\(app\.state, \"capture_preload_task\", None\),\s*"
+        r"getattr\(app\.state, \"watermark_preload_task\", None\),\s*"
         r"timeout=([\d.]+),?\s*\)",
         src,
     )

tests/test_watermark_prefetch_coldstart.py

@@ -21,12 +21,19 @@ from services import watermark
 
 @pytest.fixture(autouse=True)
 def _reset_models(monkeypatch):
+    # Reset ALL lifecycle globals (CodeRabbit, PR #1577): a stale warm-up
+    # stamp or availability cache from a prior test changes this test's
+    # conditions.
     watermark._generator = None
     watermark._detector = None
-    monkeypatch.delenv("OMNIVOICE_WM", raising=False)
+    watermark._last_used = 0.0
+    watermark._prefetched_unused = False
+    monkeypatch.setattr(watermark, "_audioseal_available", None, raising=False)
     yield
     watermark._generator = None
     watermark._detector = None
+    watermark._last_used = 0.0
+    watermark._prefetched_unused = False
 
 
 def _fake_audioseal(monkeypatch, load_s: float) -> list[int]:
@@ -152,6 +159,9 @@ def test_prefetched_model_gets_one_extra_idle_window(monkeypatch):
     generator at the first idle tick, re-imposing the cold start the prefetch
     exists to hide. It now survives ONE extra window; real use clears the
     grace entirely."""
+    # Immune to a leaked background prefetch (the CI flake): if a warm-up
+    # from an earlier app-boot fires mid-test it would re-stamp _last_used.
+    monkeypatch.setattr(watermark, "will_mark", lambda: False)
     watermark._generator = SimpleNamespace(eval=lambda: None)
     watermark._prefetched_unused = True
     watermark._last_used = 0.0  # long idle