fix(gemini): tighten native routing and streaming replay

- only use the native adapter for the canonical Gemini native endpoint
- keep custom and /openai base URLs on the OpenAI-compatible path
- preserve Hermes keepalive transport injection for native Gemini clients
- stabilize streaming tool-call replay across repeated SSE events
- add follow-up tests for base_url precedence, async streaming, and duplicate tool-call chunks
This commit is contained in:
kshitijk4poor
2026-04-20 00:41:20 +05:30
committed by Teknium
parent 3dea497b20
commit d393104bad
7 changed files with 225 additions and 56 deletions

View File

@@ -4705,6 +4705,30 @@ class AIAgent:
return bool(getattr(http_client, "is_closed", False))
return False
@staticmethod
def _build_keepalive_http_client() -> Any:
try:
import httpx as _httpx
import socket as _socket
_sock_opts = [(_socket.SOL_SOCKET, _socket.SO_KEEPALIVE, 1)]
if hasattr(_socket, "TCP_KEEPIDLE"):
_sock_opts.append((_socket.IPPROTO_TCP, _socket.TCP_KEEPIDLE, 30))
_sock_opts.append((_socket.IPPROTO_TCP, _socket.TCP_KEEPINTVL, 10))
_sock_opts.append((_socket.IPPROTO_TCP, _socket.TCP_KEEPCNT, 3))
elif hasattr(_socket, "TCP_KEEPALIVE"):
_sock_opts.append((_socket.IPPROTO_TCP, _socket.TCP_KEEPALIVE, 30))
# When a custom transport is provided, httpx won't auto-read proxy
# from env vars (allow_env_proxies = trust_env and transport is None).
# Explicitly read proxy settings to ensure HTTP_PROXY/HTTPS_PROXY work.
_proxy = _get_proxy_from_env()
return _httpx.Client(
transport=_httpx.HTTPTransport(socket_options=_sock_opts),
proxy=_proxy,
)
except Exception:
return None
def _create_openai_client(self, client_kwargs: dict, *, reason: str, shared: bool) -> Any:
from agent.auxiliary_client import _validate_base_url, _validate_proxy_env_urls
# Treat client_kwargs as read-only. Callers pass self._client_kwargs (or shallow
@@ -4746,20 +4770,26 @@ class AIAgent:
)
return client
if self.provider == "gemini":
from agent.gemini_native_adapter import GeminiNativeClient
from agent.gemini_native_adapter import GeminiNativeClient, is_native_gemini_base_url
safe_kwargs = {
k: v for k, v in client_kwargs.items()
if k in {"api_key", "base_url", "default_headers", "timeout"}
}
client = GeminiNativeClient(**safe_kwargs)
logger.info(
"Gemini native client created (%s, shared=%s) %s",
reason,
shared,
self._client_log_context(),
)
return client
base_url = str(client_kwargs.get("base_url", "") or "")
if is_native_gemini_base_url(base_url):
safe_kwargs = {
k: v for k, v in client_kwargs.items()
if k in {"api_key", "base_url", "default_headers", "timeout", "http_client"}
}
if "http_client" not in safe_kwargs:
keepalive_http = self._build_keepalive_http_client()
if keepalive_http is not None:
safe_kwargs["http_client"] = keepalive_http
client = GeminiNativeClient(**safe_kwargs)
logger.info(
"Gemini native client created (%s, shared=%s) %s",
reason,
shared,
self._client_log_context(),
)
return client
# Inject TCP keepalives so the kernel detects dead provider connections
# instead of letting them sit silently in CLOSE-WAIT (#10324). Without
# this, a peer that drops mid-stream leaves the socket in a state where
@@ -4778,28 +4808,9 @@ class AIAgent:
# Tests in ``tests/run_agent/test_create_openai_client_reuse.py`` and
# ``tests/run_agent/test_sequential_chats_live.py`` pin this invariant.
if "http_client" not in client_kwargs:
try:
import httpx as _httpx
import socket as _socket
_sock_opts = [(_socket.SOL_SOCKET, _socket.SO_KEEPALIVE, 1)]
if hasattr(_socket, "TCP_KEEPIDLE"):
# Linux
_sock_opts.append((_socket.IPPROTO_TCP, _socket.TCP_KEEPIDLE, 30))
_sock_opts.append((_socket.IPPROTO_TCP, _socket.TCP_KEEPINTVL, 10))
_sock_opts.append((_socket.IPPROTO_TCP, _socket.TCP_KEEPCNT, 3))
elif hasattr(_socket, "TCP_KEEPALIVE"):
# macOS (uses TCP_KEEPALIVE instead of TCP_KEEPIDLE)
_sock_opts.append((_socket.IPPROTO_TCP, _socket.TCP_KEEPALIVE, 30))
# When a custom transport is provided, httpx won't auto-read proxy
# from env vars (allow_env_proxies = trust_env and transport is None).
# Explicitly read proxy settings to ensure HTTP_PROXY/HTTPS_PROXY work.
_proxy = _get_proxy_from_env()
client_kwargs["http_client"] = _httpx.Client(
transport=_httpx.HTTPTransport(socket_options=_sock_opts),
proxy=_proxy,
)
except Exception:
pass # Fall through to default transport if socket opts fail
keepalive_http = self._build_keepalive_http_client()
if keepalive_http is not None:
client_kwargs["http_client"] = keepalive_http
client = OpenAI(**client_kwargs)
logger.info(
"OpenAI client created (%s, shared=%s) %s",