fix(runner): offload runner auth resolution off the event loop

_RunnerDatabricksAuth is attached to the runner's httpx.AsyncClient, which
shares the event loop with the WebSocket tunnel's keepalive coroutine. It
defined only the sync auth_flow; httpx drives a sync-only auth_flow inline on
the loop. Since #2762 made runner-local credential resolution lazy (host bearer
used until rejected), the first post-rejection factory call runs on the loop and
can block for seconds — a Databricks CLI shell-out plus, on the managed path, a
blocking httpx.Client mint POST (timeout 10s). That starves the keepalive
coroutine past the 90s ping timeout, dropping the tunnel with
1011 keepalive ping timeout (new since ~2026-07-20, climbing with daily builds).

Add async_auth_flow to _RunnerDatabricksAuth wrapping every factory/header/
invalidate call in asyncio.to_thread; httpx.AsyncClient prefers it, so slow
credential resolution runs in a worker thread and can never stall keepalive.
Auth semantics are unchanged; the sync auth_flow stays for sync-client callers.
Also run the 60s permission-hook snapshot refresh fully off-loop.

Tests: async re-mint through a real AsyncClient, and a blocking-factory guard
proving a concurrent loop task keeps ticking during token resolution.

Signed-off-by: Bryan Qiu <bryan.qiu@databricks.com>

omnigent/runner/_entry.py

@@ -255,6 +255,56 @@ class _RunnerDatabricksAuth(httpx.Auth):
                 request.headers["Authorization"] = f"Bearer {token}"
                 yield request
 
