"""Cold-hold ADMISSION CAP + abandon-on-disconnect — the blast-radius bound.

Operator incident 2026-07-27: "what is happening so that the site is getting
overwhelmed by being called by a single script?" Central serves everything from
`gunicorn --workers 3 --threads 8` = 24 request slots, and a held cold call pins
one of them for up to HUGPY_COLD_HOLD_MAX_S (1500s on the live unit). Two dozen
cold-or-doomed requests therefore park every slot and /health, /llm/workers and
the console all stop answering — the site dies wholesale while nothing is broken.

This covers the two fixes, both central-side:

  1. AN ADMISSION CAP on concurrent cold holds. Beyond N simultaneous holds a new
     arrival gets a FAST, honest refusal instead of a slot; a released hold frees
     its admission for the next arrival; a genuine slow load with a patient client
     still succeeds; WARM traffic is never refused by the cap; and — the
     acceptance test — a lightweight endpoint still answers promptly while every
     admission is taken.
  2. ABANDON-ON-DISCONNECT. The socket probe reads a hung-up peer (and ONLY a
     hung-up peer) as gone, and the WSGI->loop bridge cancels the in-flight work
     so the request slot comes back instead of sitting out the hold ceiling.

Runs under pytest AND as a plain script:
    venv/bin/python -m pytest tests/test_cold_hold_cap.py -q
    venv/bin/python tests/test_cold_hold_cap.py
"""
import asyncio
import importlib
import os
import socket
import sys
import threading
import time
import types
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

# In-process only — no cross-process comms DB side effects during tests.
os.environ.setdefault("HUGPY_COMMS_DB", "off")

ok = 0


def check(name, cond):
    global ok
    assert cond, name
    ok += 1
    print(f"  ok - {name}")


WORKER = {"id": "w1", "name": "computron", "url": "http://w1:9100"}


def _req(rid="rid-1"):
    return types.SimpleNamespace(
        request_id=rid, pool=None,
        reference_images=None, reference_images_b64=None,
        model_dump=lambda: {"messages": [{"role": "user", "content": "hi"}]},
    )


def _text_framework(remote):
    for (fw, tk) in remote.FRAMEWORK_RUNNERS:
        if tk != "image-text-to-text":
            return fw, tk
    return next(iter(remote.FRAMEWORK_RUNNERS))


async def _collect(agen):
    return [ev async for ev in agen]


def _etypes(evs):
    return [getattr(e, "type", None) for e in evs]


# ---------------------------------------------------------------------------
def _knob_checks(remote):
    """The cap is configurable and its default is sane."""
    for k in ("HUGPY_COLD_HOLD_MAX_CONCURRENT", "HUGPY_COLD_HOLD_RETRY_AFTER_S"):
        os.environ.pop(k, None)
    check("default concurrent-hold cap is 4 (half of one gunicorn process's 8 "
          "threads; 3 x 4 = 12 of the site's 24 slots)",
          remote._cold_hold_max_concurrent() == 4)
    check("default Retry-After is 20s", remote._cold_hold_retry_after_s() == 20)
    try:
        os.environ["HUGPY_COLD_HOLD_MAX_CONCURRENT"] = "9"
        check("cap is an env knob (HUGPY_COLD_HOLD_MAX_CONCURRENT)",
              remote._cold_hold_max_concurrent() == 9)
        os.environ["HUGPY_COLD_HOLD_MAX_CONCURRENT"] = "0"
        check("a non-positive cap falls back to the default (never 'no holds')",
              remote._cold_hold_max_concurrent() == 4)
        os.environ["HUGPY_COLD_HOLD_MAX_CONCURRENT"] = "banana"
        check("garbage falls back to the default (a knob can misconfigure, "
              "never break)", remote._cold_hold_max_concurrent() == 4)
        os.environ["HUGPY_COLD_HOLD_RETRY_AFTER_S"] = "45"
        check("Retry-After is an env knob", remote._cold_hold_retry_after_s() == 45)
    finally:
        for k in ("HUGPY_COLD_HOLD_MAX_CONCURRENT", "HUGPY_COLD_HOLD_RETRY_AFTER_S"):
            os.environ.pop(k, None)

    # Accounting primitives.
    check("hold counter starts at zero", remote._hold_count() == 0)
    permits = [remote._hold_try_acquire() for _ in range(5)]
    check("exactly `cap` permits are handed out, never more",
          sum(p is not None for p in permits) == 4)
    check("the 5th arrival gets NO permit", permits[4] is None)
    check("the counter reports the truth while saturated", remote._hold_count() == 4)
    permits[0].release()
    check("a released hold decrements the counter", remote._hold_count() == 3)
    permits[0].release()
    check("release() is idempotent (a double release can't free a phantom slot)",
          remote._hold_count() == 3)
    again = remote._hold_try_acquire()
    check("a released hold frees its admission for the next arrival",
          again is not None)
    for p in permits[1:4] + [again]:
        if p is not None:
            p.release()
    check("all permits returned", remote._hold_count() == 0)


