diff --git a/agent/agent_runtime_helpers.py b/agent/agent_runtime_helpers.py index 3017f00a10..4927938655 100644 --- a/agent/agent_runtime_helpers.py +++ b/agent/agent_runtime_helpers.py @@ -2860,6 +2860,12 @@ def create_openai_client(agent, client_kwargs: dict, *, reason: str, shared: boo # constructs a fresh one — no stale closed transport can be reused. # Tests in ``tests/run_agent/test_create_openai_client_reuse.py`` and # ``tests/run_agent/test_sequential_chats_live.py`` pin this invariant. + # What IS shared across those per-client wrappers is the underlying + # connection pool: ``build_keepalive_http_client`` mounts a + # process-shared ``HTTPTransport`` behind a per-client view whose + # ``close()`` is a no-op for the pool, so a closed wrapper never takes + # a sibling's (or the successor's) connections with it + # (tests/agent/test_shared_http_transport.py). if "http_client" not in client_kwargs: keepalive_http = agent._build_keepalive_http_client( client_kwargs.get("base_url", ""), verify=httpx_verify, @@ -4908,8 +4914,8 @@ def reapply_reasoning_echo_for_provider(agent, api_messages: list) -> int: ) -def _iter_httpx_pool_objects(http_client: Any): - """Yield httpcore pool objects reachable from an httpx client. +def _iter_httpx_pools_with_owner(http_client: Any): + """Yield ``(pool, owner)`` pairs reachable from an httpx client. Hermes' keepalive client (#10324 / ``_build_keepalive_http_client``) and any ``HTTP(S)_PROXY`` configuration put live connections on *mounted* @@ -4918,31 +4924,37 @@ def _iter_httpx_pool_objects(http_client: Any): ``force_close_tcp_sockets`` return 0 while a stream is still mid-recv — the interrupt logs success and the provider keeps burning the slot (#72975). + + ``owner`` is ``None`` for a pool this client owns outright, or the + ``_SharedTransport`` view id when the pool is process-shared with other + clients (``process_bootstrap.build_keepalive_http_client``). Callers must + then touch only the in-flight requests stamped with that owner. """ seen_pools: set[int] = set() - def _emit(pool: Any): + def _emit(pool: Any, owner: Any): if pool is None: return marker = id(pool) if marker in seen_pools: return seen_pools.add(marker) - yield pool + yield pool, owner def _pools_for_transport(transport: Any): if transport is None: return + owner = id(transport) if type(transport).__name__ == "_SharedTransport" else None # Normal httpx.HTTPTransport / HTTPProxy-as-transport: connections # live under ``_pool``. HTTPProxy itself *is* a ConnectionPool and # may be mounted directly — then ``_connections`` is on the # transport. pool = getattr(transport, "_pool", None) if pool is not None: - yield from _emit(pool) + yield from _emit(pool, owner) return if getattr(transport, "_connections", None) is not None: - yield from _emit(transport) + yield from _emit(transport, owner) try: yield from _pools_for_transport(getattr(http_client, "_transport", None)) @@ -4953,6 +4965,12 @@ def _iter_httpx_pool_objects(http_client: Any): return +def _iter_httpx_pool_objects(http_client: Any): + """Yield httpcore pool objects reachable from an httpx client.""" + for pool, _owner in _iter_httpx_pools_with_owner(http_client): + yield pool + + def _connection_candidates(conn: Any): """Walk nested ``_connection`` wrappers (proxy tunnel → HTTP11/2).""" seen: set[int] = set() @@ -4991,23 +5009,32 @@ def _iter_pool_sockets(client: Any): # Some SDK wrappers *are* the httpx client (or expose the pool # directly). Fall through so mount-aware discovery still runs. http_client = client - pools = list(_iter_httpx_pool_objects(http_client)) + pools = list(_iter_httpx_pools_with_owner(http_client)) except Exception: return if not pools: return + from agent.process_bootstrap import HERMES_TRANSPORT_OWNER_EXT + seen: set[int] = set() - for pool in pools: + for pool, owner in pools: # Empty-list is falsy: use ``is None`` so an empty ``_connections`` # still lets us walk in-flight ``_requests`` rather than skipping # the pool entirely. raw_conns = getattr(pool, "_connections", None) if raw_conns is None: raw_conns = getattr(pool, "_pool", None) - connections = list(raw_conns or []) + # A process-shared pool carries other clients' idle + in-flight + # connections: only this client's own in-flight requests (stamped by + # ``_SharedTransport.handle_request``) may be shut down. + connections = [] if owner is not None else list(raw_conns or []) for pool_req in list(getattr(pool, "_requests", None) or []): + if owner is not None: + exts = getattr(getattr(pool_req, "request", None), "extensions", None) or {} + if exts.get(HERMES_TRANSPORT_OWNER_EXT) != owner: + continue conn = getattr(pool_req, "connection", None) if conn is not None: connections.append(conn) diff --git a/agent/process_bootstrap.py b/agent/process_bootstrap.py index 341126c919..323dea85f6 100644 --- a/agent/process_bootstrap.py +++ b/agent/process_bootstrap.py @@ -31,6 +31,7 @@ import os import selectors import socket import sys +import threading import time import urllib.request from typing import Any, Optional @@ -43,6 +44,19 @@ from utils import base_url_hostname, normalize_proxy_url _OPENAI_CLS_CACHE = None _HAPPY_EYEBALLS_DELAY_SECONDS = 0.25 +# Process-wide pool of sync ``httpx.HTTPTransport`` objects shared by every +# keepalive client with the same (verify, proxy, happy-eyeballs) identity. +# Each delegated child AIAgent used to get its own transport = its own TLS +# pool, so a fan-out of N children held N separate socket sets to the same +# provider. Bounded: past the cap, callers get a private transport again. +_SHARED_TRANSPORTS: dict[tuple, Any] = {} +_SHARED_TRANSPORTS_LOCK = threading.Lock() +_SHARED_TRANSPORTS_MAX = 32 +# ``request.extensions`` key stamped by ``_SharedTransport.handle_request``; +# the socket-abort walker in agent_runtime_helpers uses it to find only the +# owning client's in-flight connections on a shared pool. +HERMES_TRANSPORT_OWNER_EXT = "hermes_transport_owner" + def _interleave_addrinfos(addrinfos: list[tuple]) -> list[tuple]: """Interleave resolved address families while preserving resolver order.""" @@ -418,6 +432,93 @@ def _get_proxy_for_base_url(base_url: Optional[str]) -> Optional[str]: return proxy +def _shared_transport_cls(): + """Lazily define the per-client transport view (httpx import is deferred).""" + global _SharedTransport + if _SharedTransport is not None: + return _SharedTransport + import httpx + + class _SharedTransportImpl(httpx.BaseTransport): + """Per-client view of a process-shared ``httpx.HTTPTransport``. + + ``httpx.Client.close()`` closes every mounted transport. Each OpenAI + client still owns its own ``httpx.Client`` (the #10933 contract: + closing one client must never poison the next), so the object we + mount must absorb that close while the underlying connection pool + keeps serving every other client. ``handle_request`` stamps the + owning view into ``request.extensions`` so socket-abort sweeps can + target only this client's in-flight connections on the shared pool. + """ + + __slots__ = ("_inner", "_closed") + + def __init__(self, inner: Any) -> None: + self._inner = inner + self._closed = False + + # httpx-private ``_pool`` is what our socket walkers and the + # happy-eyeballs / ssl-verify tests introspect: expose the shared one. + @property + def _pool(self) -> Any: + return getattr(self._inner, "_pool", None) + + def handle_request(self, request: Any) -> Any: + if self._closed: + raise RuntimeError("Cannot send a request, as the client has been closed.") + request.extensions[HERMES_TRANSPORT_OWNER_EXT] = id(self) + return self._inner.handle_request(request) + + def close(self) -> None: + # Deliberately does NOT close ``_inner``: it is shared. Idle + # connections are reaped by ``keepalive_expiry``; the pool lives + # for the process (see ``close_shared_transports``). + self._closed = True + + _SharedTransportImpl.__name__ = _SharedTransportImpl.__qualname__ = "_SharedTransport" + _SharedTransport = _SharedTransportImpl + return _SharedTransport + + +_SharedTransport: Any = None + + +def _shared_transport_key(base_url: str, verify: Any, proxy: Optional[str]) -> tuple: + """Identity under which sync direct transports are pooled process-wide.""" + if verify is True or verify is False: + verify_key: Any = verify + elif isinstance(verify, str): + verify_key = ("path", verify) + else: + # An ssl.SSLContext (or custom object): share only by object identity, + # which is what a caller passing the same context twice expects. + verify_key = ("id", id(verify)) + return (verify_key, proxy, _uses_codex_cloud_transport(base_url)) + + +def _get_shared_transport(key: tuple, build) -> Any: + with _SHARED_TRANSPORTS_LOCK: + transport = _SHARED_TRANSPORTS.get(key) + if transport is None: + transport = build() + if len(_SHARED_TRANSPORTS) < _SHARED_TRANSPORTS_MAX: + _SHARED_TRANSPORTS[key] = transport + return transport + + +def close_shared_transports() -> int: + """Really close every process-shared transport (test teardown / atexit).""" + with _SHARED_TRANSPORTS_LOCK: + transports = list(_SHARED_TRANSPORTS.values()) + _SHARED_TRANSPORTS.clear() + for transport in transports: + try: + transport.close() + except Exception: + pass + return len(transports) + + def build_keepalive_http_client( base_url: str = "", *, @@ -444,6 +545,14 @@ def build_keepalive_http_client( ``ssl_ca_cert`` / ``ssl_verify`` and ``HERMES_CA_BUNDLE`` settings the main client uses. It is passed on the client AND on the plain no-proxy mounts (a mounted transport owns the SSL context for its scheme). + + Every call returns a NEW ``httpx.Client`` (per-client close semantics are + what #10933 pins), but sync clients with the same + (verify, proxy, happy-eyeballs) identity mount the SAME underlying + ``HTTPTransport`` through a :class:`_SharedTransport` view, so N delegated + children share one connection pool + SSL context instead of N. Async + clients are never shared: an httpcore async pool is bound to the event + loop that first used it. """ try: import httpx @@ -462,16 +571,47 @@ def build_keepalive_http_client( client_cls = httpx.AsyncClient if async_mode else httpx.Client mounts = {} if proxy is None: - http_transport = transport_cls(verify=verify) - https_transport = transport_cls(verify=verify) - # Async transports need no explicit racing: httpcore's anyio - # backend already implements RFC 8305 natively - # (``anyio.connect_tcp(happy_eyeballs_delay=0.25)``), covered by - # tests/agent/test_codex_happy_eyeballs.py. - if not async_mode and _uses_codex_cloud_transport(base_url): - _enable_happy_eyeballs(http_transport) - _enable_happy_eyeballs(https_transport) - mounts = {"http://": http_transport, "https://": https_transport} + happy_eyeballs = not async_mode and _uses_codex_cloud_transport(base_url) + # One pool now serves every agent in the process, so its ceiling + # must cover a whole fan-out of concurrently streaming children, + # not one client. (Note: previously the mounts silently ran on + # httpx defaults — keepalive_expiry=5s — since Client-level + # ``limits`` only reach the default transport.) + direct_limits = limits if async_mode else httpx.Limits( + max_keepalive_connections=50, + max_connections=1000, + keepalive_expiry=20.0, + ) + + def _build_direct(): + transport = transport_cls(verify=verify, limits=direct_limits) + # Async transports need no explicit racing: httpcore's anyio + # backend already implements RFC 8305 natively + # (``anyio.connect_tcp(happy_eyeballs_delay=0.25)``), covered + # by tests/agent/test_codex_happy_eyeballs.py. + if happy_eyeballs: + _enable_happy_eyeballs(transport) + return transport + + if async_mode: + mounts = {"http://": _build_direct(), "https://": _build_direct()} + else: + key = _shared_transport_key(base_url, verify, proxy) + view_cls = _shared_transport_cls() + mounts = { + f"{scheme}://": view_cls( + _get_shared_transport((scheme, *key), _build_direct) + ) + for scheme in ("http", "https") + } + # Without this httpx builds a third, never-used direct + # transport (and pool + SSL context) per client. + return client_cls( + limits=limits, + timeout=timeout, + transport=mounts["https://"], + mounts=mounts, + ) return client_cls( limits=limits, timeout=timeout, @@ -506,5 +646,6 @@ __all__ = [ "_get_proxy_from_env", "_get_proxy_for_base_url", "build_keepalive_http_client", + "close_shared_transports", "enable_happy_eyeballs_on_client", ] diff --git a/agent/ssl_verify.py b/agent/ssl_verify.py index 885702185d..c62c6c81da 100644 --- a/agent/ssl_verify.py +++ b/agent/ssl_verify.py @@ -5,6 +5,7 @@ from __future__ import annotations import logging import os import ssl +import threading from pathlib import Path from typing import Any, Optional @@ -19,6 +20,27 @@ def _coerce_insecure(ssl_verify: Any) -> bool: return False +_CA_CONTEXTS: dict[str, ssl.SSLContext] = {} +_CA_CONTEXTS_LOCK = threading.Lock() + + +def _context_for_ca_bundle(ca_path: str) -> ssl.SSLContext: + """One ``SSLContext`` per CA bundle path, process-wide. + + ``ssl.create_default_context(cafile=...)`` parses the whole bundle each + call. Every AIAgent (and every delegated child) resolves verify for its + own client, so an env/config CA bundle used to cost one parsed context — + and, because sharing keys on context identity, one private connection + pool — per agent. An ``SSLContext`` is safe to share across connections. + """ + with _CA_CONTEXTS_LOCK: + ctx = _CA_CONTEXTS.get(ca_path) + if ctx is None: + ctx = ssl.create_default_context(cafile=ca_path) + _CA_CONTEXTS[ca_path] = ctx + return ctx + + def resolve_httpx_verify( *, ca_bundle: Optional[str] = None, @@ -55,7 +77,7 @@ def resolve_httpx_verify( if effective_ca: ca_path = str(Path(effective_ca).expanduser()) if os.path.isfile(ca_path): - return ssl.create_default_context(cafile=ca_path) + return _context_for_ca_bundle(ca_path) logger.warning( "CA bundle path does not exist: %s — falling back to default certificates", effective_ca, diff --git a/tests/agent/test_shared_http_transport.py b/tests/agent/test_shared_http_transport.py new file mode 100644 index 0000000000..d0c541723a --- /dev/null +++ b/tests/agent/test_shared_http_transport.py @@ -0,0 +1,185 @@ +"""Keepalive httpx clients share one HTTPTransport per (verify, proxy) identity. + +Every AIAgent (and every delegated child) gets its own ``httpx.Client`` — the +#10933 contract that closing one client must never poison the next. What is +shared underneath is the connection pool + SSL context, so a fan-out of N +children no longer holds N TLS socket sets to the same provider. +""" + +import ssl +import threading +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + +import certifi +import httpx +import pytest + +from agent import process_bootstrap +from agent.agent_runtime_helpers import _iter_pool_sockets, force_close_tcp_sockets +from agent.process_bootstrap import build_keepalive_http_client + + +@pytest.fixture +def no_proxy_env(monkeypatch): + for name in ( + "HTTPS_PROXY", "HTTP_PROXY", "ALL_PROXY", + "https_proxy", "http_proxy", "all_proxy", "NO_PROXY", "no_proxy", + ): + monkeypatch.delenv(name, raising=False) + process_bootstrap.close_shared_transports() + yield + process_bootstrap.close_shared_transports() + + +class _Handler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" # keep-alive so pooled connections persist + + def do_GET(self): # noqa: N802 + body = b"ok" + self.send_response(200) + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, *_args): + pass + + +@pytest.fixture +def local_server(): + server = ThreadingHTTPServer(("127.0.0.1", 0), _Handler) + server.daemon_threads = True + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{server.server_address[1]}" + finally: + server.shutdown() + server.server_close() + + +def _inner(client, scheme="https://"): + mount = next(t for pat, t in client._mounts.items() if str(pat.pattern) == scheme) + return mount._inner + + +def test_same_identity_clients_share_transport_but_not_client(no_proxy_env): + a = build_keepalive_http_client("https://api.example.com/v1") + b = build_keepalive_http_client("https://api.example.com/v1") + assert isinstance(a, httpx.Client) and isinstance(b, httpx.Client) + assert a is not b + assert _inner(a) is _inner(b) + assert _inner(a, "http://") is _inner(b, "http://") + # The per-client view is distinct, so each client has its own close state. + assert a._mounts is not b._mounts + a.close() + b.close() + + +def test_closing_one_client_leaves_sibling_functional(no_proxy_env, local_server): + a = build_keepalive_http_client(local_server) + b = build_keepalive_http_client(local_server) + assert _inner(a, "http://") is _inner(b, "http://") + assert a.get(local_server + "/x").status_code == 200 + a.close() + assert a.is_closed + # #10933 shape: the shared pool must still serve the surviving client and + # any successor client built after the close. + assert b.get(local_server + "/y").status_code == 200 + c = build_keepalive_http_client(local_server) + assert _inner(c, "http://") is _inner(b, "http://") + assert c.get(local_server + "/z").status_code == 200 + with pytest.raises(RuntimeError): + a.get(local_server + "/closed") + b.close() + c.close() + + +def test_pool_survives_client_close(no_proxy_env, local_server): + a = build_keepalive_http_client(local_server) + a.get(local_server + "/warm") + pool = _inner(a, "http://")._pool + before = len(pool.connections) + assert before >= 1 + a.close() + assert len(pool.connections) == before, "client close must not drain the shared pool" + + +def test_different_verify_or_proxy_get_different_transports(no_proxy_env, monkeypatch): + default = build_keepalive_http_client("https://api.example.com/v1") + insecure = build_keepalive_http_client("https://api.example.com/v1", verify=False) + ctx = ssl.create_default_context(cafile=certifi.where()) + with_ctx = build_keepalive_http_client("https://api.example.com/v1", verify=ctx) + with_ctx2 = build_keepalive_http_client("https://api.example.com/v1", verify=ctx) + codex = build_keepalive_http_client("https://chatgpt.com/backend-api/codex") + assert _inner(default) is not _inner(insecure) + assert _inner(default) is not _inner(with_ctx) + assert _inner(with_ctx) is _inner(with_ctx2) + assert _inner(with_ctx)._pool._ssl_context is ctx + assert _inner(insecure)._pool._ssl_context.check_hostname is False + # Codex cloud gets the happy-eyeballs backend, so it can't share a pool. + assert _inner(codex) is not _inner(default) + assert isinstance( + _inner(codex)._pool._network_backend, process_bootstrap._HappyEyeballsSyncBackend + ) + for c in (default, insecure, with_ctx, with_ctx2, codex): + c.close() + + monkeypatch.setenv("HTTPS_PROXY", "http://127.0.0.1:3128") + proxied = build_keepalive_http_client("https://api.example.com/v1") + # Proxy clients keep httpx's own per-client proxy transport (unshared). + assert all( + type(t).__name__ != "_SharedTransport" for t in proxied._mounts.values() if t + ) + proxied.close() + + +def test_async_clients_are_not_shared(no_proxy_env): + a = build_keepalive_http_client("https://api.example.com/v1", async_mode=True) + b = build_keepalive_http_client("https://api.example.com/v1", async_mode=True) + assert isinstance(a, httpx.AsyncClient) + ta = [t for t in a._mounts.values() if t is not None] + tb = [t for t in b._mounts.values() if t is not None] + assert all(isinstance(t, httpx.AsyncHTTPTransport) for t in ta + tb) + assert not {id(t) for t in ta} & {id(t) for t in tb} + + +def test_force_close_only_touches_owning_clients_inflight_sockets(no_proxy_env, local_server): + """A stranger-thread abort on client A must not shut down client B's + idle/in-flight connections that live on the same shared pool.""" + a = build_keepalive_http_client(local_server) + b = build_keepalive_http_client(local_server) + b.get(local_server + "/warm") # idle keepalive connection on the shared pool + pool = _inner(a, "http://")._pool + assert pool.connections + # A has nothing in flight: nothing of A's may be touched. + assert list(_iter_pool_sockets(a)) == [] + assert force_close_tcp_sockets(a) == 0 + # B's idle connection is still healthy. + assert b.get(local_server + "/again").status_code == 200 + + # Now hold a B stream open and confirm A's abort still sees zero sockets + # while B's abort sees exactly its own. + with b.stream("GET", local_server + "/stream") as resp: + assert resp.status_code == 200 + assert list(_iter_pool_sockets(a)) == [] + assert len(list(_iter_pool_sockets(b))) == 1 + a.close() + b.close() + + +def test_shared_transport_cache_is_bounded(no_proxy_env, monkeypatch): + monkeypatch.setattr(process_bootstrap, "_SHARED_TRANSPORTS_MAX", 2) + clients = [ + build_keepalive_http_client("https://api.example.com/v1", verify=False), + build_keepalive_http_client("https://api.example.com/v1"), + ] + assert len(process_bootstrap._SHARED_TRANSPORTS) == 2 + ctx = ssl.create_default_context() + extra = build_keepalive_http_client("https://api.example.com/v1", verify=ctx) + assert len(process_bootstrap._SHARED_TRANSPORTS) == 2 + # Past the cap the caller still gets a working (private) transport. + assert _inner(extra)._pool._ssl_context is ctx + for c in clients + [extra]: + c.close() + assert process_bootstrap.close_shared_transports() == 2