+    async def async_auth_flow(
+        self,
+        request: httpx.Request,
+    ) -> AsyncIterator[httpx.Request]:
+        """Async twin of :meth:`auth_flow` that never blocks the event loop.
+
+        This auth is attached to the runner's ``httpx.AsyncClient`` (see
+        :func:`serve_tunnel`'s ``server_client``), which shares the event
+        loop with the WebSocket tunnel's keepalive coroutine. httpx drives a
+        sync-only ``auth_flow`` generator *inline on the loop*, so the token
+        factory — whose first post-rejection call lazily resolves runner-local
+        auth: a Databricks CLI shell-out (~0.5s) and, on the managed path, a
+        blocking ``httpx.Client`` mint POST (up to 10s, see
+        :func:`_mint_managed_owner_token`) — would stall the loop long enough
+        to miss the server's keepalive ping and drop the tunnel with a
+        ``1011 keepalive ping timeout``. Defining ``async_auth_flow`` makes
+        httpx use *this* instead, and every potentially-blocking call is
+        offloaded with ``asyncio.to_thread`` so a slow credential resolution
+        can never starve keepalive. The auth semantics are otherwise identical
+        to :meth:`auth_flow`.
+
+        :param request: The outgoing httpx request.
+        :yields: The request with the auth header set, or unmodified when no
+            factory is configured.
+        :raises httpx.RequestError: When the factory is configured but
+            returns no token.
+        """
+        if self._server_url:
+            from omnigent.cli_auth import databricks_request_headers
+
+            headers = await asyncio.to_thread(databricks_request_headers, self._server_url)
+            request.headers.update(headers)
+        if self._factory is not None:
+            token = await asyncio.to_thread(self._factory)
+            if not token:
+                if getattr(self._factory, "declined", False):
+                    yield request
+                    return
+                raise httpx.RequestError("Databricks token refresh returned no token")
+            request.headers["Authorization"] = f"Bearer {token}"
+        response = yield request
+        if self._factory is None:
+            return
+        if _is_login_redirect_or_unauthorized(response):
+            await asyncio.to_thread(_invalidate_auth_token_factory, self._factory)
+            token = await asyncio.to_thread(self._factory)
+            if token:
+                request.headers["Authorization"] = f"Bearer {token}"
+                yield request
+
 
 def _is_login_redirect_or_unauthorized(response: httpx.Response) -> bool:
     """Return ``True`` when ``response`` is a re-auth signal.

omnigent/runner/app.py

@@ -414,10 +414,19 @@ async def _refresh_claude_permission_hook_auth(
     while True:
         await asyncio.sleep(refresh_interval_s)
         try:
-            token = await asyncio.to_thread(auth_token_factory)
-            if token:
-                headers = databricks_request_headers(server_url, bearer_token=token)
-                update_permission_hook_auth_headers(bridge_dir, headers)
+
+            def _refresh_snapshot() -> None:
+                token = auth_token_factory()
+                if token:
+                    headers = databricks_request_headers(server_url, bearer_token=token)
+                    update_permission_hook_auth_headers(bridge_dir, headers)
+
+            # Run the whole snapshot refresh off-loop: the token fetch can
+            # shell out to the Databricks CLI, and both the header build and
+            # the bridge-file write are synchronous. Keeping them on the loop
+            # (every 60s per session) needlessly competes with the tunnel's
+            # keepalive coroutine — see _RunnerDatabricksAuth.async_auth_flow.
+            await asyncio.to_thread(_refresh_snapshot)
         except asyncio.CancelledError:
             raise
         except Exception:  # noqa: BLE001 — retain the last still-valid snapshot

tests/runner/test_runner_entry.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import asyncio
+import contextlib
 import importlib
 import io
 import logging
@@ -926,6 +927,118 @@ async def test_runner_databricks_auth_end_to_end_through_mock_transport() -> Non
     assert minted_tokens == ["tok-1", "tok-2"]
 
 
+@pytest.mark.asyncio
+async def test_async_auth_flow_remints_on_login_redirect_through_async_client() -> None:
+    """``async_auth_flow`` re-mints on a 302→/oidc/ exactly like the sync flow.
+
+    ``httpx.AsyncClient`` prefers ``async_auth_flow`` over the sync
+    ``auth_flow`` when both are defined, so this pins that the async twin
+    added to keep the event loop unblocked preserves the login-redirect
+    re-mint semantics (the ``ness-tool-spin`` regression guard, async path).
+
+    :returns: None.
+    """
+    minted_tokens: list[str] = []
+    seen_authz: list[str] = []
+
+    def _factory() -> str:
+        token = f"tok-{len(minted_tokens) + 1}"
+        minted_tokens.append(token)
+        return token
+
+    def _handler(request: httpx.Request) -> httpx.Response:
+        seen_authz.append(request.headers.get("authorization", ""))
+        if len(seen_authz) == 1:
+            return httpx.Response(
+                302,
+                headers={"Location": "https://ws.example.com/oidc/oauth2/v2.0/authorize"},
+            )
+        return httpx.Response(200, json={"ok": True})
+
+    transport = httpx.MockTransport(_handler)
+    async with httpx.AsyncClient(
+        base_url="http://ap.example.com",
+        auth=_RunnerDatabricksAuth(_factory),
+        transport=transport,
+    ) as client:
+        resp = await client.post("/v1/sessions/conv_abc/mcp", json={})
+
+    assert resp.status_code == 200
+    assert seen_authz == ["Bearer tok-1", "Bearer tok-2"]
+    assert minted_tokens == ["tok-1", "tok-2"]
+
+
+@pytest.mark.asyncio
+async def test_async_auth_flow_blocking_factory_does_not_starve_event_loop() -> None:
+    """A slow (blocking) token factory must not stall the asyncio loop.
+
+    This is the regression guard for the ``1011 keepalive ping timeout``
+    tunnel drops: ``_RunnerDatabricksAuth`` is attached to the runner's
+    ``httpx.AsyncClient``, which shares the event loop with the WebSocket
+    tunnel's keepalive coroutine. If the (synchronous, sometimes multi-second
+    — CLI shell-out / mint POST) token factory runs inline on the loop, the
+    keepalive coroutine cannot answer the server's ping and the tunnel drops.
+    ``async_auth_flow`` offloads the factory via ``asyncio.to_thread``; this
+    test proves a concurrent loop task keeps ticking *while* the blocking
+    factory runs.
+
+    :returns: None.
+    """
+    factory_entered = asyncio.Event()
+
+    def _blocking_factory() -> str:
+        # Simulate the worst case: a synchronous multi-second credential
+        # resolution (Databricks CLI shell-out / 10s mint POST). If this runs
+        # on the event loop, the concurrent heartbeat below cannot advance.
+        factory_entered.set()
+        time.sleep(0.5)
+        return "tok"
+
+    heartbeats = 0
+
+    async def _heartbeat() -> None:
+        # Stand-in for the websockets keepalive coroutine: must keep ticking
+        # on the loop even while the factory is blocking in its worker thread.
+        nonlocal heartbeats
+        while True:
+            heartbeats += 1
+            await asyncio.sleep(0.02)
+
+    def _handler(request: httpx.Request) -> httpx.Response:
+        return httpx.Response(200, json={"ok": True})
+
+    hb = asyncio.create_task(_heartbeat())
+    try:
+        transport = httpx.MockTransport(_handler)
+        async with httpx.AsyncClient(
+            base_url="http://ap.example.com",
+            auth=_RunnerDatabricksAuth(_blocking_factory),
+            transport=transport,
+        ) as client:
+            # The post must run concurrently with the wait: the factory only
+            # runs (and sets factory_entered) *during* the request, so we
+            # can't await the event before launching the post.
+            post = asyncio.create_task(client.post("/v1/sessions/conv_abc/mcp", json={}))
+            await factory_entered.wait()
+            beats_before = heartbeats
+            resp = await post
+    finally:
+        hb.cancel()
+        with contextlib.suppress(asyncio.CancelledError):
+            await hb
+
+    assert resp.status_code == 200
+    # The 0.5s blocking factory spans ~25 heartbeat ticks (20ms each). If the
+    # factory had run on the loop, the heartbeat would have been frozen for
+    # its whole duration and gained ~0 ticks. Requiring several ticks proves
+    # the loop stayed live — i.e. keepalive pings would still be answered.
+    assert heartbeats - beats_before >= 5, (
+        f"event loop starved during token resolution: only "
+        f"{heartbeats - beats_before} heartbeat tick(s) while the factory "
+        f"blocked — async_auth_flow is running the factory on the loop"
+    )
+
+
 def test_runner_tunnel_binding_token_from_env_returns_none_without_token(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None: