refactor: move transport logic to a ToolboxTransport class (#344)

* add basic code

* fixes

* test fix

* new unit tests

* rename ToolboxTransport

* add py3.9 support

* fix langchain tool tests

* test fix

* lint

* fix tests

* move manage session into transport

* move warning to diff file

* avoid code duplication

* fix tests

* lint

* remove redundant tests

* make invoke method return str

* lint

* fix return type

* small refactor

* rename private method

* fix tests

* lint

Author: twishabansal

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

@@ -20,14 +20,10 @@
 from deprecated import deprecated
 
 from .itransport import ITransport
-from .mcp_transport import (
-    McpHttpTransport_v20241105,
-    McpHttpTransport_v20250326,
-    McpHttpTransport_v20250618,
-)
-from .protocol import Protocol, ToolSchema
+from .protocol import ToolSchema
 from .tool import ToolboxTool
 from .toolbox_transport import ToolboxTransport
+from .toolbox_transport import ToolboxTransport
 from .utils import identify_auth_requirements, resolve_value
 
 
@@ -40,6 +36,7 @@ class ToolboxClient:
     is not provided.
     """
 
+    __transport: ITransport
     __transport: ITransport
 
     def __init__(
@@ -60,22 +57,10 @@ def __init__(
                 If None (default), a new session is created internally. Note that
                 if a session is provided, its lifecycle (including closing)
                 should typically be managed externally.
-            client_headers: Headers to include in each request sent through this
-            client.
-            protocol: The communication protocol to use.
+            client_headers: Headers to include in each request sent through this client.
         """
-        if protocol == Protocol.TOOLBOX:
-            self.__transport = ToolboxTransport(url, session)
-        elif protocol in Protocol.get_supported_mcp_versions():
-            if protocol == Protocol.MCP_v20250618:
-                self.__transport = McpHttpTransport_v20250618(url, session, protocol)
-            elif protocol == Protocol.MCP_v20250326:
-                self.__transport = McpHttpTransport_v20250326(url, session, protocol)
-            elif protocol == Protocol.MCP_v20241105:
-                self.__transport = McpHttpTransport_v20241105(url, session, protocol)
-        else:
-            raise ValueError(f"Unsupported MCP protocol version: {protocol}")
 
+        self.__transport = ToolboxTransport(url, session)
         self.__client_headers = client_headers if client_headers is not None else {}
 
     def __parse_tool(
@@ -114,6 +99,7 @@ def __parse_tool(
         )
 
         tool = ToolboxTool(
+            transport=self.__transport,
             transport=self.__transport,
             name=name,
             description=schema.description,
@@ -160,6 +146,7 @@ async def close(self):
         is responsible for its lifecycle.
         """
         await self.__transport.close()
+        await self.__transport.close()
 
     async def load_tool(
         self,
@@ -200,6 +187,7 @@ async def load_tool(
             for name, val in self.__client_headers.items()
         }
 
+        manifest = await self.__transport.tool_get(name, resolved_headers)
         manifest = await self.__transport.tool_get(name, resolved_headers)
 
         # parse the provided definition to a tool
@@ -277,6 +265,8 @@ async def load_toolset(
 
         manifest = await self.__transport.tools_list(name, resolved_headers)
 
+        manifest = await self.__transport.tools_list(name, resolved_headers)
+
         tools: list[ToolboxTool] = []
         overall_used_auth_keys: set[str] = set()
         overall_used_bound_params: set[str] = set()

packages/toolbox-core/tests/test_tool.py

@@ -29,6 +29,9 @@
 from toolbox_core.tool import ToolboxTool
 from toolbox_core.toolbox_transport import ToolboxTransport
 from toolbox_core.utils import create_func_docstring, resolve_value
+from toolbox_core.tool import ToolboxTool
+from toolbox_core.toolbox_transport import ToolboxTransport
+from toolbox_core.utils import create_func_docstring, resolve_value
 
 TEST_BASE_URL = "http://toolbox.example.com"
 HTTPS_BASE_URL = "https://toolbox.example.com"
@@ -113,7 +116,9 @@ def toolbox_tool(
 ) -> ToolboxTool:
     """Fixture for a ToolboxTool instance with common test setup."""
     transport = ToolboxTransport(TEST_BASE_URL, http_session)
+    transport = ToolboxTransport(TEST_BASE_URL, http_session)
     return ToolboxTool(
+        transport=transport,
         transport=transport,
         name=TEST_TOOL_NAME,
         description=sample_tool_description,
@@ -232,8 +237,10 @@ 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)
+        transport = ToolboxTransport(base_url, http_session)
 
         tool_instance = ToolboxTool(
+            transport=transport,
             transport=transport,
             name=tool_name,
             description=sample_tool_description,
@@ -278,8 +285,10 @@ 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)
+        transport = ToolboxTransport(base_url, http_session)
 
         tool_instance = ToolboxTool(
+            transport=transport,
             transport=transport,
             name=tool_name,
             description=sample_tool_description,
@@ -369,8 +378,10 @@ def test_tool_init_basic(http_session, sample_tool_params, sample_tool_descripti
     with catch_warnings(record=True) as record:
         simplefilter("always")
         transport = ToolboxTransport(HTTPS_BASE_URL, http_session)
+        transport = ToolboxTransport(HTTPS_BASE_URL, http_session)
 
         tool_instance = ToolboxTool(
+            transport=transport,
             transport=transport,
             name=TEST_TOOL_NAME,
             description=sample_tool_description,
@@ -399,7 +410,9 @@ def test_tool_init_with_client_headers(
 ):
     """Tests tool initialization *with* client headers."""
     transport = ToolboxTransport(HTTPS_BASE_URL, http_session)
+    transport = ToolboxTransport(HTTPS_BASE_URL, http_session)
     tool_instance = ToolboxTool(
+        transport=transport,
         transport=transport,
         name=TEST_TOOL_NAME,
         description=sample_tool_description,
@@ -423,7 +436,9 @@ def test_tool_add_auth_token_getters_conflict_with_existing_client_header(
     whose token name conflicts with an existing client header.
     """
     transport = ToolboxTransport(HTTPS_BASE_URL, http_session)
+    transport = ToolboxTransport(HTTPS_BASE_URL, http_session)
     tool_instance = ToolboxTool(
+        transport=transport,
         transport=transport,
         name="tool_with_client_header",
         description=sample_tool_description,
@@ -473,6 +488,21 @@ async def test_auth_token_overrides_client_header(
             "X-Another-Header": "another-value",
         },
     )
+    transport = ToolboxTransport(HTTPS_BASE_URL, http_session)
+    tool_instance = ToolboxTool(
+        transport=transport,
+        name=TEST_TOOL_NAME,
+        description=sample_tool_description,
+        params=sample_tool_params,
+        required_authn_params={},
+        required_authz_tokens=[],
+        auth_service_token_getters={"test-auth": lambda: "value-from-auth-getter-123"},
+        bound_params={},
+        client_headers={
+            "test-auth_token": "value-from-client",
+            "X-Another-Header": "another-value",
+        },
+    )
     tool_name = TEST_TOOL_NAME
     base_url = HTTPS_BASE_URL
     invoke_url = f"{base_url}/api/tool/{tool_name}/invoke"
@@ -490,6 +520,7 @@ async def test_auth_token_overrides_client_header(
             method="POST",
             json=input_args,
             headers={
+                "test-auth_token": "value-from-auth-getter-123",
                 "test-auth_token": "value-from-auth-getter-123",
                 "X-Another-Header": "another-value",
             },
@@ -507,7 +538,9 @@ def test_add_auth_token_getter_unused_token(
     an unused authentication service.
     """
     transport = ToolboxTransport(HTTPS_BASE_URL, http_session)
+    transport = ToolboxTransport(HTTPS_BASE_URL, http_session)
     tool_instance = ToolboxTool(
+        transport=transport,
         transport=transport,
         name=TEST_TOOL_NAME,
         description=sample_tool_description,
@@ -526,162 +559,3 @@ def test_add_auth_token_getter_unused_token(
             next(iter(unused_auth_getters)),
             unused_auth_getters[next(iter(unused_auth_getters))],
         )
-
-
-# --- Tests for Parameter Binding ---
-
-
-@pytest.mark.asyncio
-async def test_bind_param_success(
-    http_session: ClientSession,
-    sample_tool_params: list[ParameterSchema],
-    sample_tool_description: str,
-):
-    """
-    Tests successfully binding a single parameter with a static value using bind_param.
-    """
-    transport = ToolboxTransport(HTTPS_BASE_URL, http_session)
-    original_tool = ToolboxTool(
-        transport=transport,
-        name=TEST_TOOL_NAME,
-        description=sample_tool_description,
-        params=sample_tool_params,
-        required_authn_params={},
-        required_authz_tokens=[],
-        auth_service_token_getters={},
-        bound_params={},
-        client_headers={},
-    )
-
-    # Bind the 'count' parameter
-    bound_tool = original_tool.bind_param("count", 100)
-
-    # Assert immutability and changes
-    assert bound_tool is not original_tool
-    assert "count" not in bound_tool.__signature__.parameters
-    assert "message" in bound_tool.__signature__.parameters
-    assert "count" in original_tool.__signature__.parameters
-    assert bound_tool._bound_params == {"count": 100}
-    assert original_tool._bound_params == {}
-
-    # Test invocation of the new tool
-    invoke_url = f"{HTTPS_BASE_URL}/api/tool/{TEST_TOOL_NAME}/invoke"
-    with aioresponses() as m:
-        m.post(invoke_url, status=200, payload={"result": "Success"})
-        await bound_tool(message="hello")
-
-        # Verify the payload includes both the argument and the bound parameter
-        expected_payload = {"message": "hello", "count": 100}
-        m.assert_called_once_with(
-            invoke_url, method="POST", json=expected_payload, headers={}
-        )
-
-
-@pytest.mark.asyncio
-async def test_bind_params_success_with_callable(
-    http_session: ClientSession,
-    sample_tool_params: list[ParameterSchema],
-    sample_tool_description: str,
-):
-    """
-    Tests successfully binding multiple parameters, including one with a callable.
-    """
-    transport = ToolboxTransport(HTTPS_BASE_URL, http_session)
-    tool = ToolboxTool(
-        transport=transport,
-        name=TEST_TOOL_NAME,
-        description=sample_tool_description,
-        params=sample_tool_params,
-        required_authn_params={},
-        required_authz_tokens=[],
-        auth_service_token_getters={},
-        bound_params={},
-        client_headers={},
-    )
-
-    # Bind both parameters, one with a lambda
-    bound_tool = tool.bind_params({"message": lambda: "from-callable", "count": 99})
-
-    assert "message" not in bound_tool.__signature__.parameters
-    assert "count" not in bound_tool.__signature__.parameters
-    assert len(bound_tool.__signature__.parameters) == 0
-
-    # Test invocation
-    invoke_url = f"{HTTPS_BASE_URL}/api/tool/{TEST_TOOL_NAME}/invoke"
-    with aioresponses() as m:
-        m.post(invoke_url, status=200, payload={"result": "Success"})
-        await bound_tool()  # Call with no arguments
-
-        expected_payload = {"message": "from-callable", "count": 99}
-        m.assert_called_once_with(
-            invoke_url, method="POST", json=expected_payload, headers={}
-        )
-
-
-def test_bind_param_invalid_parameter_name(toolbox_tool: ToolboxTool):
-    """
-    Tests that binding a parameter that does not exist raises a ValueError.
-    """
-    with pytest.raises(
-        ValueError, match="unable to bind parameters: no parameter named invalid_param"
-    ):
-        toolbox_tool.bind_param("invalid_param", "some_value")
-
-
-def test_bind_params_rebinding_parameter_fails(toolbox_tool: ToolboxTool):
-    """
-    Tests that attempting to re-bind an already bound parameter raises a ValueError.
-    """
-    tool_with_one_bound_param = toolbox_tool.bind_param("count", 50)
-
-    with pytest.raises(
-        ValueError, match="cannot re-bind parameter: parameter 'count' is already bound"
-    ):
-        tool_with_one_bound_param.bind_params({"count": 75})
-
-
-@pytest.mark.asyncio
-async def test_bind_param_chaining(
-    http_session: ClientSession,
-    sample_tool_params: list[ParameterSchema],
-    sample_tool_description: str,
-):
-    """
-    Tests that bind_param calls can be chained to bind multiple parameters sequentially.
-    """
-    transport = ToolboxTransport(HTTPS_BASE_URL, http_session)
-    tool = ToolboxTool(
-        transport=transport,
-        name=TEST_TOOL_NAME,
-        description=sample_tool_description,
-        params=sample_tool_params,
-        required_authn_params={},
-        required_authz_tokens=[],
-        auth_service_token_getters={},
-        bound_params={},
-        client_headers={},
-    )
-
-    # Chain the calls
-    fully_bound_tool = tool.bind_param("count", 42).bind_param(
-        "message", "chained-call"
-    )
-
-    assert len(fully_bound_tool.__signature__.parameters) == 0
-    assert fully_bound_tool._bound_params == {
-        "count": 42,
-        "message": "chained-call",
-    }
-
-    # Test invocation
-    invoke_url = f"{HTTPS_BASE_URL}/api/tool/{TEST_TOOL_NAME}/invoke"
-    with aioresponses() as m:
-        m.post(invoke_url, status=200, payload={"result": "Success"})
-        await fully_bound_tool()
-
-        m.assert_called_once_with(
-            invoke_url,
-            method="POST",
-            json={"count": 42, "message": "chained-call"},
-            headers={},
-        )