# ---------------------------------------------------------------------------
def _message_checks(remote):
    """The refusal is honest and actionable — and misattributes NOTHING."""
    err = remote.ColdHoldCapacityError("Qwen3-Coder-Next-GGUF", WORKER, 4, 4,
                                       loading=True)
    msg = err.stream_message()
    low = msg.lower()
    check("refusal carries a machine code routes can map to 503",
          err.code == "cold_load_capacity" and err.code in msg)
    check("refusal names the concurrent-load LIMIT as the reason",
          "concurrent model loads" in low and "limit" in low)
    check("refusal says the model is still LOADING (not that it failed)",
          "is still loading on" in low)
    check("refusal states plainly that nothing is broken",
          "nothing is broken" in low)
    check("refusal is actionable (tells the caller when to come back)",
          "retry in about 20s" in low)
    # The recorded doctrine: a confidently wrong error is worse than a vague one.
    for lie in ("too large", "too big", "won't fit", "out of memory", "crash",
                "unhealthy", "offline", "unreachable", "failed to load",
                "not supported", "unknown model"):
        check(f"refusal does NOT misattribute ({lie!r})", lie not in low)
    check("refusal is not mistaken for a permanent LOAD failure",
          remote._is_permanent_load_error(err) is False)
    check("refusal is not mistaken for a request-shape failure",
          remote._is_request_shape_error(err) is False)

    cold = remote.ColdHoldCapacityError("m", WORKER, 4, 4, loading=False)
    check("when the worker isn't loading it yet, the message says so honestly",
          "is not loaded yet on" in cold.stream_message())
    body = err.as_error()["error"]
    check("structured error carries code/limit/retry_after for a client",
          body["code"] == "cold_load_capacity" and body["limit"] == 4
          and body["retry_after_s"] == 20 and body["holds_in_flight"] == 4)


