From e9eb1682da5b09e1a7a7fb9b76bc71c83de4742e Mon Sep 17 00:00:00 2001 From: Lingrui Mei Date: Sat, 4 Jul 2026 08:14:17 +0800 Subject: [PATCH] abridge: configurable tunnel request window + explicit SDK retry policy MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The tunnel answered every agent request within a hard-coded 600s window (_start_tunnel already took request_timeout; Proxy.start never passed it). The window's timer starts when the request reaches the sandbox — strictly before the host's upstream call — so any upstream call approaching 600s lost the race: the agent got a 504 while the host handler kept running, and the SDK's silent default max_retries=2 let the handler occupy up to ~3x the window on a result nobody was waiting for. Long-generation models hit this in practice; the only workaround was capping the model's time budget to fit the infra. - Proxy(*clients, request_timeout=600.0) threads the window into the in-sandbox tunnel. Sizing rule documented as a LOWER bound: client timeout x (1 + max_retries), with SDK backoff and Retry-After waits on top — leave real margin. Non-positive values rejected. - AnthropicFromOpenAIClient, OpenAIClient, and AnthropicClient take max_retries and default it to 0 — retry policy is the operator's call, never a hidden multiplier behind the tunnel window. agentix-bridge-serve grows --upstream-max-retries so CLI operators can opt back in. - Forward's default deadline drops to 540s, strictly under the default window it races (an equal deadline always loses — the tunnel timer starts first). - Tests include a real end-to-end tunnel: a never-answering host 504s at ~request_timeout, so a regression re-hardcoding the window during route registration blows the time budget instead of staying green. Co-Authored-By: Claude Fable 5 --- .../agentix/bridge/clients/anthropic.py | 7 +- .../bridge/clients/anthropic_from_openai.py | 9 +- .../abridge/agentix/bridge/clients/openai.py | 7 +- plugins/abridge/agentix/bridge/forward.py | 12 +- plugins/abridge/agentix/bridge/proxy.py | 18 ++- plugins/abridge/agentix/bridge/serve.py | 7 + plugins/abridge/tests/test_tunnel_timeout.py | 135 ++++++++++++++++++ 7 files changed, 188 insertions(+), 7 deletions(-) create mode 100644 plugins/abridge/tests/test_tunnel_timeout.py diff --git a/plugins/abridge/agentix/bridge/clients/anthropic.py b/plugins/abridge/agentix/bridge/clients/anthropic.py index 16a3142..8140b56 100644 --- a/plugins/abridge/agentix/bridge/clients/anthropic.py +++ b/plugins/abridge/agentix/bridge/clients/anthropic.py @@ -85,11 +85,16 @@ def __init__( api_key: str | None = None, model: str | None = None, timeout: float = 120.0, + max_retries: int = 0, session_id: str | None = None, ) -> None: if AsyncAnthropic is None: raise ImportError(_INSTALL_HINT) - self._client = AsyncAnthropic(base_url=base_url, api_key=api_key, timeout=timeout) + # 0, not the SDK's 2: silent retries multiply tunnel-window occupancy; + # retry policy is the operator's call (see Proxy.request_timeout). + self._client = AsyncAnthropic( + base_url=base_url, api_key=api_key, timeout=timeout, max_retries=max_retries + ) self._model = model self.session_id = session_id or uuid.uuid4().hex diff --git a/plugins/abridge/agentix/bridge/clients/anthropic_from_openai.py b/plugins/abridge/agentix/bridge/clients/anthropic_from_openai.py index 64f95a3..7356419 100644 --- a/plugins/abridge/agentix/bridge/clients/anthropic_from_openai.py +++ b/plugins/abridge/agentix/bridge/clients/anthropic_from_openai.py @@ -82,12 +82,19 @@ def __init__( api_key: str | None = None, model: str | None = None, timeout: float = 120.0, + max_retries: int = 0, session_id: str | None = None, upstream_params: dict[str, Any] | None = None, ) -> None: if AsyncOpenAI is None: raise ImportError(_INSTALL_HINT) - self._client = AsyncOpenAI(base_url=base_url, api_key=api_key, timeout=timeout) + # max_retries defaults to 0, not the SDK's 2: silent retries multiply + # the handler's occupancy of the tunnel window (Proxy.request_timeout) + # behind the operator's back. Opt in explicitly and size the window to + # timeout x (1 + max_retries). + self._client = AsyncOpenAI( + base_url=base_url, api_key=api_key, timeout=timeout, max_retries=max_retries + ) self._model = model self._upstream_params = dict(upstream_params or {}) self.session_id = session_id or uuid.uuid4().hex diff --git a/plugins/abridge/agentix/bridge/clients/openai.py b/plugins/abridge/agentix/bridge/clients/openai.py index b906882..f2e9ca6 100644 --- a/plugins/abridge/agentix/bridge/clients/openai.py +++ b/plugins/abridge/agentix/bridge/clients/openai.py @@ -77,11 +77,16 @@ def __init__( model: str | None = None, extra_body: dict[str, Any] | None = None, timeout: float = 120.0, + max_retries: int = 0, session_id: str | None = None, ) -> None: if AsyncOpenAI is None: raise ImportError(_INSTALL_HINT) - self._sdk = AsyncOpenAI(base_url=base_url, api_key=api_key, timeout=timeout) + # 0, not the SDK's 2: silent retries multiply tunnel-window occupancy; + # retry policy is the operator's call (see Proxy.request_timeout). + self._sdk = AsyncOpenAI( + base_url=base_url, api_key=api_key, timeout=timeout, max_retries=max_retries + ) self._model = model self._extra_body = extra_body or {} self.session_id = session_id or uuid.uuid4().hex diff --git a/plugins/abridge/agentix/bridge/forward.py b/plugins/abridge/agentix/bridge/forward.py index ff3299d..e65a979 100644 --- a/plugins/abridge/agentix/bridge/forward.py +++ b/plugins/abridge/agentix/bridge/forward.py @@ -86,7 +86,11 @@ def __init__( target_url: str, *, paths: list[str], - timeout: float = 600.0, + # Strictly under Proxy's default 600s tunnel window: the window's + # timer starts before this forward does, so an equal deadline always + # loses the race and turns a slow-but-successful sidecar call into a + # tunnel 504. + timeout: float = 540.0, session_id: str | None = None, headers: Mapping[str, str] | None = None, ) -> None: @@ -214,7 +218,11 @@ def __init__( paths: list[str], create_path: str = "/sessions", session_id_field: str = "session_id", - timeout: float = 600.0, + # Strictly under Proxy's default 600s tunnel window: the window's + # timer starts before this forward does, so an equal deadline always + # loses the race and turns a slow-but-successful sidecar call into a + # tunnel 504. + timeout: float = 540.0, headers: Mapping[str, str] | None = None, ) -> None: super().__init__(target_url, paths=paths, timeout=timeout, headers=headers) diff --git a/plugins/abridge/agentix/bridge/proxy.py b/plugins/abridge/agentix/bridge/proxy.py index 03f8ada..52c6c99 100644 --- a/plugins/abridge/agentix/bridge/proxy.py +++ b/plugins/abridge/agentix/bridge/proxy.py @@ -490,10 +490,20 @@ class MyClient(SomeHandlerMixin, AnotherHandlerMixin): grouping; the bundled clients populate it via `populate_*_span`. """ - def __init__(self, *clients: Client) -> None: + def __init__(self, *clients: Client, request_timeout: float = 600.0) -> None: super().__init__(NAMESPACE) if not clients: raise ValueError("Proxy requires at least one client with @on-decorated handlers") + if request_timeout <= 0: + raise ValueError(f"request_timeout must be > 0, got {request_timeout!r}") + # Per-request tunnel window. Its timer starts when the agent's request + # reaches the sandbox — strictly BEFORE the host's upstream call — so + # size it to cover the slowest client's worst case. The lower bound is + # client timeout x (1 + client max_retries); SDK retry backoff and + # Retry-After waits (up to 60s each) sit ON TOP of that, so leave real + # margin, and note an explicit Sandbox call_deadline can still end the + # whole operation earlier. + self._request_timeout = float(request_timeout) self._handle: TunnelHandle | None = None self._clients = clients self._clients_closed = False @@ -632,7 +642,11 @@ async def start(self, sandbox: Sandbox) -> TunnelHandle: "sandbox.remote()/health() call — its /abridge namespace must be " "registered before the runtime client connects" ) from exc - handle = await sandbox.remote(_start_tunnel, paths=list(self.paths)) + handle = await sandbox.remote( + _start_tunnel, + paths=list(self.paths), + request_timeout=self._request_timeout, + ) except BaseException: try: await self._close_clients() diff --git a/plugins/abridge/agentix/bridge/serve.py b/plugins/abridge/agentix/bridge/serve.py index 3ffc789..3c212f2 100644 --- a/plugins/abridge/agentix/bridge/serve.py +++ b/plugins/abridge/agentix/bridge/serve.py @@ -351,6 +351,12 @@ def main(argv: list[str] | None = None) -> None: parser.add_argument("--host", default="127.0.0.1", help="bind address; expose beyond loopback deliberately") parser.add_argument("--port", type=int, default=8399) parser.add_argument("--upstream-timeout", type=float, default=180.0) + parser.add_argument( + "--upstream-max-retries", + type=int, + default=0, + help="openai SDK retries per upstream call; keep total occupancy = timeout x (1+retries) explicit", + ) parser.add_argument("--max-sessions", type=int, default=256) args = parser.parse_args(argv) if not args.upstream_base_url: @@ -365,6 +371,7 @@ def factory(session_id: str) -> AnthropicFromOpenAIClient: api_key=args.upstream_api_key, model=args.upstream_model, timeout=args.upstream_timeout, + max_retries=args.upstream_max_retries, session_id=session_id, ) diff --git a/plugins/abridge/tests/test_tunnel_timeout.py b/plugins/abridge/tests/test_tunnel_timeout.py new file mode 100644 index 0000000..5885a0e --- /dev/null +++ b/plugins/abridge/tests/test_tunnel_timeout.py @@ -0,0 +1,135 @@ +"""Tunnel request-window configuration + upstream SDK retry policy. + +The tunnel answers each agent request within `request_timeout`; the timer +starts before the host's upstream call does, so the window must strictly +cover the host client's worst case (timeout x (1 + max_retries)). These +tests pin the two halves: `Proxy` threads a configurable window into the +sandbox tunnel, and the bundled clients don't retry behind the operator's +back. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +import pytest +from agentix.bridge import Proxy, on +from agentix.bridge.clients import AnthropicFromOpenAIClient, OpenAIClient +from agentix.bridge.proxy import TunnelHandle, _make_forwarder, _start_tunnel +from starlette.requests import Request as StarletteRequest + + +class _EchoClient: + @on("/v1/echo") + async def echo(self, request: Any) -> Any: + raise AssertionError("never dispatched in these tests") + + +@dataclass +class _FakeSandbox: + remote_calls: list[tuple[Any, dict[str, Any]]] = field(default_factory=list) + + def register_namespace(self, ns: Any) -> None: + pass + + async def remote(self, fn: Any, **kwargs: Any) -> Any: + self.remote_calls.append((fn, kwargs)) + return TunnelHandle(url="http://127.0.0.1:1", port=1) + + +@pytest.mark.asyncio +async def test_proxy_start_threads_request_timeout_into_the_tunnel() -> None: + sandbox = _FakeSandbox() + proxy = Proxy(_EchoClient(), request_timeout=1234.5) + await proxy.start(sandbox) # type: ignore[arg-type] + + fn, kwargs = sandbox.remote_calls[0] + assert fn is _start_tunnel + assert kwargs["request_timeout"] == 1234.5 + + +@pytest.mark.asyncio +async def test_proxy_default_window_matches_tunnel_default() -> None: + sandbox = _FakeSandbox() + await Proxy(_EchoClient()).start(sandbox) # type: ignore[arg-type] + assert sandbox.remote_calls[0][1]["request_timeout"] == 600.0 + + +def test_proxy_rejects_nonpositive_window() -> None: + with pytest.raises(ValueError, match="request_timeout"): + Proxy(_EchoClient(), request_timeout=0) + + +@pytest.mark.asyncio +async def test_forwarder_times_out_to_504_after_the_configured_window() -> None: + async def _receive() -> dict[str, Any]: + return {"type": "http.request", "body": b"{}", "more_body": False} + + # The real ns.request raises TimeoutError when the window elapses — + # emulate that while capturing the window the forwarder threaded in. + class _TimingOutNs: + def __init__(self) -> None: + self.seen_timeout: float | None = None + + async def request(self, path: str, body: Any, *, timeout: float) -> Any: + self.seen_timeout = timeout + raise TimeoutError + + ns = _TimingOutNs() + forward = _make_forwarder(ns=ns, path="/v1/echo", request_timeout=0.05) # type: ignore[arg-type] + request = StarletteRequest( + scope={"type": "http", "method": "POST", "headers": []}, receive=_receive + ) + response = await forward(request) + assert response.status_code == 504 + assert ns.seen_timeout == 0.05 + + +@pytest.mark.asyncio +async def test_real_tunnel_returns_504_within_the_configured_window(monkeypatch) -> None: + """End-to-end through a real in-sandbox tunnel: a host that never replies + must 504 at ~request_timeout, not the 600s default — a regression that + re-hardcodes the window during route registration would blow the budget + here instead of staying green.""" + import time + + import httpx + from agentix.bridge import proxy as proxy_mod + + monkeypatch.setattr(proxy_mod._get_namespace(), "emit", lambda *a, **k: _never()) + handle = await proxy_mod._start_tunnel(paths=["/v1/echo"], request_timeout=0.3) + try: + started = time.monotonic() + async with httpx.AsyncClient() as client: + resp = await client.post(f"{handle.url}/v1/echo", json={}, timeout=5.0) + elapsed = time.monotonic() - started + assert resp.status_code == 504 + assert elapsed < 3.0 # bounded by request_timeout, nowhere near 600s + finally: + await proxy_mod._stop_tunnel(handle=handle) + + +async def _never() -> None: + """An emit that drops the message, so the host reply never arrives.""" + + +def test_bundled_clients_do_not_retry_behind_the_operators_back() -> None: + """openai SDK default max_retries=2 silently multiplies handler occupancy + behind the tunnel window; retry policy must be explicit.""" + anthropic = AnthropicFromOpenAIClient(base_url="http://up.example/v1", api_key="sk-test") + assert anthropic._client.max_retries == 0 # noqa: SLF001 + + openai_client = OpenAIClient(base_url="http://up.example/v1", api_key="sk-test") + assert openai_client._sdk.max_retries == 0 # noqa: SLF001 + + pure_anthropic = pytest.importorskip("agentix.bridge.clients.anthropic") + passthrough = pure_anthropic.AnthropicClient( + base_url="http://up.example", api_key="sk-ant-test" + ) + assert passthrough._client.max_retries == 0 # noqa: SLF001 + + tuned = AnthropicFromOpenAIClient( + base_url="http://up.example/v1", api_key="sk-test", max_retries=2 + ) + assert tuned._client.max_retries == 2 # noqa: SLF001