Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
50 changes: 30 additions & 20 deletions plugins/web/exa/provider.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@

import logging
import os
import threading
from typing import Any, Dict, List

from agent.web_search_provider import WebSearchProvider
Expand All @@ -36,6 +37,7 @@
# :mod:`tools.web_tools` so tests that do ``tools.web_tools._exa_client =
# None`` between cases see fresh state. The plugin reads/writes through
# that public module (see :func:`_get_exa_client`).
_client_lock = threading.Lock()


def _get_exa_client() -> Any:
Expand All @@ -47,41 +49,49 @@ def _get_exa_client() -> Any:
"""
import tools.web_tools as _wt

# Intentionally lock-free after publication; the slow path double-checks
# under the lock before constructing the singleton.
cached = getattr(_wt, "_exa_client", None)
if cached is not None:
return cached

from agent.web_search_provider import get_provider_env
with _client_lock:
cached = getattr(_wt, "_exa_client", None)
if cached is not None:
return cached

api_key = get_provider_env("EXA_API_KEY")
if not api_key:
raise ValueError(
"EXA_API_KEY environment variable not set. "
"Get your API key at https://exa.ai"
)
from agent.web_search_provider import get_provider_env

api_key = get_provider_env("EXA_API_KEY")
if not api_key:
raise ValueError(
"EXA_API_KEY environment variable not set. "
"Get your API key at https://exa.ai"
)

try:
from tools.lazy_deps import ensure as _lazy_ensure
try:
from tools.lazy_deps import ensure as _lazy_ensure

_lazy_ensure("search.exa", prompt=False)
except ImportError:
pass
except Exception as exc: # noqa: BLE001 — lazy_deps surfaces install hints
raise ImportError(str(exc))
_lazy_ensure("search.exa", prompt=False)
except ImportError:
pass
except Exception as exc: # noqa: BLE001 — lazy_deps surfaces install hints
raise ImportError(str(exc))

from exa_py import Exa # noqa: WPS433 — deliberately lazy
from exa_py import Exa # noqa: WPS433 — deliberately lazy

client = Exa(api_key=api_key)
client.headers["x-exa-integration"] = "hermes-agent"
_wt._exa_client = client
return client
client = Exa(api_key=api_key)
client.headers["x-exa-integration"] = "hermes-agent"
_wt._exa_client = client
return client


def _reset_client_for_tests() -> None:
"""Drop the cached Exa client so tests can re-instantiate cleanly."""
import tools.web_tools as _wt

_wt._exa_client = None
with _client_lock:
_wt._exa_client = None


class ExaWebSearchProvider(WebSearchProvider):
Expand Down
36 changes: 23 additions & 13 deletions plugins/web/parallel/provider.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@

import logging
import os
import threading
from typing import Any, Dict, List

from agent.web_search_provider import WebSearchProvider
Expand All @@ -41,6 +42,7 @@
# ``tools.web_tools._parallel_client = None`` between cases see fresh state.
# The plugin reads/writes through that public module (see
# :func:`_get_sync_client` / :func:`_get_async_client`).
_sync_client_lock = threading.Lock()


def _ensure_parallel_sdk_installed() -> None:
Expand Down Expand Up @@ -69,25 +71,32 @@ def _get_sync_client() -> Any:
"""
import tools.web_tools as _wt

# Intentionally lock-free after publication; the slow path double-checks
# under the lock before constructing the singleton.
cached = getattr(_wt, "_parallel_client", None)
if cached is not None:
return cached

from agent.web_search_provider import get_provider_env
with _sync_client_lock:
cached = getattr(_wt, "_parallel_client", None)
if cached is not None:
return cached

api_key = get_provider_env("PARALLEL_API_KEY")
if not api_key:
raise ValueError(
"PARALLEL_API_KEY environment variable not set. "
"Get your API key at https://parallel.ai"
)
from agent.web_search_provider import get_provider_env