# ---------------------------------------------------------------------------
def _admission_checks(remote):
    """_admit_cold_hold: refuse a cold call when full, never a warm one."""
    orig_ls = remote._load_state_provider
    os.environ["HUGPY_COLD_HOLD_MAX_CONCURRENT"] = "2"
    try:
        remote.set_load_state_provider(
            lambda mk, wid, since=0.0: {"healthy": False, "in_progress": True})
        held = [remote._admit_cold_hold("m", WORKER, 0.0) for _ in range(2)]
        check("cold calls take permits up to the cap",
              all(h is not None for h in held) and remote._hold_count() == 2)
        raised = None
        t0 = time.monotonic()
        try:
            remote._admit_cold_hold("m", WORKER, 0.0)
        except remote.ColdHoldCapacityError as exc:
            raised = exc
        dt = time.monotonic() - t0
        check("the N+1th cold arrival is REFUSED, not queued", raised is not None)
        # Thresholds here are deliberately generous relative to what they rule
        # out (a hold parks a slot for HUGPY_COLD_HOLD_MAX_S — 1500s live), so a
        # loaded CI box can't flake them while they still prove "immediate".
        check("the refusal is FAST (no bounded wait, no slot taken)", dt < 0.5)
        check("the refusal reports the real in-flight count", raised.held == 2)

        # WARM traffic must sail through a saturated cap — the cap bounds LOADS.
        remote.set_load_state_provider(
            lambda mk, wid, since=0.0: {"healthy": True, "in_progress": False})
        check("a WARM model is admitted (uncounted) even with the cap full",
              remote._admit_cold_hold("m", WORKER, 0.0) is None)

        # Never refuse on ignorance.
        remote.set_load_state_provider(None)
        check("with no load-state provider we cannot tell cold from warm, so we "
              "ADMIT (never refuse on ignorance)",
              remote._admit_cold_hold("m", WORKER, 0.0) is None)
        remote.set_load_state_provider(lambda *a, **k: (_ for _ in ()).throw(RuntimeError("boom")))
        check("a RAISING load-state provider degrades to admit, never to refuse",
              remote._admit_cold_hold("m", WORKER, 0.0) is None)

        for h in held:
            h.release()
        check("permits all returned", remote._hold_count() == 0)
    finally:
        remote.set_load_state_provider(orig_ls)
        os.environ.pop("HUGPY_COLD_HOLD_MAX_CONCURRENT", None)


# ---------------------------------------------------------------------------
class _Env:
    """The hold knobs a cap test needs, restored on exit."""

    KEYS = ("HUGPY_CENTRAL_GATE", "HUGPY_COLD_HOLD_POLL_S", "HUGPY_COLD_HOLD_MAX_S",
            "HUGPY_COLD_HOLD_STALL_S", "HUGPY_COLD_HOLD_MAX_CONCURRENT",
            "HUGPY_LOCAL_FALLBACK", "HUGPY_NO_LOCAL_SERVING")

    def __init__(self, **kw):
        self.kw = kw
        self.saved = {}

    def __enter__(self):
        for k in self.KEYS:
            self.saved[k] = os.environ.get(k)
            os.environ.pop(k, None)
        os.environ.update(self.kw)
        return self

    def __exit__(self, *exc):
        for k in self.KEYS:
            os.environ.pop(k, None)
            if self.saved[k] is not None:
                os.environ[k] = self.saved[k]
        return False


