diff --git a/agent_fox/core/client.py b/agent_fox/core/client.py index 72afe25c..32c80954 100644 --- a/agent_fox/core/client.py +++ b/agent_fox/core/client.py @@ -338,14 +338,17 @@ async def ai_call( async def _call() -> Any: client = create_async_anthropic_client() - return await cached_messages_create( - client, - model=model_entry.model_id, - max_tokens=max_tokens, - messages=messages, - system=system, - cache_policy=cache_policy, - ) + try: + return await cached_messages_create( + client, + model=model_entry.model_id, + max_tokens=max_tokens, + messages=messages, + system=system, + cache_policy=cache_policy, + ) + finally: + await client.close() response = await retry_api_call_async(_call, context=context) track_response_usage(response, model_entry.model_id, context) diff --git a/agent_fox/knowledge/consolidation.py b/agent_fox/knowledge/consolidation.py index 2ecc4afe..f9a4d359 100644 --- a/agent_fox/knowledge/consolidation.py +++ b/agent_fox/knowledge/consolidation.py @@ -260,24 +260,27 @@ async def _call_llm_json(model: str, prompt: str, context: dict) -> dict: from agent_fox.core.client import create_async_anthropic_client client = create_async_anthropic_client() - full_prompt = f"{prompt}\n\nContext:\n{json.dumps(context, indent=2, default=str)}" - - response = await client.messages.create( - model=model, - max_tokens=1024, - messages=[{"role": "user", "content": full_prompt}], - ) + try: + full_prompt = f"{prompt}\n\nContext:\n{json.dumps(context, indent=2, default=str)}" - block = response.content[0] - if not isinstance(block, TextBlock): - raise TypeError(f"Expected TextBlock, got {type(block).__name__}") - text = block.text.strip() - # Strip markdown code fences if present - if text.startswith("```"): - lines = text.splitlines() - text = "\n".join(line for line in lines if not line.startswith("```")).strip() + response = await client.messages.create( + model=model, + max_tokens=1024, + messages=[{"role": "user", "content": full_prompt}], + ) - return json.loads(text) + block = response.content[0] + if not isinstance(block, TextBlock): + raise TypeError(f"Expected TextBlock, got {type(block).__name__}") + text = block.text.strip() + # Strip markdown code fences if present + if text.startswith("```"): + lines = text.splitlines() + text = "\n".join(line for line in lines if not line.startswith("```")).strip() + + return json.loads(text) + finally: + await client.close() # --------------------------------------------------------------------------- diff --git a/agent_fox/nightshift/categories/quality_gate.py b/agent_fox/nightshift/categories/quality_gate.py index 0f5025c6..d08fb0cc 100644 --- a/agent_fox/nightshift/categories/quality_gate.py +++ b/agent_fox/nightshift/categories/quality_gate.py @@ -233,29 +233,34 @@ async def _run_ai_analysis( from agent_fox.core.json_extraction import extract_json_array backend = self._backend - if backend is None: + owns_backend = backend is None + if owns_backend: from agent_fox.core.client import create_async_anthropic_client backend = create_async_anthropic_client() - prompt = QUALITY_GATE_PROMPT.format(static_output=static_output) - response = await backend.messages.create( # type: ignore[attr-defined] - model=_model_id, - max_tokens=4096, - messages=[{"role": "user", "content": prompt}], - ) + try: + prompt = QUALITY_GATE_PROMPT.format(static_output=static_output) + response = await backend.messages.create( # type: ignore[attr-defined] + model=_model_id, + max_tokens=4096, + messages=[{"role": "user", "content": prompt}], + ) - # Emit cost for this auxiliary AI call (91-REQ-4.4) - from agent_fox.core.config import PricingConfig - from agent_fox.nightshift.cost_helpers import emit_auxiliary_cost + # Emit cost for this auxiliary AI call (91-REQ-4.4) + from agent_fox.core.config import PricingConfig + from agent_fox.nightshift.cost_helpers import emit_auxiliary_cost - emit_auxiliary_cost(sink, run_id, "quality_gate", response, _model_id, PricingConfig()) + emit_auxiliary_cost(sink, run_id, "quality_gate", response, _model_id, PricingConfig()) - response_text = response.content[0].text # type: ignore[attr-defined] + response_text = response.content[0].text # type: ignore[attr-defined] - items = extract_json_array(response_text) - if items is None: - raise ValueError("AI returned unparseable JSON") + items = extract_json_array(response_text) + if items is None: + raise ValueError("AI returned unparseable JSON") + finally: + if owns_backend: + await backend.close() # type: ignore[union-attr] findings: list[Finding] = [] for item in items: diff --git a/tests/unit/core/test_client_close.py b/tests/unit/core/test_client_close.py new file mode 100644 index 00000000..c95887b1 --- /dev/null +++ b/tests/unit/core/test_client_close.py @@ -0,0 +1,135 @@ +"""Regression tests for async client cleanup (issue #506). + +Verifies that ai_call() and direct create_async_anthropic_client() callers +close the AsyncAnthropic client after use, preventing RuntimeError('Event +loop is closed') tracebacks on shutdown. +""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + + +def _make_mock_client() -> MagicMock: + """Build a mock AsyncAnthropic client with a trackable close().""" + client = MagicMock() + client.close = AsyncMock() + + fake_response = MagicMock() + fake_response.content = [MagicMock(text="ok")] + fake_response.usage = MagicMock( + input_tokens=10, + output_tokens=5, + cache_read_input_tokens=0, + cache_creation_input_tokens=0, + ) + client.messages = MagicMock() + client.messages.create = AsyncMock(return_value=fake_response) + return client + + +class TestAiCallClosesClient: + """ai_call() must close the client even on success and on failure.""" + + @pytest.mark.asyncio + async def test_client_closed_on_success(self) -> None: + mock_client = _make_mock_client() + + async def fake_retry(fn, **kw): # noqa: ANN001, ANN003, ARG001 + return await fn() + + with ( + patch( + "agent_fox.core.client.create_async_anthropic_client", + return_value=mock_client, + ), + patch("agent_fox.core.models.resolve_model") as mock_resolve, + patch("agent_fox.core.retry.retry_api_call_async", side_effect=fake_retry), + patch("agent_fox.core.token_tracker.track_response_usage"), + ): + mock_resolve.return_value = MagicMock(model_id="claude-sonnet-4-6") + + from agent_fox.core.client import ai_call + + await ai_call( + model_tier="standard", + max_tokens=100, + messages=[{"role": "user", "content": "test"}], + context="test", + ) + + mock_client.close.assert_awaited_once() + + @pytest.mark.asyncio + async def test_client_closed_on_api_error(self) -> None: + mock_client = _make_mock_client() + mock_client.messages.create = AsyncMock(side_effect=RuntimeError("API error")) + + async def fake_retry(fn, **kw): # noqa: ANN001, ANN003, ARG001 + return await fn() + + with ( + patch( + "agent_fox.core.client.create_async_anthropic_client", + return_value=mock_client, + ), + patch("agent_fox.core.models.resolve_model") as mock_resolve, + patch("agent_fox.core.retry.retry_api_call_async", side_effect=fake_retry), + patch("agent_fox.core.token_tracker.track_response_usage"), + ): + mock_resolve.return_value = MagicMock(model_id="claude-sonnet-4-6") + + from agent_fox.core.client import ai_call + + with pytest.raises(RuntimeError, match="API error"): + await ai_call( + model_tier="standard", + max_tokens=100, + messages=[{"role": "user", "content": "test"}], + context="test", + ) + + mock_client.close.assert_awaited_once() + + +class TestConsolidationCallLlmJsonClosesClient: + """_call_llm_json() must close the client after use.""" + + @pytest.mark.asyncio + async def test_client_closed_after_llm_call(self) -> None: + mock_client = _make_mock_client() + from anthropic.types import TextBlock + + fake_block = TextBlock(type="text", text='{"result": "ok"}') + fake_response = MagicMock() + fake_response.content = [fake_block] + mock_client.messages.create = AsyncMock(return_value=fake_response) + + with patch( + "agent_fox.core.client.create_async_anthropic_client", + return_value=mock_client, + ): + from agent_fox.knowledge.consolidation import _call_llm_json + + result = await _call_llm_json("claude-sonnet-4-6", "test prompt", {"key": "value"}) + + mock_client.close.assert_awaited_once() + assert result == {"result": "ok"} + + @pytest.mark.asyncio + async def test_client_closed_on_llm_error(self) -> None: + mock_client = _make_mock_client() + mock_client.messages.create = AsyncMock(side_effect=RuntimeError("LLM error")) + + with patch( + "agent_fox.core.client.create_async_anthropic_client", + return_value=mock_client, + ): + from agent_fox.knowledge.consolidation import _call_llm_json + + with pytest.raises(RuntimeError, match="LLM error"): + await _call_llm_json("claude-sonnet-4-6", "test prompt", {"key": "value"}) + + mock_client.close.assert_awaited_once()