@miller-tech/uap 1.172.12 → 1.172.14

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@miller-tech/uap",
-  "version": "1.172.12",
+  "version": "1.172.14",
   "description": "Autonomous AI agent memory system with CLAUDE.md protocol enforcement",
   "type": "module",
   "main": "dist/index.js",

package/tools/agents/scripts/anthropic_proxy.py

@@ -2766,7 +2766,7 @@ def _detach_aclose(closeable) -> None:
 
 def _build_http_client() -> httpx.AsyncClient:
     """Upstream client factory — used at startup AND by pool self-healing."""
-    c = httpx.AsyncClient(
+    c = DisconnectAwareClient(
         timeout=httpx.Timeout(
             connect=10.0,  # 10s to establish connection
             read=PROXY_READ_TIMEOUT,  # configurable (default 10 min)
@@ -2970,9 +2970,30 @@ _current_request_session: contextvars.ContextVar[str | None] = contextvars.Conte
 #
 # Same ContextVar trick as the session above: request-local without threading a
 # Request through every signature. Holds Starlette's `request.is_disconnected`.
-_current_client_gone: contextvars.ContextVar[
-    Callable[[], Awaitable[bool]] | None
-] = contextvars.ContextVar("uap_current_client_gone", default=None)
+# Is the DOWNSTREAM client still there?
+#
+# Nothing used to ask. Once a turn was accepted, the guardrail loops kept
+# retrying and re-POSTing upstream on their own — so when the client died the
+# work carried on regardless. Observed live: two `uap deliver` runs were killed,
+# no socket remained on :4000, and the proxy still drove llama for minutes,
+# generating 32k tokens per orphaned turn that nobody would ever read.
+#
+# WHY NOT request.is_disconnected(). That was tried and MEASURED not to work
+# here: it returned False 16 times across 30s while the caller was demonstrably
+# gone. The cause is `@app.middleware("http")` — Starlette's BaseHTTPMiddleware
+# runs the endpoint in a separate anyio task behind its own receive channel, so
+# the endpoint polls a channel that never delivers http.disconnect. It works
+# fine in an app with no middleware, which is exactly why three rounds of
+# isolated testing said it was sound.
+#
+# So the signal is captured by a pure-ASGI middleware OUTSIDE all of that (see
+# DisconnectWatcherMiddleware) and published here as a mutable holder. A dict
+# rather than a bool because the middleware's poller task and the endpoint must
+# see the same object: the ContextVar is set before the app is invoked, so every
+# downstream task inherits the reference, and mutations are visible across them.
+_disconnect_holder: contextvars.ContextVar[dict | None] = contextvars.ContextVar(
+    "uap_disconnect_holder", default=None
+)
 
 
 class ClientGoneError(Exception):
@@ -2986,19 +3007,17 @@ class ClientGoneError(Exception):
 
 
 async def _client_gone() -> bool:
-    """True when the caller has hung up. Never raises, never blocks meaningfully.
+    """True when the caller has hung up.
 
-    A disconnect check that can itself fail or stall would be worse than no
-    check: it sits in front of every upstream call. Any error means "assume the
-    client is still there" — the pre-existing behaviour.
+    Now a flag read rather than a call into Starlette, so there is nothing to
+    raise and nothing for a fail-safe except to swallow. The previous version
+    wrapped request.is_disconnected() in `except Exception: return False`, which
+    would have hidden a failing probe as "client still present" — it turned out
+    the probe was not raising, but the shape was wrong regardless.
     """
-    probe = _current_client_gone.get()
-    if probe is None:
-        return False
-    try:
-        return bool(await probe())
-    except Exception:  # noqa: BLE001
-        return False
+    holder = _disconnect_holder.get()
+    return bool(holder and holder.get("gone"))
+
 
 # Session admission state. _admitted_sessions maps session_id -> last-seen
 # monotonic ts; OrderedDict insertion order is the LRU (oldest = front).
@@ -3364,47 +3383,50 @@ PROXY_DISCONNECT_POLL_SECS = max(
 )
 
 
-async def _post_watching_client(
-    client: httpx.AsyncClient, url: str, payload: dict, headers: dict
-) -> httpx.Response:
-    """POST upstream, cancelling it if the caller hangs up mid-generation.
+class DisconnectAwareClient(httpx.AsyncClient):
+    """An httpx client that abandons a call when the caller hangs up.
 
-    Checking only BEFORE each upstream call stops the retry loops, but does
-    nothing for a generation already in flight: a single turn is one POST, so
-    there is no next call to block and the model runs to completion — up to
-    PROXY_TOOL_TURN_MAX_TOKENS (32k, ~13 minutes here) for a client that has
-    already gone. This closes that.
+    THE CHOKE POINT. The previous attempt guarded ONE call site and claimed to
+    cover them all; there are fourteen, and the guardrail loops that caused the
+    incident call the model directly rather than through _post_with_retry. So
+    the check belongs where it cannot be missed: every high-level httpx method —
+    post(), stream(), request() — funnels through send(), so overriding send()
+    covers all of them, including any added later.
 
-    Cancelling the httpx task closes the upstream connection, and llama.cpp
-    releases the slot when its client disappears — verified directly against the
-    running server: slots went 2 -> 0 within 15s of closing the socket. So this
-    genuinely frees the GPU rather than merely letting the proxy stop waiting.
+    Cancelling matters rather than merely returning: closing the connection is
+    what makes llama.cpp release the slot (verified against the running server —
+    slots went 2 -> 0 within 15s of a raw socket close). Returning early without
+    cancelling would leave the model generating exactly as before.
 
-    No probe (background task, tests) means no watching: plain await, unchanged.
+    No probe (background task, health check, tests) means a plain send,
+    unwatched. Streaming sends are covered for their header phase; a body being
+    written to a vanished client fails on write anyway.
     """
-    probe = _current_client_gone.get()
-    if probe is None:
-        return await client.post(url, json=payload, headers=headers)
 
-    post_task = asyncio.ensure_future(client.post(url, json=payload, headers=headers))
-    try:
-        while True:
-            done, _pending = await asyncio.wait(
-                {post_task}, timeout=PROXY_DISCONNECT_POLL_SECS
-            )
-            if post_task in done:
-                return post_task.result()
-            if await _client_gone():
-                post_task.cancel()
-                try:
-                    await post_task
-                except (asyncio.CancelledError, Exception):  # noqa: BLE001
-                    pass  # cancellation is the point; the connection is closed
-                raise ClientGoneError("client disconnected mid-generation")
-    finally:
-        # Never leave the upstream call running after we stop waiting on it.
-        if not post_task.done():
-            post_task.cancel()
+    async def send(self, request, **kwargs):  # type: ignore[override]
+        holder = _disconnect_holder.get()
+        if holder is None:
+            return await super().send(request, **kwargs)
+
+        task = asyncio.ensure_future(super().send(request, **kwargs))
+        try:
+            while True:
+                done, _pending = await asyncio.wait(
+                    {task}, timeout=PROXY_DISCONNECT_POLL_SECS
+                )
+                if task in done:
+                    return task.result()
+                if await _client_gone():
+                    task.cancel()
+                    try:
+                        await task
+                    except (asyncio.CancelledError, Exception):  # noqa: BLE001
+                        pass  # cancellation is the point; the connection is closed
+                    raise ClientGoneError("client disconnected mid-generation")
+        finally:
+            # Never leave an upstream call running once we stop waiting on it.
+            if not task.done():
+                task.cancel()
 
 
 async def _post_with_retry_inner(
@@ -3418,7 +3440,7 @@ async def _post_with_retry_inner(
         try:
             _inflight_inc(client)
             try:
-                resp = await _post_watching_client(client, url, payload, headers)
+                resp = await client.post(url, json=payload, headers=headers)
             finally:
                 _inflight_dec(client)
             # Cycle 19 Option 1: if 503 "Loading model", wait for health then retry
@@ -3495,7 +3517,7 @@ async def _check_slot_hang(slot_url: str) -> bool:
     if PROXY_SLOT_HANG_TIMEOUT <= 0:
         return False
     try:
-        async with httpx.AsyncClient() as check_client:
+        async with DisconnectAwareClient() as check_client:
             resp = await check_client.get(slot_url, timeout=5.0)
             if resp.status_code != 200:
                 return False
@@ -3632,6 +3654,73 @@ async def lifespan(app: FastAPI):
     logger.info("Proxy shut down")
 
 
+class DisconnectWatcherMiddleware:
+    """Pure-ASGI middleware that records when the caller hangs up.
+
+    Sits OUTERMOST, above the auth middleware, and is deliberately pure ASGI:
+    Starlette's BaseHTTPMiddleware (`@app.middleware("http")`) re-tasks the
+    endpoint behind its own receive channel, which is precisely why
+    request.is_disconnected() reported False 16 times across 30s while the
+    caller was gone. Being outside that, this sees the real channel.
+
+    THE HARD PART is that wrapping receive() is not enough on its own. After the
+    request body is fully read nothing calls receive() again, so an
+    http.disconnect just sits in the channel unobserved. So once the body is
+    complete this starts polling the channel itself — safe precisely because the
+    app has no more body to ask for, and the only remaining message type is the
+    disconnect.
+
+    The flag lives in a mutable dict published on a ContextVar before the app is
+    invoked: downstream tasks inherit the reference when their context is copied,
+    so a mutation made by the poller is visible to the endpoint and to every
+    upstream call it makes.
+    """
+
+    def __init__(self, app):
+        self.app = app
+
+    async def __call__(self, scope, receive, send):
+        if scope.get("type") != "http":
+            return await self.app(scope, receive, send)
+
+        holder: dict = {"gone": False}
+        _disconnect_holder.set(holder)
+        poller: asyncio.Task | None = None
+
+        async def poll_for_disconnect() -> None:
+            # Only reached once the body is complete, so this cannot steal a
+            # message the app still needs.
+            try:
+                while not holder["gone"]:
+                    msg = await receive()
+                    if msg.get("type") == "http.disconnect":
+                        holder["gone"] = True
+                        return
+            except asyncio.CancelledError:
+                raise
+            except Exception:  # noqa: BLE001
+                return  # channel closed/erroring: stop watching, never crash
+
+        async def watched_receive():
+            nonlocal poller
+            msg = await receive()
+            if msg.get("type") == "http.disconnect":
+                holder["gone"] = True
+            elif msg.get("type") == "http.request" and not msg.get("more_body", False):
+                # Body complete — nobody will call receive() again, so take over.
+                if poller is None:
+                    poller = asyncio.ensure_future(poll_for_disconnect())
+            return msg
+
+        try:
+            await self.app(scope, watched_receive, send)
+        finally:
+            if poller is not None and not poller.done():
+                poller.cancel()
+
+
+
+
 app = FastAPI(
     title="UAP Anthropic Proxy",
     description="Translates Anthropic Messages API to OpenAI Chat Completions API",
@@ -3639,6 +3728,7 @@ app = FastAPI(
     lifespan=lifespan,
 )
 
+
 @app.exception_handler(ClientGoneError)
 async def _client_gone_handler(request: Request, exc: ClientGoneError):
     """The caller hung up mid-turn: stop, and say so at INFO rather than ERROR.
@@ -3693,6 +3783,18 @@ async def _pool_timeout_handler(request: Request, exc: httpx.PoolTimeout):
 _PROXY_AUTH_OPEN_PATHS = frozenset({"/health", "/", "/v1/models"})
 
 
+# OUTERMOST — and it must be registered AFTER the auth middleware below to BE
+# outermost: Starlette's add_middleware inserts at position 0, so the last one
+# registered is the outermost wrapper. (Registered before auth first time round,
+# which would have put it INSIDE the very BaseHTTPMiddleware that hides
+# http.disconnect — the exact bug it exists to work around.)
+#
+# The registration therefore trails the auth definition on purpose. Auth itself
+# is deliberately untouched.
+def _install_disconnect_watcher() -> None:
+    app.add_middleware(DisconnectWatcherMiddleware)
+
+
 @app.middleware("http")
 async def _shared_secret_auth(request: Request, call_next):
     """Gate every request behind PROXY_AUTH_TOKEN when it is set.
@@ -3751,6 +3853,10 @@ async def _shared_secret_auth(request: Request, call_next):
 # ===========================================================================
 
 
+
+# Auth is now defined, so this lands OUTSIDE it (add_middleware inserts at 0).
+_install_disconnect_watcher()
+
 def _image_block_to_openai(block: dict) -> dict | None:
     """Anthropic image block → OpenAI image_url part (data URI or URL)."""
     src = block.get("source") or {}
@@ -10905,10 +11011,6 @@ async def messages(request: Request):
     """
     global last_session_id
 
-    # Publish this request's disconnect probe for the whole call tree. Set before
-    # ANY upstream work so an already-dead client is caught on the first call.
-    _current_client_gone.set(request.is_disconnected)
-
     body = await request.json()
     is_stream = body.get("stream", False)
     model = body.get("model", "default")

package/tools/agents/tests/test_client_disconnect.py

@@ -28,33 +28,32 @@ proxy = _load_proxy()
 
 
 class ClientGoneProbeTest(unittest.TestCase):
-    def test_absent_probe_means_client_is_present(self):
-        """No probe set (background task, test harness) must not read as gone —
-        that would abandon perfectly live turns."""
-        proxy._current_client_gone.set(None)
+    """The probe is now a FLAG READ, not a call into Starlette.
+
+    It used to call request.is_disconnected() wrapped in `except Exception:
+    return False`. That was measured not to work at all behind
+    BaseHTTPMiddleware (False 16x across 30s while the caller was gone), and the
+    swallow-everything shape would have hidden a failing probe as "client still
+    present". Both are gone: there is nothing left to raise.
+    """
+
+    def test_no_holder_means_client_is_present(self):
+        """Background tasks and health checks have no request context. They must
+        not read as disconnected."""
+        proxy._disconnect_holder.set(None)
         self.assertFalse(asyncio.run(proxy._client_gone()))
 
     def test_reports_disconnect(self):
-        async def gone() -> bool:
-            return True
-
-        proxy._current_client_gone.set(gone)
+        proxy._disconnect_holder.set({"gone": True})
         self.assertTrue(asyncio.run(proxy._client_gone()))
 
     def test_reports_connected(self):
-        async def here() -> bool:
-            return False
-
-        proxy._current_client_gone.set(here)
+        proxy._disconnect_holder.set({"gone": False})
         self.assertFalse(asyncio.run(proxy._client_gone()))
 
-    def test_a_failing_probe_never_breaks_the_request(self):
-        """This check sits in front of every upstream call. A probe that can
-        raise would be worse than no probe, so failure means 'still there'."""
-        async def broken() -> bool:
-            raise RuntimeError("receive channel exploded")
-
-        proxy._current_client_gone.set(broken)
+    def test_a_malformed_holder_reads_as_present(self):
+        """Defensive: an empty dict must not be mistaken for a disconnect."""
+        proxy._disconnect_holder.set({})
         self.assertFalse(asyncio.run(proxy._client_gone()))
 
 

package/tools/agents/tests/test_disconnect_watcher.py

@@ -0,0 +1,174 @@
+#!/usr/bin/env python3
+"""Pure-ASGI disconnect watcher — the thing that finally made the abort work.
+
+Four attempts failed before this. The measured cause: `@app.middleware("http")`
+is Starlette's BaseHTTPMiddleware, which runs the endpoint in a separate anyio
+task behind its own receive channel, so request.is_disconnected() in the endpoint
+polls a channel that never delivers http.disconnect. Instrumented in production
+it returned False 16 times across 30s while the caller was demonstrably gone —
+and it worked perfectly in an app with no middleware, which is exactly why three
+rounds of isolated testing said the approach was sound.
+
+Two properties are load-bearing and each has a test that fails if it breaks:
+
+  1. ORDER. The watcher must be OUTERMOST. Starlette's add_middleware inserts at
+     position 0, so it must be registered LAST. The first implementation
+     registered it before auth, putting it inside the very middleware whose
+     behaviour it exists to work around.
+  2. POLLING. Wrapping receive() is not enough. Once the body is fully read
+     nothing calls receive() again, so an http.disconnect sits in the channel
+     unobserved. The watcher must take over polling after the final body message.
+
+Verified end to end after these were in place: the llama slot freed within 5s of
+the client hanging up, against 30s+ of continued generation before.
+"""
+
+import ast
+import asyncio
+import importlib.util
+import unittest
+from pathlib import Path
+
+PROXY_PATH = Path(__file__).resolve().parents[1] / "scripts" / "anthropic_proxy.py"
+
+
+def _load_proxy():
+    spec = importlib.util.spec_from_file_location("anthropic_proxy_watcher", PROXY_PATH)
+    assert spec is not None and spec.loader is not None
+    m = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(m)
+    return m
+
+
+proxy = _load_proxy()
+
+
+class MiddlewareOrderTest(unittest.TestCase):
+    def test_the_watcher_is_outermost(self):
+        """If it ends up inside BaseHTTPMiddleware it cannot see http.disconnect,
+        which is the whole failure this replaced."""
+        names = [getattr(mw.cls, "__name__", str(mw.cls)) for mw in proxy.app.user_middleware]
+        self.assertTrue(names, "no middleware registered at all")
+        self.assertEqual(
+            names[0],
+            "DisconnectWatcherMiddleware",
+            f"watcher is not outermost; stack is {names}",
+        )
+
+    def test_auth_middleware_is_still_installed(self):
+        """The brief was to leave auth alone. Losing it would silently open the
+        LAN token gate — a far worse outcome than the bug being fixed."""
+        names = [getattr(mw.cls, "__name__", str(mw.cls)) for mw in proxy.app.user_middleware]
+        self.assertIn("BaseHTTPMiddleware", names)
+
+    def test_probe_is_a_flag_read_not_a_starlette_call(self):
+        """_client_gone must not call back into the request object: that is the
+        path that measured False-forever behind BaseHTTPMiddleware.
+
+        Checked with ast rather than substring matching — the first version of
+        this test failed on the word appearing in the function's own docstring,
+        which is a false positive, not a finding.
+        """
+        tree = ast.parse(PROXY_PATH.read_text())
+        fn = next(
+            n for n in ast.walk(tree)
+            if isinstance(n, ast.AsyncFunctionDef) and n.name == "_client_gone"
+        )
+        attrs = {
+            n.attr for n in ast.walk(fn) if isinstance(n, ast.Attribute)
+        }
+        self.assertNotIn("is_disconnected", attrs)
+        # And it must read the holder, so the test fails if the body is gutted.
+        self.assertIn("get", attrs)
+
+
+class WatcherBehaviourTest(unittest.TestCase):
+    """Drive the middleware directly with an ASGI receive/send pair."""
+
+    def _run(self, messages, app_body):
+        """messages: what receive() yields, in order. app_body: coroutine fn(scope, receive, send)."""
+        sent = []
+        queue = list(messages)
+
+        async def receive():
+            if queue:
+                return queue.pop(0)
+            await asyncio.sleep(3600)  # channel idle, like a live connection
+
+        async def send(msg):
+            sent.append(msg)
+
+        mw = proxy.DisconnectWatcherMiddleware(app_body)
+        asyncio.run(mw({"type": "http"}, receive, send))
+        return sent
+
+    def test_records_a_disconnect_the_app_itself_reads(self):
+        seen = {}
+
+        async def app(scope, receive, send):
+            await receive()  # body
+            msg = await receive()  # disconnect
+            seen["type"] = msg["type"]
+            seen["gone"] = await proxy._client_gone()
+
+        self._run(
+            [{"type": "http.request", "body": b"{}", "more_body": False},
+             {"type": "http.disconnect"}],
+            app,
+        )
+        self.assertEqual(seen["type"], "http.disconnect")
+        self.assertTrue(seen["gone"], "watcher did not record the disconnect")
+
+    def test_records_a_disconnect_the_app_NEVER_reads(self):
+        """The real case. After the body is consumed the app never calls receive()
+        again, so the disconnect is only ever seen by the watcher's own poller.
+        Wrapping receive() alone would leave `gone` False here forever."""
+        seen = {}
+
+        async def app(scope, receive, send):
+            await receive()  # body only — then just "work", like a generation
+            for _ in range(40):
+                await asyncio.sleep(0.02)
+                if await proxy._client_gone():
+                    seen["gone"] = True
+                    return
+            seen["gone"] = False
+
+        self._run(
+            [{"type": "http.request", "body": b"{}", "more_body": False},
+             {"type": "http.disconnect"}],
+            app,
+        )
+        self.assertTrue(seen.get("gone"), "poller did not observe the unread disconnect")
+
+    def test_a_live_request_is_never_marked_gone(self):
+        """A false positive would abort healthy turns — worse than the bug."""
+        seen = {}
+
+        async def app(scope, receive, send):
+            await receive()
+            for _ in range(10):
+                await asyncio.sleep(0.02)
+            seen["gone"] = await proxy._client_gone()
+
+        self._run([{"type": "http.request", "body": b"{}", "more_body": False}], app)
+        self.assertFalse(seen["gone"])
+
+    def test_non_http_scopes_pass_straight_through(self):
+        """Lifespan and websocket scopes have no disconnect semantics here and
+        must not be wrapped."""
+        called = {}
+
+        async def app(scope, receive, send):
+            called["type"] = scope["type"]
+
+        async def go():
+            mw = proxy.DisconnectWatcherMiddleware(app)
+            await mw({"type": "lifespan"}, None, None)
+
+        asyncio.run(go())
+        self.assertEqual(called["type"], "lifespan")
+
+
+if __name__ == "__main__":
+    unittest.main()

package/tools/agents/tests/test_upstream_chokepoint.py

@@ -0,0 +1,178 @@
+#!/usr/bin/env python3
+"""Every upstream call must pass the disconnect check — structurally, not by claim.
+
+This test exists because of a specific mistake. Two releases (v1.172.11,
+v1.172.12) guarded ONE call site while the commit message asserted that "every
+upstream call funnels through _post_with_retry, so one check covers them all".
+There are fourteen upstream call sites, and the guardrail loops that caused the
+incident — malformed-tool retry, completion-contract, empty-max-tokens recovery,
+unexpected-end-turn, recipe — call the model DIRECTLY. Both releases were
+therefore near-no-ops for the workload they were written for, and every unit
+test still passed.
+
+The fix was to move the check somewhere it cannot be missed: httpx routes every
+high-level method (post, stream, request) through send(), so DisconnectAwareClient
+overrides send() and the coverage becomes a property of the client rather than a
+property of remembering.
+
+These tests pin that property. A future call site added anywhere in the proxy is
+covered automatically; a future *client* constructed bare would not be, and that
+is exactly what the enumeration below fails on.
+"""
+
+import ast
+import asyncio
+import importlib.util
+import re
+import unittest
+from pathlib import Path
+
+PROXY_PATH = Path(__file__).resolve().parents[1] / "scripts" / "anthropic_proxy.py"
+
+
+def _load_proxy():
+    spec = importlib.util.spec_from_file_location("anthropic_proxy_choke", PROXY_PATH)
+    assert spec is not None and spec.loader is not None
+    m = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(m)
+    return m
+
+
+proxy = _load_proxy()
+SRC = PROXY_PATH.read_text()
+
+
+class ChokePointEnumerationTest(unittest.TestCase):
+    def test_no_bare_httpx_client_is_ever_constructed(self):
+        """The invariant. A bare httpx.AsyncClient bypasses the override, so its
+        calls silently go unwatched — which is precisely how the first two
+        attempts failed, one call site at a time."""
+        bare = [
+            (i + 1, l.strip())
+            for i, l in enumerate(SRC.split("\n"))
+            if re.search(r"httpx\.AsyncClient\s*\(", l) and "class " not in l
+        ]
+        self.assertEqual(
+            bare, [], f"these construct an unwatched client: {bare}"
+        )
+
+    def test_the_client_actually_subclasses_httpx(self):
+        """Guards the reverse mistake: a wrapper that no longer IS an httpx
+        client would break every caller rather than fail this suite loudly."""
+        import httpx
+
+        self.assertTrue(issubclass(proxy.DisconnectAwareClient, httpx.AsyncClient))
+
+    def test_send_is_overridden_on_the_client_itself(self):
+        """send() is the funnel. If a refactor moves the logic off send() onto,
+        say, post(), then stream() and request() silently lose their guard."""
+        self.assertIn("send", vars(proxy.DisconnectAwareClient))
+
+    def test_every_upstream_call_site_is_reachable_through_send(self):
+        """Inventory, so the count cannot quietly drift.
+
+        Any client.post / client.stream / client.send in the proxy is covered by
+        the override — that is the point of a choke point. This asserts the call
+        sites still exist in numbers consistent with the audit (they were the
+        thing miscounted), and that none of them re-implements its own transport.
+        """
+        calls = re.findall(r"client\.(post|send|stream)\s*\(", SRC)
+        self.assertGreaterEqual(
+            len(calls), 10, "upstream call sites vanished — was transport re-implemented?"
+        )
+        # No call site may build its own client inline and slip the override.
+        self.assertNotRegex(SRC, r"await\s+httpx\.AsyncClient\s*\(")
+
+
+class ChokePointBehaviourTest(unittest.TestCase):
+    """The override must do the job, not merely exist."""
+
+    class _Req:
+        pass
+
+    def setUp(self):
+        proxy.PROXY_DISCONNECT_POLL_SECS = 0.05
+
+    def _client_with_send(self, behaviour):
+        """A DisconnectAwareClient whose SUPERCLASS send is stubbed, so the
+        override under test runs for real."""
+        import httpx
+
+        state = {"cancelled": False}
+
+        async def fake_send(self, request, **kwargs):
+            if behaviour == "fast":
+                return "RESPONSE"
+            try:
+                await asyncio.sleep(3600)
+            except asyncio.CancelledError:
+                state["cancelled"] = True
+                raise
+
+        original = httpx.AsyncClient.send
+        httpx.AsyncClient.send = fake_send  # type: ignore[assignment]
+        client = proxy.DisconnectAwareClient()
+        return client, state, (lambda: setattr(httpx.AsyncClient, "send", original))
+
+    def _set_probe(self, gone: bool):
+        # The guard now reads a shared mutable holder published by the ASGI
+        # watcher, rather than calling back into the request object.
+        proxy._disconnect_holder.set({"gone": gone})
+
+    def test_cancels_upstream_when_the_caller_leaves(self):
+        client, state, restore = self._client_with_send("never")
+        try:
+            self._set_probe(True)
+
+            async def run():
+                with self.assertRaises(proxy.ClientGoneError):
+                    await client.send(self._Req())
+
+            asyncio.run(run())
+            self.assertTrue(state["cancelled"], "upstream call was not cancelled")
+        finally:
+            restore()
+
+    def test_leaves_a_live_call_alone(self):
+        client, state, restore = self._client_with_send("fast")
+        try:
+            self._set_probe(False)
+            self.assertEqual(asyncio.run(client.send(self._Req())), "RESPONSE")
+            self.assertFalse(state["cancelled"])
+        finally:
+            restore()
+
+    def test_no_probe_means_unwatched(self):
+        """Health checks and background tasks have no request context and must
+        keep working rather than acquiring surprise cancellation."""
+        client, state, restore = self._client_with_send("fast")
+        try:
+            proxy._disconnect_holder.set(None)
+            self.assertEqual(asyncio.run(client.send(self._Req())), "RESPONSE")
+        finally:
+            restore()
+
+    def test_a_slow_but_live_call_is_left_pending(self):
+        """Mistaking 'slow' for 'gone' would be worse than the bug being fixed."""
+        client, state, restore = self._client_with_send("never")
+        try:
+            self._set_probe(False)
+
+            async def run():
+                task = asyncio.ensure_future(client.send(self._Req()))
+                await asyncio.sleep(0.4)  # ~8 poll intervals
+                still = not task.done()
+                task.cancel()
+                try:
+                    await task
+                except (asyncio.CancelledError, Exception):
+                    pass
+                return still
+
+            self.assertTrue(asyncio.run(run()), "a live-but-slow call was ended early")
+        finally:
+            restore()
+
+
+if __name__ == "__main__":
+    unittest.main()

package/tools/agents/tests/test_inflight_cancel.py

@@ -1,129 +0,0 @@
-#!/usr/bin/env python3
-"""Cancel the upstream generation when the caller hangs up mid-turn.
-
-Checking only BEFORE each upstream call (v1.172.11) stops the retry loops, but
-does nothing for a generation already in flight: a single turn is ONE post, so
-there is no next call to block and the model runs to completion — up to
-PROXY_TOOL_TURN_MAX_TOKENS (32k, ~13 minutes on this box) for a client that has
-already gone.
-
-That gap was caught by an end-to-end test, not by reasoning: the abandon simply
-never fired for a single long request. These tests pin the closing behaviour so
-it cannot silently regress to the pre-check-only version.
-
-Cancelling matters rather than merely returning: closing the upstream connection
-is what makes llama.cpp release the slot (verified against the running server —
-slots went 2 -> 0 within 15s of a raw socket close).
-"""
-
-import asyncio
-import importlib.util
-import unittest
-from pathlib import Path
-
-
-def _load_proxy():
-    p = Path(__file__).resolve().parents[1] / "scripts" / "anthropic_proxy.py"
-    spec = importlib.util.spec_from_file_location("anthropic_proxy_inflight", p)
-    assert spec is not None and spec.loader is not None
-    m = importlib.util.module_from_spec(spec)
-    spec.loader.exec_module(m)
-    return m
-
-
-proxy = _load_proxy()
-
-
-class _Client:
-    """Stands in for httpx.AsyncClient.post.
-
-    `never` models the real failure: the model is generating, so the POST is
-    neither erroring nor returning — it is simply not done yet.
-    """
-
-    def __init__(self, mode: str):
-        self.mode = mode
-        self.cancelled = False
-
-    async def post(self, url, json=None, headers=None):  # noqa: A002
-        if self.mode == "fast":
-            return "RESPONSE"
-        try:
-            await asyncio.sleep(3600)
-        except asyncio.CancelledError:
-            self.cancelled = True
-            raise
-        return "RESPONSE"
-
-
-def _set_probe(gone: bool):
-    async def probe() -> bool:
-        return gone
-    proxy._current_client_gone.set(probe)
-
-
-class InflightCancelTest(unittest.TestCase):
-    def setUp(self):
-        proxy.PROXY_DISCONNECT_POLL_SECS = 0.05  # keep the tests quick
-
-    def test_cancels_the_upstream_call_when_the_caller_leaves(self):
-        client = _Client("never")
-        _set_probe(True)
-
-        async def run():
-            with self.assertRaises(proxy.ClientGoneError):
-                await proxy._post_watching_client(client, "u", {}, {})
-
-        asyncio.run(run())
-        # The cancellation is the whole point: it closes the connection, which is
-        # what frees the llama slot. Returning early without it would leave the
-        # model generating exactly as before.
-        self.assertTrue(client.cancelled, "upstream POST was not cancelled")
-
-    def test_leaves_a_live_turn_completely_alone(self):
-        client = _Client("fast")
-        _set_probe(False)
-        self.assertEqual(asyncio.run(proxy._post_watching_client(client, "u", {}, {})), "RESPONSE")
-        self.assertFalse(client.cancelled)
-
-    def test_no_probe_means_a_plain_await(self):
-        """Background tasks and tests have no request context. They must keep
-        working, unwatched, rather than acquiring surprise cancellation."""
-        client = _Client("fast")
-        proxy._current_client_gone.set(None)
-        self.assertEqual(asyncio.run(proxy._post_watching_client(client, "u", {}, {})), "RESPONSE")
-
-    def test_a_slow_but_live_turn_is_left_pending(self):
-        """The poll must not mistake 'slow' for 'gone' — long generations are
-        normal, and cancelling them would be a far worse bug than the one fixed.
-
-        Asserted by letting it poll many times and checking it is STILL running.
-        (Wrapping it in wait_for would prove nothing: a wait_for timeout is the
-        caller giving up, and the cleanup correctly cancels upstream then — which
-        is how the first version of this test fooled itself.)
-        """
-        client = _Client("never")
-        _set_probe(False)
-
-        async def run():
-            task = asyncio.ensure_future(proxy._post_watching_client(client, "u", {}, {}))
-            await asyncio.sleep(0.4)  # ~8 poll intervals
-            still_running = not task.done()
-            task.cancel()
-            try:
-                await task
-            except (asyncio.CancelledError, Exception):
-                pass
-            return still_running
-
-        self.assertTrue(asyncio.run(run()), "a live-but-slow turn was ended early")
-
-    def test_poll_interval_has_a_floor(self):
-        """A zero/negative interval would spin the event loop."""
-        self.assertGreaterEqual(proxy.PROXY_DISCONNECT_POLL_SECS, 0.0)
-        mod = _load_proxy()
-        self.assertGreaterEqual(mod.PROXY_DISCONNECT_POLL_SECS, 0.5)
-
-
-if __name__ == "__main__":
-    unittest.main()