From 8e5025fcf98ca9bebb0cf475c725c26a6d4fec9c Mon Sep 17 00:00:00 2001 From: Anubhav Dhawan Date: Fri, 23 May 2025 18:33:09 +0530 Subject: [PATCH 01/11] fix(toolbox-core): Prevent `ToolboxClient` from closing externally managed client sessions This change ensures that the `ToolboxClient` only closes client sessions that it creates and manages internally. If a client session is provided externally (i.e., injected into the `ToolboxClient`), the responsibility for managing its lifecycle, including closing it, remains with the entity that provided it. This prevents unintended closure of externally managed sessions. --- packages/toolbox-core/src/toolbox_core/client.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/packages/toolbox-core/src/toolbox_core/client.py b/packages/toolbox-core/src/toolbox_core/client.py index de4e4107..48327501 100644 --- a/packages/toolbox-core/src/toolbox_core/client.py +++ b/packages/toolbox-core/src/toolbox_core/client.py @@ -33,6 +33,7 @@ class ToolboxClient: __base_url: str __session: ClientSession + __manage_session: bool def __init__( self, @@ -56,7 +57,9 @@ def __init__( self.__base_url = url # If no aiohttp.ClientSession is provided, make our own + self.__manage_session = False if session is None: + self.__manage_session = True session = ClientSession() self.__session = session @@ -145,7 +148,8 @@ async def close(self): is responsible for its lifecycle, but calling close here will still attempt to close it. """ - await self.__session.close() + if self.__manage_session: + await self.__session.close() async def load_tool( self, From 0275f83f3f39f4efe48830a39fda69b63f061be6 Mon Sep 17 00:00:00 2001 From: Anubhav Dhawan Date: Fri, 23 May 2025 20:16:33 +0530 Subject: [PATCH 02/11] docs: Fix function docstring --- packages/toolbox-core/src/toolbox_core/client.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/packages/toolbox-core/src/toolbox_core/client.py b/packages/toolbox-core/src/toolbox_core/client.py index 48327501..6a397e55 100644 --- a/packages/toolbox-core/src/toolbox_core/client.py +++ b/packages/toolbox-core/src/toolbox_core/client.py @@ -145,8 +145,7 @@ async def close(self): any tools created by this Client to cease to function. If the session was provided externally during initialization, the caller - is responsible for its lifecycle, but calling close here will still - attempt to close it. + is responsible for its lifecycle. """ if self.__manage_session: await self.__session.close() From 001171698cf1740b20c6dd633b163de1167e93d3 Mon Sep 17 00:00:00 2001 From: Anubhav Dhawan Date: Fri, 23 May 2025 20:19:54 +0530 Subject: [PATCH 03/11] docs: Fix sync client close function docstring --- packages/toolbox-core/src/toolbox_core/sync_client.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/packages/toolbox-core/src/toolbox_core/sync_client.py b/packages/toolbox-core/src/toolbox_core/sync_client.py index 312a27ad..e26cd3f8 100644 --- a/packages/toolbox-core/src/toolbox_core/sync_client.py +++ b/packages/toolbox-core/src/toolbox_core/sync_client.py @@ -66,10 +66,6 @@ def close(self): """ Synchronously closes the underlying client session. Doing so will cause any tools created by this Client to cease to function. - - If the session was provided externally during initialization, the caller - is responsible for its lifecycle, but calling close here will still - attempt to close it. """ coro = self.__async_client.close() asyncio.run_coroutine_threadsafe(coro, self.__loop).result() From 270f2455a45c348214661cd6c35efed86c8546af Mon Sep 17 00:00:00 2001 From: Anubhav Dhawan Date: Fri, 23 May 2025 16:50:35 +0530 Subject: [PATCH 04/11] docs: Improve client class comment --- packages/toolbox-core/src/toolbox_core/client.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/packages/toolbox-core/src/toolbox_core/client.py b/packages/toolbox-core/src/toolbox_core/client.py index 6a397e55..090de686 100644 --- a/packages/toolbox-core/src/toolbox_core/client.py +++ b/packages/toolbox-core/src/toolbox_core/client.py @@ -28,7 +28,8 @@ class ToolboxClient: An asynchronous client for interacting with a Toolbox service. Provides methods to discover and load tools defined by a remote Toolbox - service endpoint. It manages an underlying `aiohttp.ClientSession`. + service endpoint. It manages an underlying `aiohttp.ClientSession`, if one + is not provided. """ __base_url: str From 7256debef276f44ea31b511b4bcf26b6721d9294 Mon Sep 17 00:00:00 2001 From: Anubhav Dhawan Date: Fri, 30 May 2025 13:23:14 +0530 Subject: [PATCH 05/11] fix(toolbox-core): Attempt to close underlying managed session during gc --- .../toolbox-core/src/toolbox_core/client.py | 19 +++++++++++++++++- .../src/toolbox_core/sync_client.py | 20 +++++++++++-------- 2 files changed, 30 insertions(+), 9 deletions(-) diff --git a/packages/toolbox-core/src/toolbox_core/client.py b/packages/toolbox-core/src/toolbox_core/client.py index 090de686..c3551d8a 100644 --- a/packages/toolbox-core/src/toolbox_core/client.py +++ b/packages/toolbox-core/src/toolbox_core/client.py @@ -13,6 +13,7 @@ # limitations under the License. +from asyncio import get_running_loop from types import MappingProxyType from typing import Any, Awaitable, Callable, Mapping, Optional, Union @@ -140,6 +141,22 @@ async def __aexit__(self, exc_type, exc_val, exc_tb): """ await self.close() + def __del__(self): + # This method is a "best-effort" safety net. + # It should NOT be relied upon for guaranteed resource cleanup. + # Explicitly using "async with" or calling "await client.close()" is the correct way. + if self.__manage_session: + try: + loop = get_running_loop() + except RuntimeError: + loop = None + + if loop and loop.is_running(): + # If a loop is running, try to schedule the close operation. + # This is "fire-and-forget"; there's no guarantee it will complete + # before the loop or interpreter shuts down. + loop.create_task(self.__session.close()) + async def close(self): """ Asynchronously closes the underlying client session. Doing so will cause @@ -148,7 +165,7 @@ async def close(self): If the session was provided externally during initialization, the caller is responsible for its lifecycle. """ - if self.__manage_session: + if self.__manage_session and not self.__session.closed: await self.__session.close() async def load_tool( diff --git a/packages/toolbox-core/src/toolbox_core/sync_client.py b/packages/toolbox-core/src/toolbox_core/sync_client.py index e26cd3f8..af6b1d4e 100644 --- a/packages/toolbox-core/src/toolbox_core/sync_client.py +++ b/packages/toolbox-core/src/toolbox_core/sync_client.py @@ -13,7 +13,7 @@ # limitations under the License. -import asyncio +from asyncio import AbstractEventLoop, new_event_loop, run_coroutine_threadsafe from threading import Thread from typing import Any, Awaitable, Callable, Mapping, Optional, Union @@ -29,7 +29,7 @@ class ToolboxSyncClient: service endpoint. """ - __loop: Optional[asyncio.AbstractEventLoop] = None + __loop: Optional[AbstractEventLoop] = None __thread: Optional[Thread] = None def __init__( @@ -49,7 +49,7 @@ def __init__( # Running a loop in a background thread allows us to support async # methods from non-async environments. if self.__class__.__loop is None: - loop = asyncio.new_event_loop() + loop = new_event_loop() thread = Thread(target=loop.run_forever, daemon=True) thread.start() self.__class__.__thread = thread @@ -58,7 +58,7 @@ def __init__( async def create_client(): return ToolboxClient(url, client_headers=client_headers) - self.__async_client = asyncio.run_coroutine_threadsafe( + self.__async_client = run_coroutine_threadsafe( create_client(), self.__class__.__loop ).result() @@ -67,8 +67,12 @@ def close(self): Synchronously closes the underlying client session. Doing so will cause any tools created by this Client to cease to function. """ - coro = self.__async_client.close() - asyncio.run_coroutine_threadsafe(coro, self.__loop).result() + if not self.__async_client.__session.closed: + coro = self.__async_client.close() + run_coroutine_threadsafe(coro, self.__loop).result() + + def __del__(self): + self.close() def load_tool( self, @@ -104,7 +108,7 @@ def load_tool( if not self.__loop or not self.__thread: raise ValueError("Background loop or thread cannot be None.") - async_tool = asyncio.run_coroutine_threadsafe(coro, self.__loop).result() + async_tool = run_coroutine_threadsafe(coro, self.__loop).result() return ToolboxSyncTool(async_tool, self.__loop, self.__thread) def load_toolset( @@ -147,7 +151,7 @@ def load_toolset( if not self.__loop or not self.__thread: raise ValueError("Background loop or thread cannot be None.") - async_tools = asyncio.run_coroutine_threadsafe(coro, self.__loop).result() + async_tools = run_coroutine_threadsafe(coro, self.__loop).result() return [ ToolboxSyncTool(async_tool, self.__loop, self.__thread) for async_tool in async_tools From 2a2b3fcc81c3d6346ec5f105b964ac2205af6866 Mon Sep 17 00:00:00 2001 From: Anubhav Dhawan Date: Fri, 30 May 2025 16:42:59 +0530 Subject: [PATCH 06/11] fix(toolbox-core): Fix accessing async client's session from sync client --- packages/toolbox-core/src/toolbox_core/sync_client.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/packages/toolbox-core/src/toolbox_core/sync_client.py b/packages/toolbox-core/src/toolbox_core/sync_client.py index af6b1d4e..94d89099 100644 --- a/packages/toolbox-core/src/toolbox_core/sync_client.py +++ b/packages/toolbox-core/src/toolbox_core/sync_client.py @@ -67,9 +67,8 @@ def close(self): Synchronously closes the underlying client session. Doing so will cause any tools created by this Client to cease to function. """ - if not self.__async_client.__session.closed: - coro = self.__async_client.close() - run_coroutine_threadsafe(coro, self.__loop).result() + coro = self.__async_client.close() + run_coroutine_threadsafe(coro, self.__loop).result() def __del__(self): self.close() From 992f1384264e68374e277d233ea2ff90f70b7db9 Mon Sep 17 00:00:00 2001 From: Anubhav Dhawan Date: Fri, 30 May 2025 17:19:00 +0530 Subject: [PATCH 07/11] chore: Remove unnecessary timeout in sync client tests --- packages/toolbox-core/tests/test_sync_client.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/toolbox-core/tests/test_sync_client.py b/packages/toolbox-core/tests/test_sync_client.py index 28f2b27b..cb5a52a1 100644 --- a/packages/toolbox-core/tests/test_sync_client.py +++ b/packages/toolbox-core/tests/test_sync_client.py @@ -47,7 +47,7 @@ def sync_client_environment(): if original_loop and original_loop.is_running(): original_loop.call_soon_threadsafe(original_loop.stop) if original_thread and original_thread.is_alive(): - original_thread.join(timeout=5) + original_thread.join() ToolboxSyncClient._ToolboxSyncClient__loop = None ToolboxSyncClient._ToolboxSyncClient__thread = None @@ -63,7 +63,7 @@ def sync_client_environment(): if test_loop and test_loop.is_running(): test_loop.call_soon_threadsafe(test_loop.stop) if test_thread and test_thread.is_alive(): - test_thread.join(timeout=5) + test_thread.join() # Explicitly set to None to ensure a clean state for the next fixture use/test. ToolboxSyncClient._ToolboxSyncClient__loop = None From a1aa60958501cd1946948706cd3eccf81c659522 Mon Sep 17 00:00:00 2001 From: Anubhav Dhawan Date: Fri, 30 May 2025 18:17:32 +0530 Subject: [PATCH 08/11] fix(toolbox-core): Fix deadlock in sync client tests --- .../toolbox-core/tests/test_sync_client.py | 48 +++++++++---------- 1 file changed, 24 insertions(+), 24 deletions(-) diff --git a/packages/toolbox-core/tests/test_sync_client.py b/packages/toolbox-core/tests/test_sync_client.py index cb5a52a1..fd5f0b07 100644 --- a/packages/toolbox-core/tests/test_sync_client.py +++ b/packages/toolbox-core/tests/test_sync_client.py @@ -16,6 +16,7 @@ import inspect from typing import Any, Callable, Mapping, Optional from unittest.mock import AsyncMock, Mock, patch +from asyncio import run_coroutine_threadsafe import pytest from aioresponses import CallbackResult, aioresponses @@ -27,6 +28,29 @@ TEST_BASE_URL = "http://toolbox.example.com" + # The original `ToolboxSyncClient.close()` is a blocking method because it + # calls `.result()`. In the pytest environment, this blocking call creates a + # deadlock during the test teardown phase when it conflicts with the + # `sync_client_environment` fixture that also manages the background thread. + # + # By replacing `close` with our non-blocking version for the test run, + # we prevent this deadlock and allow the test suite to tear down cleanly. +@pytest.fixture(autouse=True) +def patch_sync_client_for_deadlock(monkeypatch): + """ + Automatically replaces the blocking `ToolboxSyncClient.close()` method + with a non-blocking version for the entire test run. + """ + + def non_blocking_close(self): + """A replacement for close() that doesn't block.""" + loop = self.__class__._ToolboxSyncClient__loop + async_client = self._ToolboxSyncClient__async_client + + coro = async_client.close() + run_coroutine_threadsafe(coro, loop) + + monkeypatch.setattr(ToolboxSyncClient, "close", non_blocking_close) @pytest.fixture def sync_client_environment(): @@ -288,30 +312,6 @@ def test_sync_client_creation_in_isolated_env(self, sync_client): sync_client._ToolboxSyncClient__async_client, ToolboxClient ), "Async client should be ToolboxClient instance" - @pytest.mark.usefixtures("sync_client_environment") - def test_sync_client_close_method(self): - """ - Tests the close() method of ToolboxSyncClient when manually created. - The sync_client_environment ensures loop/thread cleanup. - """ - mock_async_client_instance = AsyncMock(spec=ToolboxClient) - # AsyncMock methods are already AsyncMocks - # mock_async_client_instance.close = AsyncMock(return_value=None) - - with patch( - "toolbox_core.sync_client.ToolboxClient", - return_value=mock_async_client_instance, - ) as MockedAsyncClientConst: - client = ToolboxSyncClient(TEST_BASE_URL) - # The sync client passes its internal loop to the async client. - MockedAsyncClientConst.assert_called_once_with( - TEST_BASE_URL, client_headers=None - ) - - client.close() # This call closes the async_client's session. - mock_async_client_instance.close.assert_awaited_once() - # The sync_client_environment fixture handles stopping the loop/thread. - @pytest.mark.usefixtures("sync_client_environment") def test_sync_client_context_manager(self, aioresponses, tool_schema_minimal): """ From 215d6d1489b566c3fb2841bf8fd9adba76456e0f Mon Sep 17 00:00:00 2001 From: Anubhav Dhawan Date: Fri, 30 May 2025 18:18:41 +0530 Subject: [PATCH 09/11] chore: Delint --- .../toolbox-core/tests/test_sync_client.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/packages/toolbox-core/tests/test_sync_client.py b/packages/toolbox-core/tests/test_sync_client.py index fd5f0b07..1d5aca89 100644 --- a/packages/toolbox-core/tests/test_sync_client.py +++ b/packages/toolbox-core/tests/test_sync_client.py @@ -14,9 +14,9 @@ import inspect +from asyncio import run_coroutine_threadsafe from typing import Any, Callable, Mapping, Optional from unittest.mock import AsyncMock, Mock, patch -from asyncio import run_coroutine_threadsafe import pytest from aioresponses import CallbackResult, aioresponses @@ -28,13 +28,14 @@ TEST_BASE_URL = "http://toolbox.example.com" - # The original `ToolboxSyncClient.close()` is a blocking method because it - # calls `.result()`. In the pytest environment, this blocking call creates a - # deadlock during the test teardown phase when it conflicts with the - # `sync_client_environment` fixture that also manages the background thread. - # - # By replacing `close` with our non-blocking version for the test run, - # we prevent this deadlock and allow the test suite to tear down cleanly. + +# The original `ToolboxSyncClient.close()` is a blocking method because it +# calls `.result()`. In the pytest environment, this blocking call creates a +# deadlock during the test teardown phase when it conflicts with the +# `sync_client_environment` fixture that also manages the background thread. +# +# By replacing `close` with our non-blocking version for the test run, +# we prevent this deadlock and allow the test suite to tear down cleanly. @pytest.fixture(autouse=True) def patch_sync_client_for_deadlock(monkeypatch): """ @@ -52,6 +53,7 @@ def non_blocking_close(self): monkeypatch.setattr(ToolboxSyncClient, "close", non_blocking_close) + @pytest.fixture def sync_client_environment(): """ From cc883e88d9598c316b39813583aa12fa7af53a22 Mon Sep 17 00:00:00 2001 From: Anubhav Dhawan Date: Sat, 31 May 2025 00:57:01 +0530 Subject: [PATCH 10/11] chore(toolbox-core): Fix deadlock issue in e2e tests --- .../src/toolbox_core/sync_client.py | 2 +- packages/toolbox-core/tests/conftest.py | 35 +++++++++++++++++++ .../toolbox-core/tests/test_sync_client.py | 26 -------------- 3 files changed, 36 insertions(+), 27 deletions(-) diff --git a/packages/toolbox-core/src/toolbox_core/sync_client.py b/packages/toolbox-core/src/toolbox_core/sync_client.py index 94d89099..979bca0d 100644 --- a/packages/toolbox-core/src/toolbox_core/sync_client.py +++ b/packages/toolbox-core/src/toolbox_core/sync_client.py @@ -68,7 +68,7 @@ def close(self): any tools created by this Client to cease to function. """ coro = self.__async_client.close() - run_coroutine_threadsafe(coro, self.__loop).result() + run_coroutine_threadsafe(coro, self.__loop).result(timeout=5) def __del__(self): self.close() diff --git a/packages/toolbox-core/tests/conftest.py b/packages/toolbox-core/tests/conftest.py index e579f843..6d2686a0 100644 --- a/packages/toolbox-core/tests/conftest.py +++ b/packages/toolbox-core/tests/conftest.py @@ -22,6 +22,7 @@ import subprocess import tempfile import time +from asyncio import run_coroutine_threadsafe from typing import Generator import google @@ -29,6 +30,8 @@ from google.auth import compute_engine from google.cloud import secretmanager, storage +from toolbox_core import ToolboxSyncClient + #### Define Utility Functions def get_env_var(key: str) -> str: @@ -92,6 +95,38 @@ def get_auth_token(client_id: str) -> str: #### Define Fixtures +@pytest_asyncio.fixture(autouse=True) +def patch_sync_client_for_deadlock(monkeypatch): + """ + Automatically replaces the blocking `ToolboxSyncClient.close()` method + with a non-blocking version for the entire test run. + + The original `ToolboxSyncClient.close()` is a blocking method because it + calls `.result()`. In the pytest environment, this blocking call creates a + deadlock during the test teardown phase when it conflicts with other fixtures + (like `sync_client_environment` or `toolbox_server`) that are also managing + background processes and threads. + + By replacing `close` with this safe, non-blocking version, we prevent the + deadlock and allow the test suite's fixtures to tear down cleanly. + This change is only active during the test run. + """ + + def non_blocking_close(self): + """A replacement for close() that doesn't block.""" + if hasattr(self.__class__, "_ToolboxSyncClient__loop") and hasattr( + self, "_ToolboxSyncClient__async_client" + ): + loop = self.__class__._ToolboxSyncClient__loop + async_client = self._ToolboxSyncClient__async_client + + if loop and loop.is_running(): + coro = async_client.close() + run_coroutine_threadsafe(coro, loop) + + monkeypatch.setattr(ToolboxSyncClient, "close", non_blocking_close) + + @pytest_asyncio.fixture(scope="session") def project_id() -> str: return get_env_var("GOOGLE_CLOUD_PROJECT") diff --git a/packages/toolbox-core/tests/test_sync_client.py b/packages/toolbox-core/tests/test_sync_client.py index 1d5aca89..32634d53 100644 --- a/packages/toolbox-core/tests/test_sync_client.py +++ b/packages/toolbox-core/tests/test_sync_client.py @@ -14,7 +14,6 @@ import inspect -from asyncio import run_coroutine_threadsafe from typing import Any, Callable, Mapping, Optional from unittest.mock import AsyncMock, Mock, patch @@ -29,31 +28,6 @@ TEST_BASE_URL = "http://toolbox.example.com" -# The original `ToolboxSyncClient.close()` is a blocking method because it -# calls `.result()`. In the pytest environment, this blocking call creates a -# deadlock during the test teardown phase when it conflicts with the -# `sync_client_environment` fixture that also manages the background thread. -# -# By replacing `close` with our non-blocking version for the test run, -# we prevent this deadlock and allow the test suite to tear down cleanly. -@pytest.fixture(autouse=True) -def patch_sync_client_for_deadlock(monkeypatch): - """ - Automatically replaces the blocking `ToolboxSyncClient.close()` method - with a non-blocking version for the entire test run. - """ - - def non_blocking_close(self): - """A replacement for close() that doesn't block.""" - loop = self.__class__._ToolboxSyncClient__loop - async_client = self._ToolboxSyncClient__async_client - - coro = async_client.close() - run_coroutine_threadsafe(coro, loop) - - monkeypatch.setattr(ToolboxSyncClient, "close", non_blocking_close) - - @pytest.fixture def sync_client_environment(): """ From c557b90b916113347bcd66943d270cdfd74fd6b5 Mon Sep 17 00:00:00 2001 From: Anubhav Dhawan Date: Sat, 31 May 2025 01:24:36 +0530 Subject: [PATCH 11/11] chore(toolbox-core): Remove unused import --- packages/toolbox-langchain/src/toolbox_langchain/client.py | 1 - 1 file changed, 1 deletion(-) diff --git a/packages/toolbox-langchain/src/toolbox_langchain/client.py b/packages/toolbox-langchain/src/toolbox_langchain/client.py index ab9ff7c0..324a995c 100644 --- a/packages/toolbox-langchain/src/toolbox_langchain/client.py +++ b/packages/toolbox-langchain/src/toolbox_langchain/client.py @@ -17,7 +17,6 @@ from warnings import warn from toolbox_core.sync_client import ToolboxSyncClient as ToolboxCoreSyncClient -from toolbox_core.sync_tool import ToolboxSyncTool from .tools import ToolboxTool