mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-08-31 00:03:01 +03:00
fix(codex): reuse TLS contexts across requests
This commit is contained in:
@@ -7,6 +7,7 @@ from __future__ import annotations
|
|||||||
import asyncio
|
import asyncio
|
||||||
import hashlib
|
import hashlib
|
||||||
import json
|
import json
|
||||||
|
import ssl
|
||||||
from collections.abc import Awaitable, Callable
|
from collections.abc import Awaitable, Callable
|
||||||
from typing import Any, cast
|
from typing import Any, cast
|
||||||
|
|
||||||
@@ -56,6 +57,18 @@ class OpenAICodexProvider(LLMProvider):
|
|||||||
self.proxy = proxy or None
|
self.proxy = proxy or None
|
||||||
self._extra_body = dict(extra_body or {})
|
self._extra_body = dict(extra_body or {})
|
||||||
self._native_compaction_available = True
|
self._native_compaction_available = True
|
||||||
|
self._ssl_contexts: dict[bool, ssl.SSLContext] = {}
|
||||||
|
|
||||||
|
def _ssl_context(self, *, verify: bool) -> ssl.SSLContext:
|
||||||
|
"""Reuse synchronous TLS setup across requests on the shared event loop."""
|
||||||
|
context = self._ssl_contexts.get(verify)
|
||||||
|
if context is None:
|
||||||
|
context = httpx.create_ssl_context(
|
||||||
|
verify=verify,
|
||||||
|
trust_env=self.proxy is None,
|
||||||
|
)
|
||||||
|
self._ssl_contexts[verify] = context
|
||||||
|
return context
|
||||||
|
|
||||||
async def _call_codex(
|
async def _call_codex(
|
||||||
self,
|
self,
|
||||||
@@ -129,7 +142,7 @@ class OpenAICodexProvider(LLMProvider):
|
|||||||
DEFAULT_CODEX_URL,
|
DEFAULT_CODEX_URL,
|
||||||
headers,
|
headers,
|
||||||
wire_body,
|
wire_body,
|
||||||
verify=True,
|
verify=self._ssl_context(verify=True),
|
||||||
proxy=self.proxy,
|
proxy=self.proxy,
|
||||||
on_content_delta=on_content_delta if emit_deltas else None,
|
on_content_delta=on_content_delta if emit_deltas else None,
|
||||||
on_thinking_delta=on_thinking_delta if emit_deltas else None,
|
on_thinking_delta=on_thinking_delta if emit_deltas else None,
|
||||||
@@ -145,7 +158,7 @@ class OpenAICodexProvider(LLMProvider):
|
|||||||
DEFAULT_CODEX_URL,
|
DEFAULT_CODEX_URL,
|
||||||
headers,
|
headers,
|
||||||
wire_body,
|
wire_body,
|
||||||
verify=False,
|
verify=self._ssl_context(verify=False),
|
||||||
proxy=self.proxy,
|
proxy=self.proxy,
|
||||||
on_content_delta=on_content_delta if emit_deltas else None,
|
on_content_delta=on_content_delta if emit_deltas else None,
|
||||||
on_thinking_delta=on_thinking_delta if emit_deltas else None,
|
on_thinking_delta=on_thinking_delta if emit_deltas else None,
|
||||||
@@ -411,7 +424,7 @@ async def _request_codex(
|
|||||||
url: str,
|
url: str,
|
||||||
headers: dict[str, str],
|
headers: dict[str, str],
|
||||||
body: dict[str, Any],
|
body: dict[str, Any],
|
||||||
verify: bool,
|
verify: ssl.SSLContext | bool,
|
||||||
proxy: str | None = None,
|
proxy: str | None = None,
|
||||||
on_content_delta: Callable[[str], Awaitable[None]] | None = None,
|
on_content_delta: Callable[[str], Awaitable[None]] | None = None,
|
||||||
on_thinking_delta: Callable[[str], Awaitable[None]] | None = None,
|
on_thinking_delta: Callable[[str], Awaitable[None]] | None = None,
|
||||||
|
|||||||
@@ -1,6 +1,8 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
import io
|
import io
|
||||||
|
import ssl
|
||||||
from types import SimpleNamespace
|
from types import SimpleNamespace
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
@@ -42,6 +44,46 @@ def test_codex_default_model_matches_curated_flagship() -> None:
|
|||||||
assert OpenAICodexProvider().get_default_model() == spec.builtin_models[0].id
|
assert OpenAICodexProvider().get_default_model() == spec.builtin_models[0].id
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_codex_provider_reuses_tls_context_for_concurrent_requests(monkeypatch) -> None:
|
||||||
|
_mock_codex_token(monkeypatch)
|
||||||
|
proxy = "http://127.0.0.1:23458"
|
||||||
|
context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
|
||||||
|
context_calls: list[tuple[bool, bool]] = []
|
||||||
|
request_contexts: list[object] = []
|
||||||
|
|
||||||
|
def fake_create_ssl_context(
|
||||||
|
*,
|
||||||
|
verify: bool,
|
||||||
|
cert: object = None,
|
||||||
|
trust_env: bool = True,
|
||||||
|
) -> ssl.SSLContext:
|
||||||
|
_ = cert
|
||||||
|
context_calls.append((verify, trust_env))
|
||||||
|
return context
|
||||||
|
|
||||||
|
async def fake_request(_url, _headers, _body, *, verify, **_kwargs):
|
||||||
|
request_contexts.append(verify)
|
||||||
|
await asyncio.sleep(0)
|
||||||
|
return provider_base.LLMResponse(content="ok")
|
||||||
|
|
||||||
|
monkeypatch.setattr(
|
||||||
|
"nanobot.providers.openai_codex_provider.httpx.create_ssl_context",
|
||||||
|
fake_create_ssl_context,
|
||||||
|
)
|
||||||
|
monkeypatch.setattr("nanobot.providers.openai_codex_provider._request_codex", fake_request)
|
||||||
|
|
||||||
|
provider = OpenAICodexProvider(proxy=proxy)
|
||||||
|
responses = await asyncio.gather(*(
|
||||||
|
provider.chat([{"role": "user", "content": f"request {index}"}])
|
||||||
|
for index in range(3)
|
||||||
|
))
|
||||||
|
|
||||||
|
assert [response.content for response in responses] == ["ok", "ok", "ok"]
|
||||||
|
assert context_calls == [(True, False)]
|
||||||
|
assert request_contexts == [context, context, context]
|
||||||
|
|
||||||
|
|
||||||
class _WarningCaptureLogger:
|
class _WarningCaptureLogger:
|
||||||
def __init__(self) -> None:
|
def __init__(self) -> None:
|
||||||
self.calls: list[tuple[str, tuple[Any, ...]]] = []
|
self.calls: list[tuple[str, tuple[Any, ...]]] = []
|
||||||
|
|||||||
Reference in New Issue
Block a user