def _runner_checks(remote):
    """run(): N admitted, N+1 refused fast, a slow load still succeeds."""
    fw, tk = _text_framework(remote)
    Runner = remote.make_delegating_runner(fw, tk)
    runner = Runner(types.SimpleNamespace(model_key="cold-model"))

    orig_select, orig_run = remote._select, remote._worker_run_once
    orig_ls = remote._load_state_provider
    remote._select = lambda mk, pool=None, task=None, **kw: (dict(WORKER), None)
    remote.set_load_state_provider(
        lambda mk, wid, since=0.0: {"healthy": False, "in_progress": True})

    try:
        with _Env(HUGPY_CENTRAL_GATE="off", HUGPY_COLD_HOLD_POLL_S="0.02",
                  HUGPY_COLD_HOLD_MAX_S="1.0", HUGPY_COLD_HOLD_STALL_S="1.0",
                  HUGPY_COLD_HOLD_MAX_CONCURRENT="3"):

            live = {"n": 0, "max": 0}
            lock = threading.Lock()

            async def run_forever_cold(worker, payload, result_type, request_id,
                                       model_key):
                with lock:
                    live["n"] += 1
                    live["max"] = max(live["max"], live["n"])
                try:
                    await asyncio.sleep(0.05)
                    raise RuntimeError("RemoteProtocolError: Server disconnected")
                finally:
                    with lock:
                        live["n"] -= 1

            remote._worker_run_once = run_forever_cold

            async def _six():
                # Six simultaneous cold calls, cap 3.
                async def one(i):
                    t0 = time.monotonic()
                    try:
                        await runner.run(_req(f"c{i}"))
                        return ("ok", time.monotonic() - t0)
                    except remote.ColdHoldCapacityError as exc:
                        return ("refused", time.monotonic() - t0, exc)
                    except RuntimeError as exc:
                        return ("held", time.monotonic() - t0, str(exc))
                return await asyncio.gather(*[one(i) for i in range(6)])

            out = asyncio.run(_six())
            refused = [r for r in out if r[0] == "refused"]
            held = [r for r in out if r[0] == "held"]
            check("exactly `cap` concurrent cold calls are admitted (3)",
                  len(held) == 3)
            check("every arrival beyond the cap is refused (3)", len(refused) == 3)
            check("the refusals are FAST — not a parked slot",
                  all(r[1] < 0.5 for r in refused))
            check("a refusal returns far sooner than a hold does (the whole point)",
                  max(r[1] for r in refused) * 2 < min(h[1] for h in held))
            check("the admitted holds really did hold (they outlast the refusals)",
                  all(h[1] > 0.5 for h in held))
            check("refusals carry the honest capacity message",
                  all("cold_load_capacity" in r[2].stream_message() for r in refused))
            check("no relay was ever fired for a refused call (no worker load "
                  "kicked, no slot consumed)", live["max"] <= 3)
            check("all admissions were returned when the holds ended",
                  remote._hold_count() == 0)

            # A released hold frees its admission for the NEXT arrival.
            after = asyncio.run(_six())
            check("after the first wave drained, a fresh wave is admitted again "
                  "(a released hold frees its slot)",
                  len([r for r in after if r[0] == "held"]) == 3)
            check("hold counter back to zero between waves",
                  remote._hold_count() == 0)

        # -- a GENUINE slow load with a patient client still succeeds ---------
        with _Env(HUGPY_CENTRAL_GATE="off", HUGPY_COLD_HOLD_POLL_S="0.02",
                  HUGPY_COLD_HOLD_MAX_S="10", HUGPY_COLD_HOLD_STALL_S="10",
                  HUGPY_COLD_HOLD_MAX_CONCURRENT="3"):
            calls = {"n": 0}

            async def slow_then_ok(worker, payload, result_type, request_id, model_key):
                calls["n"] += 1
                if calls["n"] <= 5:
                    await asyncio.sleep(0.03)
                    raise RuntimeError("RemoteProtocolError: Server disconnected")
                return {"ok": True, "text": "done", "request_id": request_id,
                        "model_key": model_key}

            remote._worker_run_once = slow_then_ok
            res = asyncio.run(runner.run(_req("patient")))
            check("a genuine slow load with a patient client STILL succeeds "
                  "under the cap", res.get("ok") is True and calls["n"] == 6)
            check("its admission was returned on success", remote._hold_count() == 0)

            # …and it never consumed an admission for the whole generation:
            # a warm one-shot with the cap already full is served uncapped.
            remote.set_load_state_provider(
                lambda mk, wid, since=0.0: {"healthy": True, "in_progress": False})
            blockers = [remote._hold_try_acquire() for _ in range(3)]
            calls["n"] = 6

            async def warm_ok(worker, payload, result_type, request_id, model_key):
                return {"ok": True, "text": "warm", "request_id": request_id,
                        "model_key": model_key}

            remote._worker_run_once = warm_ok
            res = asyncio.run(runner.run(_req("warm")))
            check("WARM traffic is never refused by the cold-hold cap",
                  res.get("ok") is True)
            for b in blockers:
                if b is not None:
                    b.release()
    finally:
        remote._select, remote._worker_run_once = orig_select, orig_run
        remote.set_load_state_provider(orig_ls)


