move manage session into transport

Author: twishabansal

packages/toolbox-core/src/toolbox_core/client.py

@@ -57,12 +57,7 @@ def __init__(
             client_headers: Headers to include in each request sent through this client.
         """
 
-        # If no aiohttp.ClientSession is provided, make our own
-        manage_session = False
-        if session is None:
-            manage_session = True
-            session = ClientSession()
-        self.__transport = ToolboxTransport(url, session, manage_session)
+        self.__transport = ToolboxTransport(url, session)
         self.__client_headers = client_headers if client_headers is not None else {}
 
     def __parse_tool(

packages/toolbox-core/src/toolbox_core/toolbox_transport.py

@@ -23,10 +23,16 @@
 class ToolboxTransport(ITransport):
     """Transport for the native Toolbox protocol."""
 
-    def __init__(self, base_url: str, session: ClientSession, manage_session: bool):
+    def __init__(self, base_url: str, session: Optional[ClientSession]):
         self.__base_url = base_url
-        self.__session = session
-        self.__manage_session = manage_session
+
+        # If no aiohttp.ClientSession is provided, make our own
+        self.__manage_session = False
+        if session is not None:
+            self.__session = session
+        else:
+            self.__manage_session = True
+            self.__session = ClientSession()
 
     @property
     def base_url(self) -> str:

packages/toolbox-core/tests/test_tool.py

@@ -112,7 +112,7 @@ def toolbox_tool(
     sample_tool_description: str,
 ) -> ToolboxTool:
     """Fixture for a ToolboxTool instance with common test setup."""
-    transport = ToolboxTransport(TEST_BASE_URL, http_session, False)
+    transport = ToolboxTransport(TEST_BASE_URL, http_session)
     return ToolboxTool(
         transport=transport,
         name=TEST_TOOL_NAME,
@@ -231,7 +231,7 @@ async def test_tool_creation_callable_and_run(
 
     with aioresponses() as m:
         m.post(invoke_url, status=200, payload=mock_server_response_body)
-        transport = ToolboxTransport(base_url, http_session, False)
+        transport = ToolboxTransport(base_url, http_session)
 
         tool_instance = ToolboxTool(
             transport=transport,
@@ -277,7 +277,7 @@ async def test_tool_run_with_pydantic_validation_error(
 
     with aioresponses() as m:
         m.post(invoke_url, status=200, payload={"result": "Should not be called"})
-        transport = ToolboxTransport(base_url, http_session, False)
+        transport = ToolboxTransport(base_url, http_session)
 
         tool_instance = ToolboxTool(
             transport=transport,
@@ -368,7 +368,7 @@ def test_tool_init_basic(http_session, sample_tool_params, sample_tool_descripti
     """Tests basic tool initialization without headers or auth."""
     with catch_warnings(record=True) as record:
         simplefilter("always")
-        transport = ToolboxTransport(HTTPS_BASE_URL, http_session, False)
+        transport = ToolboxTransport(HTTPS_BASE_URL, http_session)
 
         tool_instance = ToolboxTool(
             transport=transport,
@@ -398,7 +398,7 @@ def test_tool_init_with_client_headers(
     http_session, sample_tool_params, sample_tool_description, static_client_header
 ):
     """Tests tool initialization *with* client headers."""
-    transport = ToolboxTransport(HTTPS_BASE_URL, http_session, False)
+    transport = ToolboxTransport(HTTPS_BASE_URL, http_session)
     tool_instance = ToolboxTool(
         transport=transport,
         name=TEST_TOOL_NAME,
@@ -422,7 +422,7 @@ def test_tool_init_header_auth_conflict(
 ):
     """Tests ValueError on init if client header conflicts with auth token."""
     conflicting_client_header = {auth_header_key: "some-client-value"}
-    transport = ToolboxTransport(HTTPS_BASE_URL, http_session, False)
+    transport = ToolboxTransport(HTTPS_BASE_URL, http_session)
 
     with pytest.raises(
         ValueError, match=f"Client header\\(s\\) `{auth_header_key}` already registered"
@@ -449,7 +449,7 @@ def test_tool_add_auth_token_getters_conflict_with_existing_client_header(
     Tests ValueError when add_auth_token_getters introduces an auth service
     whose token name conflicts with an existing client header.
     """
-    transport = ToolboxTransport(HTTPS_BASE_URL, http_session, False)
+    transport = ToolboxTransport(HTTPS_BASE_URL, http_session)
     tool_instance = ToolboxTool(
         transport=transport,
         name="tool_with_client_header",
@@ -485,7 +485,7 @@ def test_add_auth_token_getters_unused_token(
     Tests ValueError when add_auth_token_getters is called with a getter for
     an unused authentication service.
     """
-    transport = ToolboxTransport(HTTPS_BASE_URL, http_session, False)
+    transport = ToolboxTransport(HTTPS_BASE_URL, http_session)
     tool_instance = ToolboxTool(
         transport=transport,
         name=TEST_TOOL_NAME,
@@ -514,7 +514,7 @@ def test_add_auth_token_getter_unused_token(
     Tests ValueError when add_auth_token_getters is called with a getter for
     an unused authentication service.
     """
-    transport = ToolboxTransport(HTTPS_BASE_URL, http_session, False)
+    transport = ToolboxTransport(HTTPS_BASE_URL, http_session)
     tool_instance = ToolboxTool(
         transport=transport,
         name=TEST_TOOL_NAME,
@@ -673,7 +673,7 @@ def test_tool_init_http_warning_when_sensitive_info_over_http(
         "Sending ID token over HTTP. User data may be exposed. "
         "Use HTTPS for secure communication."
     )
-    transport = ToolboxTransport(TEST_BASE_URL, http_session, False)
+    transport = ToolboxTransport(TEST_BASE_URL, http_session)
     init_kwargs = {
         "transport": transport,
         "name": "http_warning_tool",
@@ -704,7 +704,7 @@ def test_tool_init_no_http_warning_if_https(
     """
     with catch_warnings(record=True) as record:
         simplefilter("always")
-        transport = ToolboxTransport(HTTPS_BASE_URL, http_session, False)
+        transport = ToolboxTransport(HTTPS_BASE_URL, http_session)
 
         ToolboxTool(
             transport=transport,
@@ -733,7 +733,7 @@ def test_tool_init_no_http_warning_if_no_sensitive_info_on_http(
     """
     with catch_warnings(record=True) as record:
         simplefilter("always")
-        transport = ToolboxTransport(TEST_BASE_URL, http_session, False)
+        transport = ToolboxTransport(TEST_BASE_URL, http_session)
 
         ToolboxTool(
             transport=transport,

packages/toolbox-core/tests/test_toolbox_transport.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import AsyncGenerator, Union
+from typing import AsyncGenerator, Optional, Union
 from unittest.mock import AsyncMock
 
 import pytest
@@ -58,7 +58,7 @@ def mock_manifest_dict() -> dict:
 @pytest.mark.asyncio
 async def test_base_url_property(http_session: ClientSession):
     """Tests that the base_url property returns the correct URL."""
-    transport = ToolboxTransport(TEST_BASE_URL, http_session, False)
+    transport = ToolboxTransport(TEST_BASE_URL, http_session)
     assert transport.base_url == TEST_BASE_URL
 
 
@@ -67,7 +67,7 @@ async def test_tool_get_success(http_session: ClientSession, mock_manifest_dict:
     """Tests a successful tool_get call."""
     url = f"{TEST_BASE_URL}/api/tool/{TEST_TOOL_NAME}"
     headers = {"X-Test-Header": "value"}
-    transport = ToolboxTransport(TEST_BASE_URL, http_session, False)
+    transport = ToolboxTransport(TEST_BASE_URL, http_session)
 
     with aioresponses() as m:
         m.get(url, status=200, payload=mock_manifest_dict)
@@ -84,7 +84,7 @@ async def test_tool_get_success(http_session: ClientSession, mock_manifest_dict:
 async def test_tool_get_failure(http_session: ClientSession):
     """Tests a failing tool_get call and ensures it raises RuntimeError."""
     url = f"{TEST_BASE_URL}/api/tool/{TEST_TOOL_NAME}"
-    transport = ToolboxTransport(TEST_BASE_URL, http_session, False)
+    transport = ToolboxTransport(TEST_BASE_URL, http_session)
 
     with aioresponses() as m:
         m.get(url, status=500, body="Internal Server Error")
@@ -111,7 +111,7 @@ async def test_tools_list_success(
 ):
     """Tests successful tools_list calls with and without a toolset name."""
     url = f"{TEST_BASE_URL}{expected_path}"
-    transport = ToolboxTransport(TEST_BASE_URL, http_session, False)
+    transport = ToolboxTransport(TEST_BASE_URL, http_session)
 
     with aioresponses() as m:
         m.get(url, status=200, payload=mock_manifest_dict)
@@ -129,7 +129,7 @@ async def test_tool_invoke_success(http_session: ClientSession):
     args = {"param1": "value1"}
     headers = {"Authorization": "Bearer token"}
     response_payload = {"result": "success"}
-    transport = ToolboxTransport(TEST_BASE_URL, http_session, False)
+    transport = ToolboxTransport(TEST_BASE_URL, http_session)
 
     with aioresponses() as m:
         m.post(url, status=200, payload=response_payload)
@@ -144,7 +144,7 @@ async def test_tool_invoke_failure(http_session: ClientSession):
     """Tests a failing tool_invoke call where the server returns an error payload."""
     url = f"{TEST_BASE_URL}/api/tool/{TEST_TOOL_NAME}/invoke"
     response_payload = {"error": "Invalid arguments"}
-    transport = ToolboxTransport(TEST_BASE_URL, http_session, False)
+    transport = ToolboxTransport(TEST_BASE_URL, http_session)
 
     with aioresponses() as m:
         m.post(url, status=400, payload=response_payload)
@@ -155,25 +155,27 @@ async def test_tool_invoke_failure(http_session: ClientSession):
 
 
 @pytest.mark.asyncio
-@pytest.mark.parametrize(
-    "manage_session, is_closed, should_call_close",
-    [
-        (True, False, True),
-        (False, False, False),
-        (True, True, False),
-    ],
-)
-async def test_close_behavior(
-    manage_session: bool, is_closed: bool, should_call_close: bool
-):
-    """Tests the close method under different conditions."""
+async def test_close_does_not_close_unmanaged_session():
+    """
+    Tests that close() does NOT affect a session that was provided externally
+    (i.e., an unmanaged session).
+    """
     mock_session = AsyncMock(spec=ClientSession)
-    mock_session.closed = is_closed
-    transport = ToolboxTransport(TEST_BASE_URL, mock_session, manage_session)
+    mock_session.closed = False
 
+    transport = ToolboxTransport(TEST_BASE_URL, mock_session)
     await transport.close()
+    mock_session.close.assert_not_called()
 
-    if should_call_close:
-        mock_session.close.assert_awaited_once()
-    else:
-        mock_session.close.assert_not_awaited()
+
+@pytest.mark.asyncio
+async def test_close_closes_managed_session():
+    """
+    Tests that close() successfully closes a session that was created and
+    managed internally by the transport.
+    """
+    transport = ToolboxTransport(TEST_BASE_URL, session=None)
+
+    await transport.close()
+    internal_session = transport._ToolboxTransport__session
+    assert internal_session.closed is True