feat(librarian): bounded retries, timeout wiring, and client reuse
- 2-attempt short-backoff retry for GETs and the read-only POST /query/* and /rag/search endpoints only; wiki writes are never retried (duplicate-page risk) - honor the defined-but-ignored LIBRARY_DESK_TIMEOUT config instead of hardcoded 60s/30s per-call values - hold ONE shared httpx.AsyncClient per librarian run via library_client_session (contextvar), instead of constructing a client per tool call; nested sessions are no-ops and custom targets still get their own client - read tools raise ModelRetry on transient HTTP errors (transport errors, 5xx, 429) so Agent(retries=2) engages; write tools keep returning safe failure messages Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Claude Fable 5
@@ -0,0 +1,232 @@
|
||||
"""
|
||||
Tests for bounded retries, timeout wiring, and client reuse in
|
||||
LibraryDeskClient, plus ModelRetry escalation from the read tools.
|
||||
"""
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from pydantic_ai import ModelRetry
|
||||
|
||||
from src.agents.librarian.client import (
|
||||
LibraryDeskClient,
|
||||
library_client_session,
|
||||
)
|
||||
from src.agents.librarian.tools import (
|
||||
create_wiki_page,
|
||||
hybrid_search,
|
||||
search_wiki,
|
||||
)
|
||||
from src.core.config import config
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _no_backoff(monkeypatch):
|
||||
"""Skip the retry backoff sleep in tests."""
|
||||
monkeypatch.setattr("src.agents.librarian.client._RETRY_BACKOFF_SECONDS", 0)
|
||||
|
||||
|
||||
def _ok_response(payload: dict) -> MagicMock:
|
||||
response = MagicMock()
|
||||
response.status_code = 200
|
||||
response.json.return_value = payload
|
||||
response.raise_for_status = MagicMock()
|
||||
return response
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client_with_mock():
|
||||
client = LibraryDeskClient(base_url="http://test:8089", api_key="test-key")
|
||||
client._client = AsyncMock(spec=httpx.AsyncClient)
|
||||
return client
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestBoundedRetries:
|
||||
"""2-attempt retry for GETs and read-only POST /query/*, /rag/search."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_retries_once_on_transport_error(self, client_with_mock):
|
||||
mock_httpx = client_with_mock._client
|
||||
mock_httpx.get.side_effect = [
|
||||
httpx.ConnectError("Connection refused"),
|
||||
_ok_response({"results": []}),
|
||||
]
|
||||
|
||||
results = await client_with_mock.search_wiki("docker", user="u")
|
||||
|
||||
assert results == []
|
||||
assert mock_httpx.get.call_count == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_gives_up_after_two_attempts(self, client_with_mock):
|
||||
mock_httpx = client_with_mock._client
|
||||
mock_httpx.get.side_effect = httpx.ConnectError("Connection refused")
|
||||
|
||||
with pytest.raises(httpx.ConnectError):
|
||||
await client_with_mock.search_wiki("docker", user="u")
|
||||
|
||||
assert mock_httpx.get.call_count == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_query_hybrid_retries_on_503(self, client_with_mock):
|
||||
mock_httpx = client_with_mock._client
|
||||
bad = MagicMock()
|
||||
bad.status_code = 503
|
||||
mock_httpx.post.side_effect = [
|
||||
bad,
|
||||
_ok_response({"results": [], "keywords": {}, "context": ""}),
|
||||
]
|
||||
|
||||
response = await client_with_mock.hybrid_search("docker", user="u")
|
||||
|
||||
assert response.results == []
|
||||
assert mock_httpx.post.call_count == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_wiki_write_is_never_retried(self, client_with_mock):
|
||||
"""POST /wiki/pages must not retry - it could duplicate pages."""
|
||||
mock_httpx = client_with_mock._client
|
||||
mock_httpx.post.side_effect = httpx.ConnectError("Connection refused")
|
||||
|
||||
with pytest.raises(httpx.ConnectError):
|
||||
await client_with_mock.create_wiki_page(
|
||||
title="T", path="/t", content="c", user="u"
|
||||
)
|
||||
|
||||
assert mock_httpx.post.call_count == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_smart_create_is_never_retried(self, client_with_mock):
|
||||
mock_httpx = client_with_mock._client
|
||||
mock_httpx.post.side_effect = httpx.ConnectError("Connection refused")
|
||||
|
||||
with pytest.raises(httpx.ConnectError):
|
||||
await client_with_mock.smart_create_wiki_page(
|
||||
topic="T", tags=["x"], user="u"
|
||||
)
|
||||
|
||||
assert mock_httpx.post.call_count == 1
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestTimeoutWiring:
|
||||
"""LIBRARY_DESK_TIMEOUT config replaces the hardcoded 60s/30s."""
|
||||
|
||||
def test_default_timeout_from_config(self):
|
||||
client = LibraryDeskClient()
|
||||
|
||||
assert client.timeout == config.LIBRARY_DESK_TIMEOUT
|
||||
|
||||
def test_explicit_timeout_wins(self):
|
||||
client = LibraryDeskClient(timeout=5)
|
||||
|
||||
assert client.timeout == 5
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestClientReuse:
|
||||
"""One shared HTTP connection per librarian run."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_clients_share_connection_inside_session(self):
|
||||
async with library_client_session():
|
||||
async with LibraryDeskClient() as c1:
|
||||
http1 = c1._client
|
||||
# shared connection survives client exit
|
||||
assert http1 is not None
|
||||
assert not http1.is_closed
|
||||
|
||||
async with LibraryDeskClient() as c2:
|
||||
assert c2._client is http1
|
||||
|
||||
# session close tears the shared connection down
|
||||
assert http1.is_closed
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_nested_sessions_are_noops(self):
|
||||
async with library_client_session():
|
||||
async with LibraryDeskClient() as c1:
|
||||
http1 = c1._client
|
||||
async with library_client_session():
|
||||
async with LibraryDeskClient() as c2:
|
||||
assert c2._client is http1
|
||||
# inner session exit must not close the shared connection
|
||||
assert not http1.is_closed
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_client_owns_connection_outside_session(self):
|
||||
async with LibraryDeskClient() as client:
|
||||
http_client = client._client
|
||||
|
||||
assert http_client.is_closed
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_custom_target_does_not_reuse_shared(self):
|
||||
async with library_client_session():
|
||||
async with LibraryDeskClient() as shared_client:
|
||||
shared_http = shared_client._client
|
||||
async with LibraryDeskClient(base_url="http://other:9999") as custom:
|
||||
assert custom._client is not shared_http
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestModelRetryEscalation:
|
||||
"""Read tools raise ModelRetry on transient errors so Agent(retries=2) engages."""
|
||||
|
||||
def _patched_client(self, mock_client):
|
||||
factory = MagicMock()
|
||||
factory.return_value.__aenter__ = AsyncMock(return_value=mock_client)
|
||||
factory.return_value.__aexit__ = AsyncMock(return_value=None)
|
||||
return patch("src.agents.librarian.tools.LibraryDeskClient", factory)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_tool_raises_model_retry_on_transport_error(self):
|
||||
mock_client = AsyncMock()
|
||||
mock_client.hybrid_search.side_effect = httpx.ConnectError(
|
||||
"Connection refused"
|
||||
)
|
||||
|
||||
with self._patched_client(mock_client):
|
||||
with pytest.raises(ModelRetry):
|
||||
await hybrid_search("docker")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_tool_raises_model_retry_on_5xx(self):
|
||||
request = httpx.Request("GET", "http://test:8089/wiki/search")
|
||||
response = httpx.Response(502, request=request)
|
||||
mock_client = AsyncMock()
|
||||
mock_client.search_wiki.side_effect = httpx.HTTPStatusError(
|
||||
"bad gateway", request=request, response=response
|
||||
)
|
||||
|
||||
with self._patched_client(mock_client):
|
||||
with pytest.raises(ModelRetry):
|
||||
await search_wiki("docker")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_tool_returns_safe_message_on_non_transient(self):
|
||||
mock_client = AsyncMock()
|
||||
mock_client.hybrid_search.side_effect = ValueError("bad parse")
|
||||
|
||||
with self._patched_client(mock_client):
|
||||
result = await hybrid_search("docker")
|
||||
|
||||
assert "unable" in result
|
||||
assert "bad parse" not in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_write_tool_never_raises_model_retry(self):
|
||||
mock_client = AsyncMock()
|
||||
mock_client.create_wiki_page.side_effect = httpx.ConnectError(
|
||||
"Connection refused"
|
||||
)
|
||||
|
||||
with self._patched_client(mock_client):
|
||||
result = await create_wiki_page(
|
||||
title="T", path="/t", content="c", tags=["x"]
|
||||
)
|
||||
|
||||
assert "unable" in result
|
||||
assert "Connection refused" not in result
|
||||
Reference in New Issue
Block a user