# ---------------------------------------------------------------------------
def _stream_checks(remote):
    """stream(): refusal is an honest ErrorEvent; a warm stream frees its slot."""
    fw, tk = _text_framework(remote)
    Runner = remote.make_delegating_runner(fw, tk)
    runner = Runner(types.SimpleNamespace(model_key="cold-model"))
    TokenEvent = remote.TokenEvent
    DoneEvent = remote.DoneEvent

    orig_select, orig_ws = remote._select, remote._worker_stream
    orig_ls = remote._load_state_provider
    remote._select = lambda mk, pool=None, task=None, **kw: (dict(WORKER), None)
    remote.set_load_state_provider(
        lambda mk, wid, since=0.0: {"healthy": False, "in_progress": True})
    try:
        with _Env(HUGPY_CENTRAL_GATE="off", HUGPY_COLD_HOLD_POLL_S="0.02",
                  HUGPY_COLD_HOLD_MAX_S="1.0", HUGPY_COLD_HOLD_STALL_S="1.0",
                  HUGPY_COLD_HOLD_MAX_CONCURRENT="1"):

            async def ws_cold(worker, payload, rid):
                await asyncio.sleep(0.05)
                raise RuntimeError("server disconnected (swapping)")
                yield  # pragma: no cover — marks this a generator

            remote._worker_stream = ws_cold

            async def _two():
                return await asyncio.gather(_collect(runner.stream(_req("s-a"))),
                                            _collect(runner.stream(_req("s-b"))))

            a, b = asyncio.run(_two())
            msgs = [getattr(e, "message", "") for e in a + b
                    if getattr(e, "type", None) == "error"]
            caps = [m for m in msgs if "cold_load_capacity" in m]
            check("stream(): the arrival past the cap gets ONE honest capacity "
                  "error event", len(caps) == 1)
            check("stream(): the refused caller got no token and no loading "
                  "status (it was never admitted)",
                  min(len(_etypes(a)), len(_etypes(b))) <= 2)
            check("stream(): admissions returned after both streams ended",
                  remote._hold_count() == 0)

            # A WARM stream must not squat an admission for its whole generation.
            seen = {"held_at_first_token": None}

            async def ws_warm(worker, payload, rid):
                yield TokenEvent(request_id=rid, text="a")
                await asyncio.sleep(0.05)
                yield TokenEvent(request_id=rid, text="b")
                yield DoneEvent(request_id=rid, input_tokens=1, output_chunks=2,
                                finish_reason="stop")

            remote._worker_stream = ws_warm

            async def _drive():
                first = True
                async for ev in runner.stream(_req("warm-s")):
                    if first and getattr(ev, "type", None) == "token":
                        first = False
                        seen["held_at_first_token"] = remote._hold_count()

            asyncio.run(_drive())
            check("a WARM stream releases its admission at the FIRST TOKEN — a "
                  "long generation never squats a cold-hold slot",
                  seen["held_at_first_token"] == 0)
            check("stream(): counter clean after a warm generation",
                  remote._hold_count() == 0)
    finally:
        remote._select, remote._worker_stream = orig_select, orig_ws
        remote.set_load_state_provider(orig_ls)


