diff --git a/aiofmp/base.py b/aiofmp/base.py index 68c470a..b67c455 100644 --- a/aiofmp/base.py +++ b/aiofmp/base.py @@ -140,6 +140,15 @@ def __init__( # Session management self._session: aiohttp.ClientSession | None = None self._session_owner = True + # Reference count of active start()/`async with` scopes. A shared + # client (e.g. the MCP server's global singleton) is entered + # concurrently by many in-flight requests, each wrapping its work in + # `async with client:`. Closing the session on the FIRST scope exit + # would tear it out from under the others ("Connector is closed", + # then a None session -> "'NoneType' object has no attribute 'get'"). + # Instead we only close once the LAST concurrent scope exits. + self._session_refcount = 0 + self._session_lock = asyncio.Lock() # Rate limiting self._request_semaphore = asyncio.Semaphore(max_concurrent_requests) @@ -171,19 +180,40 @@ async def __aexit__(self, exc_type, exc_val, exc_tb): await self.close() async def start(self): - """Start the client session if not already started""" - if self._session is None: - timeout = aiohttp.ClientTimeout(total=self.timeout) - self._session = aiohttp.ClientSession(timeout=timeout) - self._session_owner = True - logger.debug("FMP client session started") + """Start the client session if not already started. + + Reference-counted: concurrent ``async with client:`` scopes on a shared + client all share a single session. See :meth:`close`. + """ + async with self._session_lock: + if self._session is None: + timeout = aiohttp.ClientTimeout(total=self.timeout) + self._session = aiohttp.ClientSession(timeout=timeout) + self._session_owner = True + logger.debug("FMP client session started") + # Count the scope only after a session is guaranteed to exist, so a + # failed session creation can't leak a reference (a raising + # __aenter__ means __aexit__/close() never runs to balance it). + self._session_refcount += 1 async def close(self): - """Close the client session""" - if self._session_owner and self._session: - await self._session.close() - self._session = None - logger.debug("FMP client session closed") + """Close the client session once the last active scope exits. + + Decrements the scope reference count and only tears down the + underlying session when it reaches zero, so overlapping requests on a + shared client keep a live session for the duration of their own scope. + """ + async with self._session_lock: + if self._session_refcount > 0: + self._session_refcount -= 1 + if ( + self._session_refcount == 0 + and self._session_owner + and self._session is not None + ): + await self._session.close() + self._session = None + logger.debug("FMP client session closed") async def _make_request( self, endpoint: str, params: dict[str, Any] | None = None, method: str = "GET" diff --git a/tests/test_base.py b/tests/test_base.py new file mode 100644 index 0000000..d49e910 --- /dev/null +++ b/tests/test_base.py @@ -0,0 +1,100 @@ +""" +Unit tests for FMPBaseClient session lifecycle. + +Regression coverage for the shared-session concurrency race: the MCP server +reuses a single global client that many in-flight requests enter concurrently +via ``async with client:``. The first scope to exit must NOT close the session +out from under the others (which produced "Connector is closed" and +"'NoneType' object has no attribute 'get'" in production). +""" + +import asyncio +from unittest.mock import patch + +import pytest + +from aiofmp.base import FMPBaseClient + + +class _FakeSession: + """Minimal stand-in for aiohttp.ClientSession that tracks close().""" + + instances: list["_FakeSession"] = [] + + def __init__(self, *args, **kwargs): + self.closed = False + _FakeSession.instances.append(self) + + async def close(self): + self.closed = True + + +@pytest.fixture(autouse=True) +def _reset_fake_sessions(): + _FakeSession.instances.clear() + yield + _FakeSession.instances.clear() + + +class TestSessionLifecycle: + """Reference-counted start()/close() for safe concurrent use.""" + + @pytest.mark.asyncio + async def test_overlapping_scopes_share_one_session(self): + client = FMPBaseClient(api_key="test") + with patch("aiofmp.base.aiohttp.ClientSession", _FakeSession): + await client.start() # scope A enters + session = client._session + assert session is not None + + await client.start() # scope B enters (concurrent) + assert client._session is session # reuse, not a second session + assert len(_FakeSession.instances) == 1 + + # Scope A exits — session MUST survive for still-active scope B. + await client.close() + assert client._session is session + assert session.closed is False + + # Scope B exits — last user gone, now it closes. + await client.close() + assert client._session is None + assert session.closed is True + + @pytest.mark.asyncio + async def test_new_session_created_after_full_drain(self): + client = FMPBaseClient(api_key="test") + with patch("aiofmp.base.aiohttp.ClientSession", _FakeSession): + await client.start() + first = client._session + await client.close() + assert client._session is None + + # A later request re-opens a fresh session. + await client.start() + assert client._session is not None + assert client._session is not first + await client.close() + + @pytest.mark.asyncio + async def test_concurrent_enter_exit_keeps_session_alive_mid_flight(self): + """Mimics repro: many concurrent `async with` scopes, one shared client.""" + client = FMPBaseClient(api_key="test") + observed_closed_mid_request = [] + + with patch("aiofmp.base.aiohttp.ClientSession", _FakeSession): + + async def one_request(): + async with client: # start() / close() + # Yield control so other scopes interleave enter/exit here. + await asyncio.sleep(0) + # The session we're about to "use" must be open. + observed_closed_mid_request.append(client._session is None) + await asyncio.sleep(0) + + await asyncio.gather(*[one_request() for _ in range(20)]) + + # No request ever saw a torn-down session while inside its own scope. + assert not any(observed_closed_mid_request) + # And everything is cleaned up after the last scope exits. + assert client._session is None