_ensure_parallel_sdk_installed()
from parallel import Parallel # noqa: WPS433 — deliberately lazy
api_key = get_provider_env("PARALLEL_API_KEY")
if not api_key:
raise ValueError(
"PARALLEL_API_KEY environment variable not set. "
"Get your API key at https://parallel.ai"
)

client = Parallel(api_key=api_key)
_wt._parallel_client = client
return client
_ensure_parallel_sdk_installed()
from parallel import Parallel # noqa: WPS433 — deliberately lazy

client = Parallel(api_key=api_key)
_wt._parallel_client = client
return client


def _get_async_client() -> Any:
Expand Down Expand Up @@ -126,7 +135,8 @@ def _reset_clients_for_tests() -> None:
"""
import tools.web_tools as _wt

_wt._parallel_client = None
with _sync_client_lock:
_wt._parallel_client = None
_wt._async_parallel_client = None


Expand Down
122 changes: 122 additions & 0 deletions tests/plugins/web/test_provider_client_concurrency.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,122 @@
from concurrent.futures import ThreadPoolExecutor
import sys
import threading
import types

import pytest

import tools.web_tools as web_tools
from plugins.web.exa import provider as exa_provider
from plugins.web.parallel import provider as parallel_provider


@pytest.mark.parametrize(
(
"provider_module",
"lock_name",
"getter_name",
"cache_name",
"env_name",
"sdk_module_name",
"constructor_name",
"needs_headers",
),
[
pytest.param(
parallel_provider,
"_sync_client_lock",
"_get_sync_client",
"_parallel_client",
"PARALLEL_API_KEY",
"parallel",
"Parallel",
False,
id="parallel-sync",
),
pytest.param(
exa_provider,
"_client_lock",
"_get_exa_client",
"_exa_client",
"EXA_API_KEY",
"exa_py",
"Exa",
True,
id="exa",
),
],
)
def test_concurrent_first_use_constructs_one_client(
monkeypatch,
provider_module,
lock_name,
getter_name,
cache_name,
env_name,
sdk_module_name,
constructor_name,
needs_headers,
):
workers = 8
callers_ready = threading.Barrier(workers, timeout=5)
release_constructor = threading.Event()
state_changed = threading.Condition()
lock_attempts = [0]
constructed = []

class ObservedLock:
def __init__(self):
self._lock = threading.Lock()

def __enter__(self):
with state_changed:
lock_attempts[0] += 1
state_changed.notify_all()
self._lock.acquire()
return self

def __exit__(self, *_args):
self._lock.release()

def constructor(*_args, **_kwargs):
client = types.SimpleNamespace(headers={}) if needs_headers else object()
with state_changed:
constructed.append(client)
state_changed.notify_all()
assert release_constructor.wait(timeout=5)
return client

sdk_module = types.ModuleType(sdk_module_name)
setattr(sdk_module, constructor_name, constructor)
monkeypatch.setitem(sys.modules, sdk_module_name, sdk_module)
monkeypatch.setenv(env_name, "test-key")
monkeypatch.setattr(web_tools, cache_name, None)
monkeypatch.setattr(provider_module, lock_name, ObservedLock())
if provider_module is parallel_provider:
monkeypatch.setattr(
parallel_provider, "_ensure_parallel_sdk_installed", lambda: None
)
else:
from tools import lazy_deps

monkeypatch.setattr(lazy_deps, "ensure", lambda *_args, **_kwargs: None)

getter = getattr(provider_module, getter_name)

def get_client(_index):
callers_ready.wait()
return getter()

with ThreadPoolExecutor(max_workers=workers) as pool:
futures = [pool.submit(get_client, index) for index in range(workers)]
with state_changed:
reached_boundary = state_changed.wait_for(
lambda: lock_attempts[0] == workers or len(constructed) == workers,
timeout=5,
)
release_constructor.set()
returned = [future.result(timeout=5) for future in futures]

assert reached_boundary
assert len(constructed) == 1
assert all(client is returned[0] for client in returned)
Loading