# ---------------------------------------------------------------------------
def _acceptance_checks(remote):
    """THE ACCEPTANCE TEST: a lightweight endpoint still answers while every
    cold-hold admission is taken.

    Models the real constraint: one gunicorn process = 8 request threads shared
    by /health, the console and the API. Twenty-four cold calls arrive (the
    script-iterating-every-model shape). Without a cap all 8 threads block in
    holds and a health check waits behind them; with the cap at most 3 threads
    can be held and the other 5 stay free, so health answers promptly.
    """
    fw, tk = _text_framework(remote)
    Runner = remote.make_delegating_runner(fw, tk)
    runner = Runner(types.SimpleNamespace(model_key="cold-model"))

    orig_select, orig_run = remote._select, remote._worker_run_once
    orig_ls = remote._load_state_provider
    remote._select = lambda mk, pool=None, task=None, **kw: (dict(WORKER), None)
    remote.set_load_state_provider(
        lambda mk, wid, since=0.0: {"healthy": False, "in_progress": True})

    GUNICORN_THREADS = 8
    HOLD_S = 1.5
    try:
        with _Env(HUGPY_CENTRAL_GATE="off", HUGPY_COLD_HOLD_POLL_S="0.05",
                  HUGPY_COLD_HOLD_MAX_S=str(HOLD_S),
                  HUGPY_COLD_HOLD_STALL_S=str(HOLD_S),
                  HUGPY_COLD_HOLD_MAX_CONCURRENT="3"):

            async def always_cold(worker, payload, result_type, request_id, model_key):
                await asyncio.sleep(0.05)
                raise RuntimeError("RemoteProtocolError: Server disconnected")

            remote._worker_run_once = always_cold
            outcomes = []
            out_lock = threading.Lock()

            def one_llm_request(i):
                try:
                    asyncio.run(runner.run(_req(f"flood-{i}")))
                    tag = "ok"
                except remote.ColdHoldCapacityError:
                    tag = "refused"
                except Exception:
                    tag = "held"
                with out_lock:
                    outcomes.append(tag)

            def health():
                """A lightweight endpoint: touches no model, just needs a thread."""
                return {"status": "ok"}

            pool = ThreadPoolExecutor(max_workers=GUNICORN_THREADS)
            try:
                flood = [pool.submit(one_llm_request, i) for i in range(24)]
                time.sleep(0.25)          # let the flood take what it can
                held_now = remote._hold_count()
                t0 = time.monotonic()
                hp = pool.submit(health)
                body = hp.result(timeout=120)
                latency = time.monotonic() - t0
                for f in flood:
                    f.result(timeout=120)
            finally:
                pool.shutdown(wait=True)

            check("with 24 cold calls in flight, holds never exceed the cap",
                  held_now <= 3)
            # The bar is "sooner than a SINGLE hold would have finished". The
            # measured uncapped control on this box was 4.5s at HOLD_S=1.5 —
            # i.e. health queued behind three full rounds of holds.
            check("ACCEPTANCE: /health-class endpoint still answers while the "
                  f"cap is saturated (latency {latency:.3f}s < one hold "
                  f"{HOLD_S}s)",
                  body == {"status": "ok"} and latency < HOLD_S)
            check("the flood was bounded — most arrivals were refused fast, not "
                  "parked in a slot",
                  outcomes.count("refused") >= 24 - GUNICORN_THREADS)
            check("every request got a definite answer (nothing left hanging)",
                  len(outcomes) == 24)
            check("all admissions returned after the flood", remote._hold_count() == 0)
    finally:
        remote._select, remote._worker_run_once = orig_select, orig_run
        remote.set_load_state_provider(orig_ls)


# ---------------------------------------------------------------------------
def _probe_checks(cl):
    """The socket probe: a hung-up peer, and ONLY a hung-up peer, reads gone."""
    a, b = socket.socketpair()
    try:
        p = cl.SocketProbe(a)
        check("a live peer reads CONNECTED", p.gone() is False)
        b.send(b"POST /v1/chat/completions")
        check("pipelined bytes are NOT a disconnect", p.gone() is False)
    finally:
        a.close()
        b.close()

    a, b = socket.socketpair()
    try:
        p = cl.SocketProbe(a)
        check("still connected before the peer hangs up", p.gone() is False)
        b.close()
        time.sleep(0.05)
        check("a peer that hung up reads GONE (FIN -> zero-byte peek)",
              p.gone() is True)
        check("gone latches (a second read is still gone)", p.gone() is True)
    finally:
        a.close()

    # Degradation: everything we cannot honestly probe reads as CONNECTED.
    check("no environ -> no probe (nothing changes)",
          cl.probe_for_environ(None) is None)
    check("a WSGI server that publishes no socket -> no probe",
          cl.probe_for_environ({"REQUEST_METHOD": "POST"}) is None)
    check("a non-socket object -> no probe",
          cl.probe_for_environ({"gunicorn.socket": object()}) is None)

    class _Exploding:
        def recv(self, *a, **k):
            raise ValueError("non-zero flags not allowed on an SSL socket")

        def fileno(self):
            return 0

    check("a probe whose recv EXPLODES reads connected, never gone "
          "(abandoning a live caller would be the worse bug)",
          cl.SocketProbe(_Exploding()).gone() is False)

    os.environ["HUGPY_CLIENT_DISCONNECT_ABANDON"] = "off"
    try:
        a, b = socket.socketpair()
        check("the feature has an off switch (HUGPY_CLIENT_DISCONNECT_ABANDON=off)",
              cl.enabled() is False and cl.probe_for_environ(
                  {"gunicorn.socket": a}) is None)
        a.close()
        b.close()
    finally:
        os.environ.pop("HUGPY_CLIENT_DISCONNECT_ABANDON", None)

    check("poll interval defaults to 2s", cl.poll_s() == 2.0)
    os.environ["HUGPY_CLIENT_DISCONNECT_POLL_S"] = "0.5"
    try:
        check("poll interval is an env knob", cl.poll_s() == 0.5)
    finally:
        os.environ.pop("HUGPY_CLIENT_DISCONNECT_POLL_S", None)


# ---------------------------------------------------------------------------
def _abandon_checks(cl, rt):
    """The WSGI->loop bridge gives the request slot back when the caller leaves."""
    os.environ["HUGPY_CLIENT_DISCONNECT_POLL_S"] = "0.05"
    a, b = socket.socketpair()
    unwound = {"n": 0}
    try:
        async def held_call():
            try:
                await asyncio.sleep(30)      # a 25-minute cold hold, in miniature
                return "served"
            finally:
                unwound["n"] += 1            # the `finally` that frees relay slots

        cl.bind(cl.SocketProbe(a))
        b.close()                            # the caller gives up
        t0 = time.monotonic()
        raised = None
        try:
            rt.run(held_call())
        except cl.ClientGone as exc:
            raised = exc
        dt = time.monotonic() - t0
        check("a disconnected caller's work is ABANDONED, not held to the ceiling",
              raised is not None)
        check(f"the request slot comes back promptly ({dt:.2f}s, not 30s)", dt < 2.0)
        time.sleep(0.2)
        check("cancellation unwound the coroutine's finally (relay slots and "
              "cold-hold admissions are released)", unwound["n"] == 1)
    finally:
        cl.clear()
        a.close()
        os.environ.pop("HUGPY_CLIENT_DISCONNECT_POLL_S", None)

    # A LIVE caller is never abandoned — the patient-client guarantee.
    a, b = socket.socketpair()
    os.environ["HUGPY_CLIENT_DISCONNECT_POLL_S"] = "0.05"
    try:
        async def slow_but_wanted():
            await asyncio.sleep(0.6)
            return "served"

        cl.bind(cl.SocketProbe(a))
        check("a slow call with a LIVE client still completes normally",
              rt.run(slow_but_wanted()) == "served")
    finally:
        cl.clear()
        a.close()
        b.close()
        os.environ.pop("HUGPY_CLIENT_DISCONNECT_POLL_S", None)

    # No probe bound (background work, worker agent, internal drains) => unchanged.
    async def plain():
        return 7

    cl.clear()
    check("with no probe bound, run() is byte-identical to before", rt.run(plain()) == 7)

    # iter_sync: the /v1 non-streaming drain never writes, so a failed write can
    # never reveal the disconnect — the poll must.
    a, b = socket.socketpair()
    os.environ["HUGPY_CLIENT_DISCONNECT_POLL_S"] = "0.05"
    drained = {"closed": False}
    try:
        async def forever():
            try:
                yield b"first\n"
                await asyncio.sleep(30)
                yield b"never\n"
            finally:
                drained["closed"] = True

        cl.bind(cl.SocketProbe(a))
        b.close()
        t0 = time.monotonic()
        got = list(rt.iter_sync(forever()))
        dt = time.monotonic() - t0
        check("iter_sync ends the drain when the caller disconnects "
              f"({dt:.2f}s)", got == [b"first\n"] and dt < 2.0)
        check("the abandoned async generator was closed (httpx relay released)",
              drained["closed"] is True)
    finally:
        cl.clear()
        a.close()
        os.environ.pop("HUGPY_CLIENT_DISCONNECT_POLL_S", None)

    # Heartbeats still pace correctly when a probe is present (finer poll must
    # not turn every poll tick into a keepalive write).
    a, b = socket.socketpair()
    os.environ["HUGPY_CLIENT_DISCONNECT_POLL_S"] = "0.02"
    try:
        async def one_slow_event():
            await asyncio.sleep(0.35)
            yield b"data\n"

        cl.bind(cl.SocketProbe(a))
        out = list(rt.iter_sync(one_slow_event(), heartbeat=b":ka\n",
                                heartbeat_secs=0.1))
        beats = out.count(b":ka\n")
        check(f"heartbeats still pace on heartbeat_secs, not on the poll "
              f"({beats} beats over 0.35s at 0.1s)", 2 <= beats <= 4)
        check("the real event still arrived", out[-1] == b"data\n")
    finally:
        cl.clear()
        a.close()
        b.close()
        os.environ.pop("HUGPY_CLIENT_DISCONNECT_POLL_S", None)


# ---------------------------------------------------------------------------
def _route_checks():
    """The refusal reaches an HTTP client as a 503 with Retry-After."""
    try:
        v1 = importlib.import_module(
            "abstract_hugpy_dev.flask_app.app.routes.v1_routes")
        prompt = importlib.import_module(
            "abstract_hugpy_dev.flask_app.app.routes.prompt_routes")
    except Exception as exc:  # pragma: no cover — heavy import in a bare env
        print(f"  ~ skip route checks ({type(exc).__name__}: {exc})")
        return
    remote = importlib.import_module("abstract_hugpy_dev.managers.resolvers.remote")
    err = remote.ColdHoldCapacityError("m", WORKER, 4, 4, loading=True)

    check("v1 classifies the capacity refusal (503 branch, not 500)",
          v1._is_capacity_message(err.stream_message()) is True)
    check("v1 does not confuse it with an ordinary error",
          v1._is_capacity_message("worker exploded") is False)
    check("v1 reads Retry-After from the same knob the message quotes",
          v1._capacity_retry_after() == remote._cold_hold_retry_after_s())

    from flask import Flask
    app = Flask(__name__)
    with app.app_context():
        out = v1._openai_error("busy", 503, "server_busy", retry_after=20)
        check("v1 error helper emits a Retry-After header on a 503",
              len(out) == 3 and out[1] == 503 and out[2]["Retry-After"] == "20")
        out = v1._openai_error("nope", 400)
        check("v1 error helper is unchanged without Retry-After",
              len(out) == 2 and out[1] == 400)
        ref = prompt._capacity_refusal(err)
        check("/prompt turns the refusal into 503 + Retry-After (never a 500)",
              ref is not None and ref[1] == 503 and ref[2]["Retry-After"] == "20")
        body = ref[0].get_json()
        check("/prompt 503 body carries the honest message and the limit",
              "cold_load_capacity" in body["error"] and body["limit"] == 4)
        check("/prompt leaves ordinary failures alone (still a 500 path)",
              prompt._capacity_refusal(RuntimeError("boom")) is None)

    from abstract_hugpy_dev._platform.client_liveness import ClientGone
    check("/prompt recognises a client-disconnect as a non-failure",
          prompt._client_gone(ClientGone("gone")) is True
          and prompt._client_gone(RuntimeError("boom")) is False)


# ---------------------------------------------------------------------------
def test_cold_hold_cap():
    global ok
    ok = 0
    remote = importlib.import_module("abstract_hugpy_dev.managers.resolvers.remote")
    cl = importlib.import_module("abstract_hugpy_dev._platform.client_liveness")
    rt = importlib.import_module("abstract_hugpy_dev._platform.async_runtime")
    _knob_checks(remote)
    _message_checks(remote)
    _admission_checks(remote)
    _runner_checks(remote)
    _stream_checks(remote)
    _acceptance_checks(remote)
    _probe_checks(cl)
    _abandon_checks(cl, rt)
    _route_checks()
    print(f"\nall {ok} checks passed")


if __name__ == "__main__":
    test_cold_hold_cap()
