From 04f7962ca9b5d557200983b192c6cec031c45501 Mon Sep 17 00:00:00 2001 From: Twisha Bansal Date: Mon, 25 Aug 2025 18:08:39 +0530 Subject: [PATCH 01/23] add basic code --- .../toolbox-core/src/toolbox_core/client.py | 45 ++++--------- .../src/toolbox_core/itransport.py | 45 +++++++++++++ .../toolbox-core/src/toolbox_core/tool.py | 40 ++++-------- .../src/toolbox_core/toolbox_transport.py | 63 +++++++++++++++++++ 4 files changed, 133 insertions(+), 60 deletions(-) create mode 100644 packages/toolbox-core/src/toolbox_core/itransport.py create mode 100644 packages/toolbox-core/src/toolbox_core/toolbox_transport.py diff --git a/packages/toolbox-core/src/toolbox_core/client.py b/packages/toolbox-core/src/toolbox_core/client.py index 3f55f7181..302fd30d6 100644 --- a/packages/toolbox-core/src/toolbox_core/client.py +++ b/packages/toolbox-core/src/toolbox_core/client.py @@ -19,8 +19,10 @@ from aiohttp import ClientSession from deprecated import deprecated -from .protocol import ManifestSchema, ToolSchema +from .itransport import ITransport +from .protocol import ToolSchema from .tool import ToolboxTool +from .toolbox_transport import ToolboxHttpTransport from .utils import identify_auth_requirements, resolve_value @@ -33,9 +35,7 @@ class ToolboxClient: is not provided. """ - __base_url: str - __session: ClientSession - __manage_session: bool + __transport: ITransport def __init__( self, @@ -56,15 +56,13 @@ def __init__( should typically be managed externally. client_headers: Headers to include in each request sent through this client. """ - self.__base_url = url # If no aiohttp.ClientSession is provided, make our own - self.__manage_session = False + manage_session = False if session is None: - self.__manage_session = True + manage_session = True session = ClientSession() - self.__session = session - + self.__transport = ToolboxHttpTransport(url, session, manage_session) self.__client_headers = client_headers if client_headers is not None else {} def __parse_tool( @@ -103,8 +101,7 @@ def __parse_tool( ) tool = ToolboxTool( - session=self.__session, - base_url=self.__base_url, + transport=self.__transport, name=name, description=schema.description, # create a read-only values to prevent mutation @@ -149,8 +146,7 @@ async def close(self): If the session was provided externally during initialization, the caller is responsible for its lifecycle. """ - if self.__manage_session and not self.__session.closed: - await self.__session.close() + await self.__transport.close() async def load_tool( self, @@ -191,16 +187,7 @@ async def load_tool( for name, val in self.__client_headers.items() } - # request the definition of the tool from the server - url = f"{self.__base_url}/api/tool/{name}" - async with self.__session.get(url, headers=resolved_headers) as response: - if not response.ok: - error_text = await response.text() - raise RuntimeError( - f"API request failed with status {response.status} ({response.reason}). Server response: {error_text}" - ) - json = await response.json() - manifest: ManifestSchema = ManifestSchema(**json) + manifest = await self.__transport.tools_list(name, resolved_headers) # parse the provided definition to a tool if name not in manifest.tools: @@ -274,16 +261,8 @@ async def load_toolset( header_name: await resolve_value(original_headers[header_name]) for header_name in original_headers } - # Request the definition of the toolset from the server - url = f"{self.__base_url}/api/toolset/{name or ''}" - async with self.__session.get(url, headers=resolved_headers) as response: - if not response.ok: - error_text = await response.text() - raise RuntimeError( - f"API request failed with status {response.status} ({response.reason}). Server response: {error_text}" - ) - json = await response.json() - manifest: ManifestSchema = ManifestSchema(**json) + + manifest = await self.__transport.tools_list(name, resolved_headers) tools: list[ToolboxTool] = [] overall_used_auth_keys: set[str] = set() diff --git a/packages/toolbox-core/src/toolbox_core/itransport.py b/packages/toolbox-core/src/toolbox_core/itransport.py new file mode 100644 index 000000000..ef9e546ff --- /dev/null +++ b/packages/toolbox-core/src/toolbox_core/itransport.py @@ -0,0 +1,45 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from abc import ABC, abstractmethod +from typing import Mapping, Optional + +from .protocol import ManifestSchema + + +class ITransport(ABC): + """Defines the contract for a 'smart' transport that handles both + protocol formatting and network communication. + """ + + @abstractmethod + async def tools_list( + self, + toolset_name: Optional[str] = None, + headers: Optional[Mapping[str, str]] = None, + ) -> ManifestSchema: + """Lists available tools from the server.""" + pass + + @abstractmethod + async def tool_invoke( + self, tool_name: str, arguments: dict, headers: Mapping[str, str] + ) -> dict: + """Invokes a specific tool on the server.""" + pass + + @abstractmethod + async def close(self): + """Closes any underlying connections.""" + pass diff --git a/packages/toolbox-core/src/toolbox_core/tool.py b/packages/toolbox-core/src/toolbox_core/tool.py index d9ec28839..f93cb98bd 100644 --- a/packages/toolbox-core/src/toolbox_core/tool.py +++ b/packages/toolbox-core/src/toolbox_core/tool.py @@ -20,8 +20,7 @@ from typing import Any, Awaitable, Callable, Mapping, Optional, Sequence, Union from warnings import warn -from aiohttp import ClientSession - +from .itransport import ITransport from .protocol import ParameterSchema from .utils import ( create_func_docstring, @@ -46,8 +45,7 @@ class ToolboxTool: def __init__( self, - session: ClientSession, - base_url: str, + transport: ITransport, name: str, description: str, params: Sequence[ParameterSchema], @@ -68,8 +66,7 @@ def __init__( Toolbox server. Args: - session: The `aiohttp.ClientSession` used for making API requests. - base_url: The base URL of the Toolbox server API. + transport: The transport used for making API requests. name: The name of the remote tool. description: The description of the remote tool. params: The args of the tool. @@ -84,9 +81,7 @@ def __init__( client_headers: Client specific headers bound to the tool. """ # used to invoke the toolbox API - self.__session: ClientSession = session - self.__base_url: str = base_url - self.__url = f"{base_url}/api/tool/{name}/invoke" + self.__transport = transport self.__description = description self.__params = params self.__pydantic_model = params_to_pydantic_model(name, self.__params) @@ -137,9 +132,7 @@ def __init__( # these over HTTP exposes the data to interception and unauthorized # access. Always use HTTPS to ensure secure communication and protect # user privacy. - if ( - required_authn_params or required_authz_tokens or client_headers - ) and not self.__url.startswith("https://"): + if required_authn_params or required_authz_tokens or client_headers: warn( "Sending ID token over HTTP. User data may be exposed. Use HTTPS for secure communication." ) @@ -184,8 +177,7 @@ def _client_headers( def __copy( self, - session: Optional[ClientSession] = None, - base_url: Optional[str] = None, + transport: Optional[ITransport] = None, name: Optional[str] = None, description: Optional[str] = None, params: Optional[Sequence[ParameterSchema]] = None, @@ -205,8 +197,7 @@ def __copy( Creates a copy of the ToolboxTool, overriding specific fields. Args: - session: The `aiohttp.ClientSession` used for making API requests. - base_url: The base URL of the Toolbox server API. + transport: The transport used for making API requests. name: The name of the remote tool. description: The description of the remote tool. params: The args of the tool. @@ -222,8 +213,7 @@ def __copy( """ check = lambda val, default: val if val is not None else default return ToolboxTool( - session=check(session, self.__session), - base_url=check(base_url, self.__base_url), + transport=check(transport, self.__transport), name=check(name, self.__name__), description=check(description, self.__description), params=check(params, self.__params), @@ -304,15 +294,11 @@ async def __call__(self, *args: Any, **kwargs: Any) -> str: token_getter ) - async with self.__session.post( - self.__url, - json=payload, - headers=headers, - ) as resp: - body = await resp.json() - if not resp.ok: - err = body.get("error", f"unexpected status from server: {resp.status}") - raise Exception(err) + body = await self.__transport.tool_invoke( + self.__name__, + payload, + headers, + ) return body.get("result", body) def add_auth_token_getters( diff --git a/packages/toolbox-core/src/toolbox_core/toolbox_transport.py b/packages/toolbox-core/src/toolbox_core/toolbox_transport.py new file mode 100644 index 000000000..a95068f47 --- /dev/null +++ b/packages/toolbox-core/src/toolbox_core/toolbox_transport.py @@ -0,0 +1,63 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Mapping, Optional + +from aiohttp import ClientSession + +from .itransport import ITransport +from .protocol import ManifestSchema + + +class ToolboxHttpTransport(ITransport): + """Transport for the native Toolbox protocol.""" + + def __init__(self, base_url: str, session: ClientSession, manage_session: bool): + self.__base_url = base_url + self.__session = session + self.__manage_session = manage_session + + async def tools_list( + self, + toolset_name: Optional[str] = None, + headers: Optional[Mapping[str, str]] = None, + ) -> ManifestSchema: + url = f"{self.__base_url}/api/toolset/{toolset_name or ''}" + async with self.__session.get(url, headers=headers) as response: + if not response.ok: + error_text = await response.text() + raise RuntimeError( + f"API request failed with status {response.status} ({response.reason}). Server response: {error_text}" + ) + json = await response.json() + return ManifestSchema(**json) + + async def tool_invoke( + self, tool_name: str, arguments: dict, headers: Mapping[str, str] + ) -> dict: + url = f"{self.__base_url}/api/tool/{tool_name}/invoke" + async with self.__session.post( + url, + json=arguments, + headers=headers, + ) as resp: + body = await resp.json() + if not resp.ok: + err = body.get("error", f"unexpected status from server: {resp.status}") + raise Exception(err) + return body + + async def close(self): + if self.__manage_session and not self.__session.closed: + await self.__session.close() From 9ba5bbcab7bbd9b4a898a4b9eb88a1f992ea8cd3 Mon Sep 17 00:00:00 2001 From: Twisha Bansal Date: Mon, 25 Aug 2025 18:24:46 +0530 Subject: [PATCH 02/23] fixes --- packages/toolbox-core/src/toolbox_core/client.py | 2 +- .../src/toolbox_core/toolbox_transport.py | 13 +++++++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/packages/toolbox-core/src/toolbox_core/client.py b/packages/toolbox-core/src/toolbox_core/client.py index 302fd30d6..649617e7b 100644 --- a/packages/toolbox-core/src/toolbox_core/client.py +++ b/packages/toolbox-core/src/toolbox_core/client.py @@ -187,7 +187,7 @@ async def load_tool( for name, val in self.__client_headers.items() } - manifest = await self.__transport.tools_list(name, resolved_headers) + manifest = await self.__transport.tool_get(name, resolved_headers) # parse the provided definition to a tool if name not in manifest.tools: diff --git a/packages/toolbox-core/src/toolbox_core/toolbox_transport.py b/packages/toolbox-core/src/toolbox_core/toolbox_transport.py index a95068f47..96388b919 100644 --- a/packages/toolbox-core/src/toolbox_core/toolbox_transport.py +++ b/packages/toolbox-core/src/toolbox_core/toolbox_transport.py @@ -28,6 +28,19 @@ def __init__(self, base_url: str, session: ClientSession, manage_session: bool): self.__session = session self.__manage_session = manage_session + async def tool_get( + self, tool_name: str, headers: Optional[Mapping[str, str]] = None + ) -> ManifestSchema: + url = f"{self.__base_url}/api/tool/{tool_name}" + async with self.__session.get(url, headers=headers) as response: + if not response.ok: + error_text = await response.text() + raise RuntimeError( + f"API request failed with status {response.status} ({response.reason}). Server response: {error_text}" + ) + json = await response.json() + return ManifestSchema(**json) + async def tools_list( self, toolset_name: Optional[str] = None, From 34cb50a6ff2decf466a11051e8cd6a6083b29b17 Mon Sep 17 00:00:00 2001 From: Twisha Bansal Date: Mon, 25 Aug 2025 19:24:27 +0530 Subject: [PATCH 03/23] test fix --- .../src/toolbox_core/itransport.py | 13 +++++ .../toolbox-core/src/toolbox_core/tool.py | 4 +- .../src/toolbox_core/toolbox_transport.py | 5 ++ packages/toolbox-core/tests/test_tool.py | 55 ++++++++++--------- 4 files changed, 49 insertions(+), 28 deletions(-) diff --git a/packages/toolbox-core/src/toolbox_core/itransport.py b/packages/toolbox-core/src/toolbox_core/itransport.py index ef9e546ff..7d4a87e5a 100644 --- a/packages/toolbox-core/src/toolbox_core/itransport.py +++ b/packages/toolbox-core/src/toolbox_core/itransport.py @@ -23,6 +23,19 @@ class ITransport(ABC): protocol formatting and network communication. """ + @property + @abstractmethod + def base_url(self) -> str: + """The base URL for the transport.""" + pass + + @abstractmethod + async def tool_get( + self, tool_name: str, headers: Optional[Mapping[str, str]] = None + ) -> ManifestSchema: + """Gets a single tool from the server.""" + pass + @abstractmethod async def tools_list( self, diff --git a/packages/toolbox-core/src/toolbox_core/tool.py b/packages/toolbox-core/src/toolbox_core/tool.py index f93cb98bd..aa6ea260c 100644 --- a/packages/toolbox-core/src/toolbox_core/tool.py +++ b/packages/toolbox-core/src/toolbox_core/tool.py @@ -132,7 +132,9 @@ def __init__( # these over HTTP exposes the data to interception and unauthorized # access. Always use HTTPS to ensure secure communication and protect # user privacy. - if required_authn_params or required_authz_tokens or client_headers: + if self.__transport.base_url.startswith("http://") and ( + required_authn_params or required_authz_tokens or client_headers + ): warn( "Sending ID token over HTTP. User data may be exposed. Use HTTPS for secure communication." ) diff --git a/packages/toolbox-core/src/toolbox_core/toolbox_transport.py b/packages/toolbox-core/src/toolbox_core/toolbox_transport.py index 96388b919..ed62e407e 100644 --- a/packages/toolbox-core/src/toolbox_core/toolbox_transport.py +++ b/packages/toolbox-core/src/toolbox_core/toolbox_transport.py @@ -28,6 +28,11 @@ def __init__(self, base_url: str, session: ClientSession, manage_session: bool): self.__session = session self.__manage_session = manage_session + @property + def base_url(self) -> str: + """The base URL for the transport.""" + return self.__base_url + async def tool_get( self, tool_name: str, headers: Optional[Mapping[str, str]] = None ) -> ManifestSchema: diff --git a/packages/toolbox-core/tests/test_tool.py b/packages/toolbox-core/tests/test_tool.py index 2e918c3d9..5a52cc941 100644 --- a/packages/toolbox-core/tests/test_tool.py +++ b/packages/toolbox-core/tests/test_tool.py @@ -26,7 +26,9 @@ from pydantic import ValidationError from toolbox_core.protocol import ParameterSchema -from toolbox_core.tool import ToolboxTool, create_func_docstring, resolve_value +from toolbox_core.tool import ToolboxTool +from toolbox_core.toolbox_transport import ToolboxHttpTransport +from toolbox_core.utils import create_func_docstring, resolve_value TEST_BASE_URL = "http://toolbox.example.com" HTTPS_BASE_URL = "https://toolbox.example.com" @@ -110,9 +112,9 @@ def toolbox_tool( sample_tool_description: str, ) -> ToolboxTool: """Fixture for a ToolboxTool instance with common test setup.""" + transport = ToolboxHttpTransport(TEST_BASE_URL, http_session, False) return ToolboxTool( - session=http_session, - base_url=TEST_BASE_URL, + transport=transport, name=TEST_TOOL_NAME, description=sample_tool_description, params=sample_tool_params, @@ -229,10 +231,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 = ToolboxHttpTransport(base_url, http_session, False) tool_instance = ToolboxTool( - session=http_session, - base_url=base_url, + transport=transport, name=tool_name, description=sample_tool_description, params=sample_tool_params, @@ -275,10 +277,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 = ToolboxHttpTransport(base_url, http_session, False) tool_instance = ToolboxTool( - session=http_session, - base_url=base_url, + transport=transport, name=tool_name, description=sample_tool_description, params=sample_tool_params, @@ -366,10 +368,10 @@ 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 = ToolboxHttpTransport(HTTPS_BASE_URL, http_session, False) tool_instance = ToolboxTool( - session=http_session, - base_url=HTTPS_BASE_URL, + transport=transport, name=TEST_TOOL_NAME, description=sample_tool_description, params=sample_tool_params, @@ -396,9 +398,9 @@ 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 = ToolboxHttpTransport(HTTPS_BASE_URL, http_session, False) tool_instance = ToolboxTool( - session=http_session, - base_url=HTTPS_BASE_URL, + transport=transport, name=TEST_TOOL_NAME, description=sample_tool_description, params=sample_tool_params, @@ -420,13 +422,13 @@ 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 = ToolboxHttpTransport(HTTPS_BASE_URL, http_session, False) with pytest.raises( ValueError, match=f"Client header\\(s\\) `{auth_header_key}` already registered" ): ToolboxTool( - session=http_session, - base_url=HTTPS_BASE_URL, + transport=transport, name="auth_conflict_tool", description=sample_tool_description, params=sample_tool_auth_params, @@ -447,9 +449,9 @@ 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 = ToolboxHttpTransport(HTTPS_BASE_URL, http_session, False) tool_instance = ToolboxTool( - session=http_session, - base_url=HTTPS_BASE_URL, + transport=transport, name="tool_with_client_header", description=sample_tool_description, params=sample_tool_params, @@ -483,9 +485,9 @@ 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 = ToolboxHttpTransport(HTTPS_BASE_URL, http_session, False) tool_instance = ToolboxTool( - session=http_session, - base_url=HTTPS_BASE_URL, + transport=transport, name=TEST_TOOL_NAME, description=sample_tool_description, params=sample_tool_params, @@ -512,9 +514,9 @@ 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 = ToolboxHttpTransport(HTTPS_BASE_URL, http_session, False) tool_instance = ToolboxTool( - session=http_session, - base_url=HTTPS_BASE_URL, + transport=transport, name=TEST_TOOL_NAME, description=sample_tool_description, params=sample_tool_params, @@ -667,14 +669,13 @@ def test_tool_init_http_warning_when_sensitive_info_over_http( Tests that a UserWarning is issued if client headers, auth params, or auth tokens are present and the base_url is HTTP. """ - expected_warning_message = ( + expected_warning_message: str = ( "Sending ID token over HTTP. User data may be exposed. " "Use HTTPS for secure communication." ) - + transport = ToolboxHttpTransport(TEST_BASE_URL, http_session, False) init_kwargs = { - "session": http_session, - "base_url": TEST_BASE_URL, + "transport": transport, "name": "http_warning_tool", "description": sample_tool_description, "params": sample_tool_params, @@ -703,10 +704,10 @@ def test_tool_init_no_http_warning_if_https( """ with catch_warnings(record=True) as record: simplefilter("always") + transport = ToolboxHttpTransport(HTTPS_BASE_URL, http_session, False) ToolboxTool( - session=http_session, - base_url=HTTPS_BASE_URL, + transport=transport, name="https_tool", description=sample_tool_description, params=sample_tool_params, @@ -732,10 +733,10 @@ def test_tool_init_no_http_warning_if_no_sensitive_info_on_http( """ with catch_warnings(record=True) as record: simplefilter("always") + transport = ToolboxHttpTransport(TEST_BASE_URL, http_session, False) ToolboxTool( - session=http_session, - base_url=TEST_BASE_URL, + transport=transport, name="http_tool_no_sensitive", description=sample_tool_description, params=sample_tool_params, From ea417bacac699443e8c23e179bbd2b3b85a120cc Mon Sep 17 00:00:00 2001 From: Twisha Bansal Date: Mon, 25 Aug 2025 19:33:38 +0530 Subject: [PATCH 04/23] new unit tests --- .../tests/test_toolbox_transport.py | 179 ++++++++++++++++++ 1 file changed, 179 insertions(+) create mode 100644 packages/toolbox-core/tests/test_toolbox_transport.py diff --git a/packages/toolbox-core/tests/test_toolbox_transport.py b/packages/toolbox-core/tests/test_toolbox_transport.py new file mode 100644 index 000000000..7ec012363 --- /dev/null +++ b/packages/toolbox-core/tests/test_toolbox_transport.py @@ -0,0 +1,179 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import AsyncGenerator +from unittest.mock import AsyncMock + +import pytest +import pytest_asyncio +from aiohttp import ClientSession +from aioresponses import aioresponses + +from toolbox_core.protocol import ManifestSchema +from toolbox_core.toolbox_transport import ToolboxHttpTransport + +TEST_BASE_URL = "http://fake-toolbox-server.com" +TEST_TOOL_NAME = "test_tool" + + +@pytest_asyncio.fixture +async def http_session() -> AsyncGenerator[ClientSession, None]: + """Provides a real aiohttp ClientSession that is closed after the test.""" + async with ClientSession() as session: + yield session + + +@pytest.fixture +def mock_manifest_dict() -> dict: + """Provides a valid sample dictionary for a ManifestSchema response.""" + tool_definition = { + "name": TEST_TOOL_NAME, + "description": "A test tool", + "parameters": [ + { + "name": "param1", + "type": "string", + "description": "The first parameter.", + "required": True, + } + ], + } + return { + "serverVersion": "1.0.0", + "tools": {TEST_TOOL_NAME: tool_definition}, + } + + +@pytest.mark.asyncio +async def test_base_url_property(http_session: ClientSession): + """Tests that the base_url property returns the correct URL.""" + transport = ToolboxHttpTransport(TEST_BASE_URL, http_session, False) + assert transport.base_url == TEST_BASE_URL + + +@pytest.mark.asyncio +async def test_tool_get_success(http_session: ClientSession, mock_manifest_dict: dict): + """Tests a successful tool_get call.""" + url = f"{TEST_BASE_URL}/api/tool/{TEST_TOOL_NAME}" + headers = {"X-Test-Header": "value"} + transport = ToolboxHttpTransport(TEST_BASE_URL, http_session, False) + + with aioresponses() as m: + m.get(url, status=200, payload=mock_manifest_dict) + result = await transport.tool_get(TEST_TOOL_NAME, headers=headers) + + assert isinstance(result, ManifestSchema) + assert result.serverVersion == "1.0.0" + # FIX: Check for a valid attribute like 'description' instead of 'name' + assert result.tools[TEST_TOOL_NAME].description == "A test tool" + m.assert_called_once_with(url, headers=headers) + + +@pytest.mark.asyncio +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 = ToolboxHttpTransport(TEST_BASE_URL, http_session, False) + + with aioresponses() as m: + m.get(url, status=500, body="Internal Server Error") + with pytest.raises(RuntimeError) as exc_info: + await transport.tool_get(TEST_TOOL_NAME) + + assert "API request failed with status 500" in str(exc_info.value) + assert "Internal Server Error" in str(exc_info.value) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "toolset_name, expected_path", + [ + ("my_toolset", "/api/toolset/my_toolset"), + (None, "/api/toolset/"), + ], +) +async def test_tools_list_success( + http_session: ClientSession, + mock_manifest_dict: dict, + toolset_name: str | None, + expected_path: str, +): + """Tests successful tools_list calls with and without a toolset name.""" + url = f"{TEST_BASE_URL}{expected_path}" + transport = ToolboxHttpTransport(TEST_BASE_URL, http_session, False) + + with aioresponses() as m: + m.get(url, status=200, payload=mock_manifest_dict) + result = await transport.tools_list(toolset_name=toolset_name) + + assert isinstance(result, ManifestSchema) + # FIX: Add headers=None to match the actual call signature + m.assert_called_once_with(url, headers=None) + + +@pytest.mark.asyncio +async def test_tool_invoke_success(http_session: ClientSession): + """Tests a successful tool_invoke call.""" + url = f"{TEST_BASE_URL}/api/tool/{TEST_TOOL_NAME}/invoke" + args = {"param1": "value1"} + headers = {"Authorization": "Bearer token"} + response_payload = {"result": "success"} + transport = ToolboxHttpTransport(TEST_BASE_URL, http_session, False) + + with aioresponses() as m: + m.post(url, status=200, payload=response_payload) + result = await transport.tool_invoke(TEST_TOOL_NAME, args, headers) + + assert result == response_payload + m.assert_called_once_with(url, method="POST", json=args, headers=headers) + + +@pytest.mark.asyncio +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 = ToolboxHttpTransport(TEST_BASE_URL, http_session, False) + + with aioresponses() as m: + m.post(url, status=400, payload=response_payload) + with pytest.raises(Exception) as exc_info: + await transport.tool_invoke(TEST_TOOL_NAME, {}, {}) + + assert str(exc_info.value) == "Invalid arguments" + + +@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.""" + mock_session = AsyncMock(spec=ClientSession) + mock_session.closed = is_closed + transport = ToolboxHttpTransport(TEST_BASE_URL, mock_session, manage_session) + + await transport.close() + + if should_call_close: + mock_session.close.assert_awaited_once() + else: + mock_session.close.assert_not_awaited() From 7c0f9172a2801f04788b3a166bb0fff1a7a45474 Mon Sep 17 00:00:00 2001 From: Twisha Bansal Date: Mon, 25 Aug 2025 19:40:34 +0530 Subject: [PATCH 05/23] rename ToolboxTransport --- .../toolbox-core/src/toolbox_core/client.py | 4 +-- .../src/toolbox_core/toolbox_transport.py | 2 +- packages/toolbox-core/tests/test_tool.py | 26 +++++++++---------- .../tests/test_toolbox_transport.py | 16 ++++++------ 4 files changed, 24 insertions(+), 24 deletions(-) diff --git a/packages/toolbox-core/src/toolbox_core/client.py b/packages/toolbox-core/src/toolbox_core/client.py index 649617e7b..5db0c957b 100644 --- a/packages/toolbox-core/src/toolbox_core/client.py +++ b/packages/toolbox-core/src/toolbox_core/client.py @@ -22,7 +22,7 @@ from .itransport import ITransport from .protocol import ToolSchema from .tool import ToolboxTool -from .toolbox_transport import ToolboxHttpTransport +from .toolbox_transport import ToolboxTransport from .utils import identify_auth_requirements, resolve_value @@ -62,7 +62,7 @@ def __init__( if session is None: manage_session = True session = ClientSession() - self.__transport = ToolboxHttpTransport(url, session, manage_session) + self.__transport = ToolboxTransport(url, session, manage_session) self.__client_headers = client_headers if client_headers is not None else {} def __parse_tool( diff --git a/packages/toolbox-core/src/toolbox_core/toolbox_transport.py b/packages/toolbox-core/src/toolbox_core/toolbox_transport.py index ed62e407e..06b61bfc2 100644 --- a/packages/toolbox-core/src/toolbox_core/toolbox_transport.py +++ b/packages/toolbox-core/src/toolbox_core/toolbox_transport.py @@ -20,7 +20,7 @@ from .protocol import ManifestSchema -class ToolboxHttpTransport(ITransport): +class ToolboxTransport(ITransport): """Transport for the native Toolbox protocol.""" def __init__(self, base_url: str, session: ClientSession, manage_session: bool): diff --git a/packages/toolbox-core/tests/test_tool.py b/packages/toolbox-core/tests/test_tool.py index 5a52cc941..449b94405 100644 --- a/packages/toolbox-core/tests/test_tool.py +++ b/packages/toolbox-core/tests/test_tool.py @@ -27,7 +27,7 @@ from toolbox_core.protocol import ParameterSchema from toolbox_core.tool import ToolboxTool -from toolbox_core.toolbox_transport import ToolboxHttpTransport +from toolbox_core.toolbox_transport import ToolboxTransport from toolbox_core.utils import create_func_docstring, resolve_value TEST_BASE_URL = "http://toolbox.example.com" @@ -112,7 +112,7 @@ def toolbox_tool( sample_tool_description: str, ) -> ToolboxTool: """Fixture for a ToolboxTool instance with common test setup.""" - transport = ToolboxHttpTransport(TEST_BASE_URL, http_session, False) + transport = ToolboxTransport(TEST_BASE_URL, http_session, False) 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 = ToolboxHttpTransport(base_url, http_session, False) + transport = ToolboxTransport(base_url, http_session, False) 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 = ToolboxHttpTransport(base_url, http_session, False) + transport = ToolboxTransport(base_url, http_session, False) 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 = ToolboxHttpTransport(HTTPS_BASE_URL, http_session, False) + transport = ToolboxTransport(HTTPS_BASE_URL, http_session, False) 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 = ToolboxHttpTransport(HTTPS_BASE_URL, http_session, False) + transport = ToolboxTransport(HTTPS_BASE_URL, http_session, False) 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 = ToolboxHttpTransport(HTTPS_BASE_URL, http_session, False) + transport = ToolboxTransport(HTTPS_BASE_URL, http_session, False) 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 = ToolboxHttpTransport(HTTPS_BASE_URL, http_session, False) + transport = ToolboxTransport(HTTPS_BASE_URL, http_session, False) 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 = ToolboxHttpTransport(HTTPS_BASE_URL, http_session, False) + transport = ToolboxTransport(HTTPS_BASE_URL, http_session, False) 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 = ToolboxHttpTransport(HTTPS_BASE_URL, http_session, False) + transport = ToolboxTransport(HTTPS_BASE_URL, http_session, False) 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 = ToolboxHttpTransport(TEST_BASE_URL, http_session, False) + transport = ToolboxTransport(TEST_BASE_URL, http_session, False) 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 = ToolboxHttpTransport(HTTPS_BASE_URL, http_session, False) + transport = ToolboxTransport(HTTPS_BASE_URL, http_session, False) 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 = ToolboxHttpTransport(TEST_BASE_URL, http_session, False) + transport = ToolboxTransport(TEST_BASE_URL, http_session, False) ToolboxTool( transport=transport, diff --git a/packages/toolbox-core/tests/test_toolbox_transport.py b/packages/toolbox-core/tests/test_toolbox_transport.py index 7ec012363..ca2d7d79d 100644 --- a/packages/toolbox-core/tests/test_toolbox_transport.py +++ b/packages/toolbox-core/tests/test_toolbox_transport.py @@ -21,7 +21,7 @@ from aioresponses import aioresponses from toolbox_core.protocol import ManifestSchema -from toolbox_core.toolbox_transport import ToolboxHttpTransport +from toolbox_core.toolbox_transport import ToolboxTransport TEST_BASE_URL = "http://fake-toolbox-server.com" TEST_TOOL_NAME = "test_tool" @@ -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 = ToolboxHttpTransport(TEST_BASE_URL, http_session, False) + transport = ToolboxTransport(TEST_BASE_URL, http_session, False) 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 = ToolboxHttpTransport(TEST_BASE_URL, http_session, False) + transport = ToolboxTransport(TEST_BASE_URL, http_session, False) 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 = ToolboxHttpTransport(TEST_BASE_URL, http_session, False) + transport = ToolboxTransport(TEST_BASE_URL, http_session, False) 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 = ToolboxHttpTransport(TEST_BASE_URL, http_session, False) + transport = ToolboxTransport(TEST_BASE_URL, http_session, False) 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 = ToolboxHttpTransport(TEST_BASE_URL, http_session, False) + transport = ToolboxTransport(TEST_BASE_URL, http_session, False) 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 = ToolboxHttpTransport(TEST_BASE_URL, http_session, False) + transport = ToolboxTransport(TEST_BASE_URL, http_session, False) with aioresponses() as m: m.post(url, status=400, payload=response_payload) @@ -169,7 +169,7 @@ async def test_close_behavior( """Tests the close method under different conditions.""" mock_session = AsyncMock(spec=ClientSession) mock_session.closed = is_closed - transport = ToolboxHttpTransport(TEST_BASE_URL, mock_session, manage_session) + transport = ToolboxTransport(TEST_BASE_URL, mock_session, manage_session) await transport.close() From 434e91d43b70c9f5720a506ab6e63116bffe38c2 Mon Sep 17 00:00:00 2001 From: Twisha Bansal Date: Mon, 25 Aug 2025 19:42:27 +0530 Subject: [PATCH 06/23] add py3.9 support --- packages/toolbox-core/tests/test_toolbox_transport.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/toolbox-core/tests/test_toolbox_transport.py b/packages/toolbox-core/tests/test_toolbox_transport.py index ca2d7d79d..10bd3071c 100644 --- a/packages/toolbox-core/tests/test_toolbox_transport.py +++ b/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 +from typing import AsyncGenerator, Union from unittest.mock import AsyncMock import pytest @@ -106,7 +106,7 @@ async def test_tool_get_failure(http_session: ClientSession): async def test_tools_list_success( http_session: ClientSession, mock_manifest_dict: dict, - toolset_name: str | None, + toolset_name: Union[str, None], expected_path: str, ): """Tests successful tools_list calls with and without a toolset name.""" From 31d0f9e69a9236b2154920bbda2fd36b6d579c5b Mon Sep 17 00:00:00 2001 From: Twisha Bansal Date: Mon, 25 Aug 2025 20:01:21 +0530 Subject: [PATCH 07/23] fix langchain tool tests --- .../tests/test_async_tools.py | 71 ++++++++----------- 1 file changed, 28 insertions(+), 43 deletions(-) diff --git a/packages/toolbox-langchain/tests/test_async_tools.py b/packages/toolbox-langchain/tests/test_async_tools.py index 96bd7660b..33661e1b4 100644 --- a/packages/toolbox-langchain/tests/test_async_tools.py +++ b/packages/toolbox-langchain/tests/test_async_tools.py @@ -12,14 +12,15 @@ # See the License for the specific language governing permissions and # limitations under the License. -import types # For MappingProxyType -from unittest.mock import AsyncMock, Mock, patch +import types +from unittest.mock import AsyncMock, patch import pytest import pytest_asyncio from pydantic import ValidationError from toolbox_core.protocol import ParameterSchema as CoreParameterSchema from toolbox_core.tool import ToolboxTool as ToolboxCoreTool +from toolbox_core.toolbox_transport import ToolboxTransport from toolbox_langchain.async_tools import AsyncToolboxTool @@ -67,9 +68,11 @@ def _create_core_tool_from_dict( else: tool_constructor_params.append(p_schema) + transport = ToolboxTransport( + base_url=url, session=session, manage_session=False + ) return ToolboxCoreTool( - session=session, - base_url=url, + transport=transport, name=name, description=schema_dict["description"], params=tool_constructor_params, @@ -86,10 +89,10 @@ def _create_core_tool_from_dict( @patch("aiohttp.ClientSession") async def toolbox_tool(self, MockClientSession, tool_schema_dict): mock_session = MockClientSession.return_value - mock_response = mock_session.post.return_value.__aenter__.return_value - mock_response.raise_for_status = Mock() - mock_response.json = AsyncMock(return_value={"result": "test-result"}) - mock_response.status = 200 # *** Fix: Set status for the mock response *** + mock_session.post.return_value.__aenter__.return_value.json = AsyncMock( + return_value={"result": "test-result"} + ) + mock_session.post.return_value.__aenter__.return_value.ok = True core_tool_instance = self._create_core_tool_from_dict( session=mock_session, @@ -104,10 +107,10 @@ async def toolbox_tool(self, MockClientSession, tool_schema_dict): @patch("aiohttp.ClientSession") async def auth_toolbox_tool(self, MockClientSession, auth_tool_schema_dict): mock_session = MockClientSession.return_value - mock_response = mock_session.post.return_value.__aenter__.return_value - mock_response.raise_for_status = Mock() - mock_response.json = AsyncMock(return_value={"result": "test-result"}) - mock_response.status = 200 # *** Fix: Set status for the mock response *** + mock_session.post.return_value.__aenter__.return_value.json = AsyncMock( + return_value={"result": "test-result"} + ) + mock_session.post.return_value.__aenter__.return_value.ok = True core_tool_instance = self._create_core_tool_from_dict( session=mock_session, @@ -121,8 +124,6 @@ async def auth_toolbox_tool(self, MockClientSession, auth_tool_schema_dict): @patch("aiohttp.ClientSession") async def test_toolbox_tool_init(self, MockClientSession, tool_schema_dict): mock_session = MockClientSession.return_value - mock_response = mock_session.post.return_value.__aenter__.return_value - mock_response.status = 200 core_tool_instance = self._create_core_tool_from_dict( session=mock_session, name="test_tool", @@ -173,8 +174,6 @@ async def test_toolbox_tool_bind_params_duplicate(self, toolbox_tool): async def test_toolbox_tool_bind_params_invalid_params(self, auth_toolbox_tool): auth_core_tool = auth_toolbox_tool._AsyncToolboxTool__core_tool - # Verify that 'param1' is not in the list of bindable parameters for the core tool - # because it requires authentication. assert "param1" not in [p.name for p in auth_core_tool._ToolboxTool__params] with pytest.raises( ValueError, match="unable to bind parameters: no parameter named param1" @@ -227,19 +226,6 @@ async def test_toolbox_tool_add_unused_auth_token_getter_raises_error( in str(excinfo.value) ) - valid_lambda = lambda: "test-token" - with pytest.raises(ValueError) as excinfo_mixed: - auth_toolbox_tool.add_auth_token_getters( - { - "test-auth-source": valid_lambda, - "another-auth-source": unused_lambda, - } - ) - assert ( - "Authentication source(s) `another-auth-source` unused by tool `test_tool`" - in str(excinfo_mixed.value) - ) - async def test_toolbox_tool_add_auth_token_getters_duplicate( self, auth_toolbox_tool ): @@ -263,7 +249,8 @@ async def test_toolbox_tool_call(self, toolbox_tool): result = await toolbox_tool.ainvoke({"param1": "test-value", "param2": 123}) assert result == "test-result" core_tool = toolbox_tool._AsyncToolboxTool__core_tool - core_tool._ToolboxTool__session.post.assert_called_once_with( + session = core_tool._ToolboxTool__transport._ToolboxTransport__session + session.post.assert_called_once_with( "http://test_url/api/tool/test_tool/invoke", json={"param1": "test-value", "param2": 123}, headers={}, @@ -283,7 +270,8 @@ async def test_toolbox_tool_call_with_bound_params( result = await tool.ainvoke({"param2": 123}) assert result == "test-result" core_tool = tool._AsyncToolboxTool__core_tool - core_tool._ToolboxTool__session.post.assert_called_once_with( + session = core_tool._ToolboxTool__transport._ToolboxTransport__session + session.post.assert_called_once_with( "http://test_url/api/tool/test_tool/invoke", json={"param1": expected_value, "param2": 123}, headers={}, @@ -296,7 +284,8 @@ async def test_toolbox_tool_call_with_auth_tokens(self, auth_toolbox_tool): result = await tool.ainvoke({"param2": 123}) assert result == "test-result" core_tool = tool._AsyncToolboxTool__core_tool - core_tool._ToolboxTool__session.post.assert_called_once_with( + session = core_tool._ToolboxTool__transport._ToolboxTransport__session + session.post.assert_called_once_with( "https://test-url/api/tool/test_tool/invoke", json={"param2": 123}, headers={"test-auth-source_token": "test-token"}, @@ -306,7 +295,9 @@ async def test_toolbox_tool_call_with_auth_tokens_insecure( self, auth_toolbox_tool, auth_tool_schema_dict ): core_tool_of_auth_tool = auth_toolbox_tool._AsyncToolboxTool__core_tool - mock_session = core_tool_of_auth_tool._ToolboxTool__session + mock_session = ( + core_tool_of_auth_tool._ToolboxTool__transport._ToolboxTransport__session + ) with pytest.warns( UserWarning, @@ -327,16 +318,10 @@ async def test_toolbox_tool_call_with_auth_tokens_insecure( result = await tool_with_getter.ainvoke({"param2": 123}) assert result == "test-result" - modified_core_tool_in_new_tool = tool_with_getter._AsyncToolboxTool__core_tool - assert ( - modified_core_tool_in_new_tool._ToolboxTool__base_url == "http://test-url" - ) - assert ( - modified_core_tool_in_new_tool._ToolboxTool__url - == "http://test-url/api/tool/test_tool/invoke" - ) - - modified_core_tool_in_new_tool._ToolboxTool__session.post.assert_called_once_with( + modified_core_tool = tool_with_getter._AsyncToolboxTool__core_tool + session = modified_core_tool._ToolboxTool__transport._ToolboxTransport__session + assert modified_core_tool._ToolboxTool__transport.base_url == "http://test-url" + session.post.assert_called_once_with( "http://test-url/api/tool/test_tool/invoke", json={"param2": 123}, headers={"test-auth-source_token": "test-token"}, From b50513b6b51d985cc77d6b73dea1f5b66129e428 Mon Sep 17 00:00:00 2001 From: Twisha Bansal Date: Mon, 25 Aug 2025 20:17:03 +0530 Subject: [PATCH 08/23] test fix --- .../tests/test_async_tools.py | 85 ++++++++----------- 1 file changed, 34 insertions(+), 51 deletions(-) diff --git a/packages/toolbox-llamaindex/tests/test_async_tools.py b/packages/toolbox-llamaindex/tests/test_async_tools.py index 251c88cdd..c4b5729b3 100644 --- a/packages/toolbox-llamaindex/tests/test_async_tools.py +++ b/packages/toolbox-llamaindex/tests/test_async_tools.py @@ -12,8 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -import types # For MappingProxyType -from unittest.mock import AsyncMock, Mock, patch +import types +from unittest.mock import AsyncMock, patch import pytest import pytest_asyncio @@ -21,6 +21,7 @@ from pydantic import ValidationError from toolbox_core.protocol import ParameterSchema as CoreParameterSchema from toolbox_core.tool import ToolboxTool as ToolboxCoreTool +from toolbox_core.toolbox_transport import ToolboxTransport from toolbox_llamaindex.async_tools import AsyncToolboxTool @@ -68,9 +69,11 @@ def _create_core_tool_from_dict( else: tool_constructor_params.append(p_schema) + transport = ToolboxTransport( + base_url=url, session=session, manage_session=False + ) return ToolboxCoreTool( - session=session, - base_url=url, + transport=transport, name=name, description=schema_dict["description"], params=tool_constructor_params, @@ -87,10 +90,10 @@ def _create_core_tool_from_dict( @patch("aiohttp.ClientSession") async def toolbox_tool(self, MockClientSession, tool_schema_dict): mock_session = MockClientSession.return_value - mock_response = mock_session.post.return_value.__aenter__.return_value - mock_response.raise_for_status = Mock() - mock_response.json = AsyncMock(return_value={"result": "test-result"}) - mock_response.status = 200 # *** Fix: Set status for the mock response *** + mock_session.post.return_value.__aenter__.return_value.json = AsyncMock( + return_value={"result": "test-result"} + ) + mock_session.post.return_value.__aenter__.return_value.ok = True core_tool_instance = self._create_core_tool_from_dict( session=mock_session, @@ -105,11 +108,10 @@ async def toolbox_tool(self, MockClientSession, tool_schema_dict): @patch("aiohttp.ClientSession") async def auth_toolbox_tool(self, MockClientSession, auth_tool_schema_dict): mock_session = MockClientSession.return_value - mock_response = mock_session.post.return_value.__aenter__.return_value - mock_response.raise_for_status = Mock() - mock_response.json = AsyncMock(return_value={"result": "test-result"}) - mock_response.status = 200 # *** Fix: Set status for the mock response *** - + mock_session.post.return_value.__aenter__.return_value.json = AsyncMock( + return_value={"result": "test-result"} + ) + mock_session.post.return_value.__aenter__.return_value.ok = True core_tool_instance = self._create_core_tool_from_dict( session=mock_session, name="test_tool", @@ -122,8 +124,6 @@ async def auth_toolbox_tool(self, MockClientSession, auth_tool_schema_dict): @patch("aiohttp.ClientSession") async def test_toolbox_tool_init(self, MockClientSession, tool_schema_dict): mock_session = MockClientSession.return_value - mock_response = mock_session.post.return_value.__aenter__.return_value - mock_response.status = 200 core_tool_instance = self._create_core_tool_from_dict( session=mock_session, name="test_tool", @@ -174,8 +174,6 @@ async def test_toolbox_tool_bind_params_duplicate(self, toolbox_tool): async def test_toolbox_tool_bind_params_invalid_params(self, auth_toolbox_tool): auth_core_tool = auth_toolbox_tool._AsyncToolboxTool__core_tool - # Verify that 'param1' is not in the list of bindable parameters for the core tool - # because it requires authentication. assert "param1" not in [p.name for p in auth_core_tool._ToolboxTool__params] with pytest.raises( ValueError, match="unable to bind parameters: no parameter named param1" @@ -228,19 +226,6 @@ async def test_toolbox_tool_add_unused_auth_token_getter_raises_error( in str(excinfo.value) ) - valid_lambda = lambda: "test-token" - with pytest.raises(ValueError) as excinfo_mixed: - auth_toolbox_tool.add_auth_token_getters( - { - "test-auth-source": valid_lambda, - "another-auth-source": unused_lambda, - } - ) - assert ( - "Authentication source(s) `another-auth-source` unused by tool `test_tool`" - in str(excinfo_mixed.value) - ) - async def test_toolbox_tool_add_auth_token_getters_duplicate( self, auth_toolbox_tool ): @@ -267,11 +252,11 @@ async def test_toolbox_tool_call(self, toolbox_tool): tool_name="test_tool", raw_input={"param1": "test-value", "param2": 123}, raw_output="test-result", - is_error=False, ) core_tool = toolbox_tool._AsyncToolboxTool__core_tool - core_tool._ToolboxTool__session.post.assert_called_once_with( + session = core_tool._ToolboxTool__transport._ToolboxHttpTransport__session + session.post.assert_called_once_with( "http://test_url/api/tool/test_tool/invoke", json={"param1": "test-value", "param2": 123}, headers={}, @@ -294,10 +279,10 @@ async def test_toolbox_tool_call_with_bound_params( tool_name="test_tool", raw_input={"param2": 123}, raw_output="test-result", - is_error=False, ) core_tool = tool._AsyncToolboxTool__core_tool - core_tool._ToolboxTool__session.post.assert_called_once_with( + session = core_tool._ToolboxTool__transport._ToolboxHttpTransport__session + session.post.assert_called_once_with( "http://test_url/api/tool/test_tool/invoke", json={"param1": expected_value, "param2": 123}, headers={}, @@ -313,11 +298,11 @@ async def test_toolbox_tool_call_with_auth_tokens(self, auth_toolbox_tool): tool_name="test_tool", raw_input={"param2": 123}, raw_output="test-result", - is_error=False, ) core_tool = tool._AsyncToolboxTool__core_tool - core_tool._ToolboxTool__session.post.assert_called_once_with( + session = core_tool._ToolboxTool__transport._ToolboxHttpTransport__session + session.post.assert_called_once_with( "https://test-url/api/tool/test_tool/invoke", json={"param2": 123}, headers={"test-auth-source_token": "test-token"}, @@ -327,7 +312,9 @@ async def test_toolbox_tool_call_with_auth_tokens_insecure( self, auth_toolbox_tool, auth_tool_schema_dict ): core_tool_of_auth_tool = auth_toolbox_tool._AsyncToolboxTool__core_tool - mock_session = core_tool_of_auth_tool._ToolboxTool__session + mock_session = ( + core_tool_of_auth_tool._ToolboxTool__transport._ToolboxTransport__session + ) with pytest.warns( UserWarning, @@ -340,9 +327,9 @@ async def test_toolbox_tool_call_with_auth_tokens_insecure( url="http://test-url", ) - insecure_auth_langchain_tool = AsyncToolboxTool(core_tool=insecure_core_tool) + insecure_auth_llamaindex_tool = AsyncToolboxTool(core_tool=insecure_core_tool) - tool_with_getter = insecure_auth_langchain_tool.add_auth_token_getters( + tool_with_getter = insecure_auth_llamaindex_tool.add_auth_token_getters( {"test-auth-source": lambda: "test-token"} ) result = await tool_with_getter.acall(param2=123) @@ -351,19 +338,15 @@ async def test_toolbox_tool_call_with_auth_tokens_insecure( tool_name="test_tool", raw_input={"param2": 123}, raw_output="test-result", - is_error=False, - ) - - modified_core_tool_in_new_tool = tool_with_getter._AsyncToolboxTool__core_tool - assert ( - modified_core_tool_in_new_tool._ToolboxTool__base_url == "http://test-url" - ) - assert ( - modified_core_tool_in_new_tool._ToolboxTool__url - == "http://test-url/api/tool/test_tool/invoke" ) - modified_core_tool_in_new_tool._ToolboxTool__session.post.assert_called_once_with( + modified_core_tool = tool_with_getter._AsyncToolboxTool__core_tool + transport = modified_core_tool._ToolboxTool__transport + session = transport._ToolboxHttpTransport__session + + assert transport.base_url == "http://test-url" + + session.post.assert_called_once_with( "http://test-url/api/tool/test_tool/invoke", json={"param2": 123}, headers={"test-auth-source_token": "test-token"}, @@ -383,4 +366,4 @@ async def test_toolbox_tool_call_with_invalid_input(self, toolbox_tool): async def test_toolbox_tool_run_not_implemented(self, toolbox_tool): with pytest.raises(NotImplementedError): - toolbox_tool.call() + toolbox_tool.call() \ No newline at end of file From c222d25ba0c19d9a301de66d75394d55d39a5e04 Mon Sep 17 00:00:00 2001 From: Twisha Bansal Date: Mon, 25 Aug 2025 20:18:56 +0530 Subject: [PATCH 09/23] lint --- packages/toolbox-llamaindex/tests/test_async_tools.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/packages/toolbox-llamaindex/tests/test_async_tools.py b/packages/toolbox-llamaindex/tests/test_async_tools.py index c4b5729b3..fb10e61c4 100644 --- a/packages/toolbox-llamaindex/tests/test_async_tools.py +++ b/packages/toolbox-llamaindex/tests/test_async_tools.py @@ -343,9 +343,9 @@ async def test_toolbox_tool_call_with_auth_tokens_insecure( modified_core_tool = tool_with_getter._AsyncToolboxTool__core_tool transport = modified_core_tool._ToolboxTool__transport session = transport._ToolboxHttpTransport__session - + assert transport.base_url == "http://test-url" - + session.post.assert_called_once_with( "http://test-url/api/tool/test_tool/invoke", json={"param2": 123}, @@ -366,4 +366,4 @@ async def test_toolbox_tool_call_with_invalid_input(self, toolbox_tool): async def test_toolbox_tool_run_not_implemented(self, toolbox_tool): with pytest.raises(NotImplementedError): - toolbox_tool.call() \ No newline at end of file + toolbox_tool.call() From b2447409b4e2337749126df98f8d8a013e40ec9d Mon Sep 17 00:00:00 2001 From: Twisha Bansal Date: Mon, 25 Aug 2025 20:23:54 +0530 Subject: [PATCH 10/23] fix tests --- packages/toolbox-llamaindex/tests/test_async_tools.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/packages/toolbox-llamaindex/tests/test_async_tools.py b/packages/toolbox-llamaindex/tests/test_async_tools.py index fb10e61c4..a817f542f 100644 --- a/packages/toolbox-llamaindex/tests/test_async_tools.py +++ b/packages/toolbox-llamaindex/tests/test_async_tools.py @@ -255,7 +255,7 @@ async def test_toolbox_tool_call(self, toolbox_tool): ) core_tool = toolbox_tool._AsyncToolboxTool__core_tool - session = core_tool._ToolboxTool__transport._ToolboxHttpTransport__session + session = core_tool._ToolboxTool__transport._ToolboxTransport__session session.post.assert_called_once_with( "http://test_url/api/tool/test_tool/invoke", json={"param1": "test-value", "param2": 123}, @@ -281,7 +281,7 @@ async def test_toolbox_tool_call_with_bound_params( raw_output="test-result", ) core_tool = tool._AsyncToolboxTool__core_tool - session = core_tool._ToolboxTool__transport._ToolboxHttpTransport__session + session = core_tool._ToolboxTool__transport._ToolboxTransport__session session.post.assert_called_once_with( "http://test_url/api/tool/test_tool/invoke", json={"param1": expected_value, "param2": 123}, @@ -301,7 +301,7 @@ async def test_toolbox_tool_call_with_auth_tokens(self, auth_toolbox_tool): ) core_tool = tool._AsyncToolboxTool__core_tool - session = core_tool._ToolboxTool__transport._ToolboxHttpTransport__session + session = core_tool._ToolboxTool__transport._ToolboxTransport__session session.post.assert_called_once_with( "https://test-url/api/tool/test_tool/invoke", json={"param2": 123}, @@ -342,7 +342,7 @@ async def test_toolbox_tool_call_with_auth_tokens_insecure( modified_core_tool = tool_with_getter._AsyncToolboxTool__core_tool transport = modified_core_tool._ToolboxTool__transport - session = transport._ToolboxHttpTransport__session + session = transport._ToolboxTransport__session assert transport.base_url == "http://test-url" From d93f4dd2c69765f0b489b56281360607ab9ead8a Mon Sep 17 00:00:00 2001 From: Twisha Bansal Date: Thu, 28 Aug 2025 14:05:25 +0530 Subject: [PATCH 11/23] move manage session into transport --- .../toolbox-core/src/toolbox_core/client.py | 7 +-- .../src/toolbox_core/toolbox_transport.py | 12 +++-- packages/toolbox-core/tests/test_tool.py | 24 ++++----- .../tests/test_toolbox_transport.py | 52 ++++++++++--------- 4 files changed, 49 insertions(+), 46 deletions(-) diff --git a/packages/toolbox-core/src/toolbox_core/client.py b/packages/toolbox-core/src/toolbox_core/client.py index 5db0c957b..abef0b732 100644 --- a/packages/toolbox-core/src/toolbox_core/client.py +++ b/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( diff --git a/packages/toolbox-core/src/toolbox_core/toolbox_transport.py b/packages/toolbox-core/src/toolbox_core/toolbox_transport.py index 06b61bfc2..bf195bf78 100644 --- a/packages/toolbox-core/src/toolbox_core/toolbox_transport.py +++ b/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: diff --git a/packages/toolbox-core/tests/test_tool.py b/packages/toolbox-core/tests/test_tool.py index 449b94405..44051c504 100644 --- a/packages/toolbox-core/tests/test_tool.py +++ b/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, diff --git a/packages/toolbox-core/tests/test_toolbox_transport.py b/packages/toolbox-core/tests/test_toolbox_transport.py index 10bd3071c..7c8485398 100644 --- a/packages/toolbox-core/tests/test_toolbox_transport.py +++ b/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 From f9d8e64ea6200912afa8fcecb0e6538e96fd108e Mon Sep 17 00:00:00 2001 From: Twisha Bansal Date: Thu, 28 Aug 2025 14:47:18 +0530 Subject: [PATCH 12/23] move warning to diff file --- .../toolbox-core/src/toolbox_core/tool.py | 11 -- .../src/toolbox_core/toolbox_transport.py | 9 ++ packages/toolbox-core/tests/test_tool.py | 127 ------------------ .../tests/test_toolbox_transport.py | 47 ++++++- 4 files changed, 55 insertions(+), 139 deletions(-) diff --git a/packages/toolbox-core/src/toolbox_core/tool.py b/packages/toolbox-core/src/toolbox_core/tool.py index aa6ea260c..5678da794 100644 --- a/packages/toolbox-core/src/toolbox_core/tool.py +++ b/packages/toolbox-core/src/toolbox_core/tool.py @@ -128,17 +128,6 @@ def __init__( # map of client headers to their value/callable/coroutine self.__client_headers = client_headers - # ID tokens contain sensitive user information (claims). Transmitting - # these over HTTP exposes the data to interception and unauthorized - # access. Always use HTTPS to ensure secure communication and protect - # user privacy. - if self.__transport.base_url.startswith("http://") and ( - required_authn_params or required_authz_tokens or client_headers - ): - warn( - "Sending ID token over HTTP. User data may be exposed. Use HTTPS for secure communication." - ) - @property def _name(self) -> str: return self.__name__ diff --git a/packages/toolbox-core/src/toolbox_core/toolbox_transport.py b/packages/toolbox-core/src/toolbox_core/toolbox_transport.py index bf195bf78..4ff7520c4 100644 --- a/packages/toolbox-core/src/toolbox_core/toolbox_transport.py +++ b/packages/toolbox-core/src/toolbox_core/toolbox_transport.py @@ -13,6 +13,7 @@ # limitations under the License. from typing import Mapping, Optional +from warnings import warn from aiohttp import ClientSession @@ -70,6 +71,14 @@ async def tools_list( async def tool_invoke( self, tool_name: str, arguments: dict, headers: Mapping[str, str] ) -> dict: + # ID tokens contain sensitive user information (claims). Transmitting + # these over HTTP exposes the data to interception and unauthorized + # access. Always use HTTPS to ensure secure communication and protect + # user privacy. + if self.base_url.startswith("http://") and headers: + warn( + "Sending data token over HTTP. User data may be exposed. Use HTTPS for secure communication." + ) url = f"{self.__base_url}/api/tool/{tool_name}/invoke" async with self.__session.post( url, diff --git a/packages/toolbox-core/tests/test_tool.py b/packages/toolbox-core/tests/test_tool.py index 44051c504..214947b89 100644 --- a/packages/toolbox-core/tests/test_tool.py +++ b/packages/toolbox-core/tests/test_tool.py @@ -622,130 +622,3 @@ def test_toolbox_tool_underscore_client_headers_property(toolbox_tool: ToolboxTo # Verify immutability with pytest.raises(TypeError): client_headers["new_header"] = "new_value" - - -# --- Test for the HTTP Warning --- -@pytest.mark.parametrize( - "trigger_condition_params", - [ - {"client_headers": {"X-Some-Header": "value"}}, - {"required_authn_params": {"param1": ["auth-service1"]}}, - {"required_authz_tokens": ["auth-service2"]}, - { - "client_headers": {"X-Some-Header": "value"}, - "required_authn_params": {"param1": ["auth-service1"]}, - }, - { - "client_headers": {"X-Some-Header": "value"}, - "required_authz_tokens": ["auth-service2"], - }, - { - "required_authn_params": {"param1": ["auth-service1"]}, - "required_authz_tokens": ["auth-service2"], - }, - { - "client_headers": {"X-Some-Header": "value"}, - "required_authn_params": {"param1": ["auth-service1"]}, - "required_authz_tokens": ["auth-service2"], - }, - ], - ids=[ - "client_headers_only", - "authn_params_only", - "authz_tokens_only", - "headers_and_authn", - "headers_and_authz", - "authn_and_authz", - "all_three_conditions", - ], -) -def test_tool_init_http_warning_when_sensitive_info_over_http( - http_session: ClientSession, - sample_tool_params: list[ParameterSchema], - sample_tool_description: str, - trigger_condition_params: dict, -): - """ - Tests that a UserWarning is issued if client headers, auth params, or - auth tokens are present and the base_url is HTTP. - """ - expected_warning_message: str = ( - "Sending ID token over HTTP. User data may be exposed. " - "Use HTTPS for secure communication." - ) - transport = ToolboxTransport(TEST_BASE_URL, http_session) - init_kwargs = { - "transport": transport, - "name": "http_warning_tool", - "description": sample_tool_description, - "params": sample_tool_params, - "required_authn_params": {}, - "required_authz_tokens": [], - "auth_service_token_getters": {}, - "bound_params": {}, - "client_headers": {}, - } - # Apply the specific conditions for this parametrized test - init_kwargs.update(trigger_condition_params) - - with pytest.warns(UserWarning, match=expected_warning_message): - ToolboxTool(**init_kwargs) - - -def test_tool_init_no_http_warning_if_https( - http_session: ClientSession, - sample_tool_params: list[ParameterSchema], - sample_tool_description: str, - static_client_header: dict, -): - """ - Tests that NO UserWarning is issued if client headers are present but - the base_url is HTTPS. - """ - with catch_warnings(record=True) as record: - simplefilter("always") - transport = ToolboxTransport(HTTPS_BASE_URL, http_session) - - ToolboxTool( - transport=transport, - name="https_tool", - description=sample_tool_description, - params=sample_tool_params, - required_authn_params={}, - required_authz_tokens=[], - auth_service_token_getters={}, - bound_params={}, - client_headers=static_client_header, - ) - assert ( - len(record) == 0 - ), f"Expected no warnings, but got: {[f'{w.category.__name__}: {w.message}' for w in record]}" - - -def test_tool_init_no_http_warning_if_no_sensitive_info_on_http( - http_session: ClientSession, - sample_tool_params: list[ParameterSchema], - sample_tool_description: str, -): - """ - Tests that NO UserWarning is issued if the URL is HTTP but there are - no client headers, auth params, or auth tokens. - """ - with catch_warnings(record=True) as record: - simplefilter("always") - transport = ToolboxTransport(TEST_BASE_URL, http_session) - - ToolboxTool( - transport=transport, - name="http_tool_no_sensitive", - description=sample_tool_description, - params=sample_tool_params, - required_authn_params={}, - required_authz_tokens=[], - auth_service_token_getters={}, - bound_params={}, - client_headers={}, - ) - assert ( - len(record) == 0 - ), f"Expected no warnings, but got: {[f'{w.category.__name__}: {w.message}' for w in record]}" diff --git a/packages/toolbox-core/tests/test_toolbox_transport.py b/packages/toolbox-core/tests/test_toolbox_transport.py index 7c8485398..a40c79c2a 100644 --- a/packages/toolbox-core/tests/test_toolbox_transport.py +++ b/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, Optional, Union +from typing import AsyncGenerator, Mapping, Optional, Union from unittest.mock import AsyncMock import pytest @@ -154,6 +154,48 @@ async def test_tool_invoke_failure(http_session: ClientSession): assert str(exc_info.value) == "Invalid arguments" +@pytest.mark.asyncio +@pytest.mark.parametrize( + "base_url, headers, should_warn", + [ + ( + "http://fake-toolbox-server.com", + {"Authorization": "Bearer token"}, + True, + ), + ( + "https://fake-toolbox-server.com", + {"Authorization": "Bearer token"}, + False, + ), + ("http://fake-toolbox-server.com", {}, False), + ("http://fake-toolbox-server.com", None, False), + ], +) +async def test_tool_invoke_http_warning( + http_session: ClientSession, + base_url: str, + headers: Optional[Mapping[str, str]], + should_warn: bool, +): + """Tests the HTTP security warning logic in tool_invoke.""" + url = f"{base_url}/api/tool/{TEST_TOOL_NAME}/invoke" + args = {"param1": "value1"} + response_payload = {"result": "success"} + transport = ToolboxTransport(base_url, http_session) + + with aioresponses() as m: + m.post(url, status=200, payload=response_payload) + + if should_warn: + with pytest.warns(UserWarning, match="Sending data token over HTTP"): + await transport.tool_invoke(TEST_TOOL_NAME, args, headers) + else: + # By not using pytest.warns, we assert that no warnings are raised. + # The test will fail if an unexpected UserWarning occurs. + await transport.tool_invoke(TEST_TOOL_NAME, args, headers) + + @pytest.mark.asyncio async def test_close_does_not_close_unmanaged_session(): """ @@ -175,6 +217,9 @@ async def test_close_closes_managed_session(): managed internally by the transport. """ transport = ToolboxTransport(TEST_BASE_URL, session=None) + # Access the internal session before closing to check its state + internal_session = transport._ToolboxTransport__session + assert internal_session.closed is False await transport.close() internal_session = transport._ToolboxTransport__session From d6361ed1e421a5f3031892f6d5cca5e3f9a01754 Mon Sep 17 00:00:00 2001 From: Twisha Bansal Date: Thu, 28 Aug 2025 14:49:47 +0530 Subject: [PATCH 13/23] avoid code duplication --- .../src/toolbox_core/toolbox_transport.py | 21 +++++++++---------- 1 file changed, 10 insertions(+), 11 deletions(-) diff --git a/packages/toolbox-core/src/toolbox_core/toolbox_transport.py b/packages/toolbox-core/src/toolbox_core/toolbox_transport.py index 4ff7520c4..7fa6a006c 100644 --- a/packages/toolbox-core/src/toolbox_core/toolbox_transport.py +++ b/packages/toolbox-core/src/toolbox_core/toolbox_transport.py @@ -40,10 +40,10 @@ def base_url(self) -> str: """The base URL for the transport.""" return self.__base_url - async def tool_get( - self, tool_name: str, headers: Optional[Mapping[str, str]] = None + async def _get_manifest( + self, url: str, headers: Optional[Mapping[str, str]] ) -> ManifestSchema: - url = f"{self.__base_url}/api/tool/{tool_name}" + """Helper method to perform GET requests and parse the ManifestSchema.""" async with self.__session.get(url, headers=headers) as response: if not response.ok: error_text = await response.text() @@ -53,20 +53,19 @@ async def tool_get( json = await response.json() return ManifestSchema(**json) + async def tool_get( + self, tool_name: str, headers: Optional[Mapping[str, str]] = None + ) -> ManifestSchema: + url = f"{self.__base_url}/api/tool/{tool_name}" + return await self._get_manifest(url, headers) + async def tools_list( self, toolset_name: Optional[str] = None, headers: Optional[Mapping[str, str]] = None, ) -> ManifestSchema: url = f"{self.__base_url}/api/toolset/{toolset_name or ''}" - async with self.__session.get(url, headers=headers) as response: - if not response.ok: - error_text = await response.text() - raise RuntimeError( - f"API request failed with status {response.status} ({response.reason}). Server response: {error_text}" - ) - json = await response.json() - return ManifestSchema(**json) + return await self._get_manifest(url, headers) async def tool_invoke( self, tool_name: str, arguments: dict, headers: Mapping[str, str] From 0f3bacc878735c722252a5e426fc9ef4a9144d50 Mon Sep 17 00:00:00 2001 From: Twisha Bansal Date: Thu, 28 Aug 2025 14:54:23 +0530 Subject: [PATCH 14/23] fix tests --- packages/toolbox-langchain/tests/test_async_tools.py | 2 +- packages/toolbox-llamaindex/tests/test_async_tools.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/toolbox-langchain/tests/test_async_tools.py b/packages/toolbox-langchain/tests/test_async_tools.py index 33661e1b4..6aac2af15 100644 --- a/packages/toolbox-langchain/tests/test_async_tools.py +++ b/packages/toolbox-langchain/tests/test_async_tools.py @@ -69,7 +69,7 @@ def _create_core_tool_from_dict( tool_constructor_params.append(p_schema) transport = ToolboxTransport( - base_url=url, session=session, manage_session=False + base_url=url, session=session ) return ToolboxCoreTool( transport=transport, diff --git a/packages/toolbox-llamaindex/tests/test_async_tools.py b/packages/toolbox-llamaindex/tests/test_async_tools.py index a817f542f..49c37e27c 100644 --- a/packages/toolbox-llamaindex/tests/test_async_tools.py +++ b/packages/toolbox-llamaindex/tests/test_async_tools.py @@ -70,7 +70,7 @@ def _create_core_tool_from_dict( tool_constructor_params.append(p_schema) transport = ToolboxTransport( - base_url=url, session=session, manage_session=False + base_url=url, session=session ) return ToolboxCoreTool( transport=transport, From 0a79f5af28fc000feec5fb34570938278d09e919 Mon Sep 17 00:00:00 2001 From: Twisha Bansal Date: Thu, 28 Aug 2025 14:54:55 +0530 Subject: [PATCH 15/23] lint --- packages/toolbox-langchain/tests/test_async_tools.py | 4 +--- packages/toolbox-llamaindex/tests/test_async_tools.py | 4 +--- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/packages/toolbox-langchain/tests/test_async_tools.py b/packages/toolbox-langchain/tests/test_async_tools.py index 6aac2af15..26080ea1b 100644 --- a/packages/toolbox-langchain/tests/test_async_tools.py +++ b/packages/toolbox-langchain/tests/test_async_tools.py @@ -68,9 +68,7 @@ def _create_core_tool_from_dict( else: tool_constructor_params.append(p_schema) - transport = ToolboxTransport( - base_url=url, session=session - ) + transport = ToolboxTransport(base_url=url, session=session) return ToolboxCoreTool( transport=transport, name=name, diff --git a/packages/toolbox-llamaindex/tests/test_async_tools.py b/packages/toolbox-llamaindex/tests/test_async_tools.py index 49c37e27c..6237a0c40 100644 --- a/packages/toolbox-llamaindex/tests/test_async_tools.py +++ b/packages/toolbox-llamaindex/tests/test_async_tools.py @@ -69,9 +69,7 @@ def _create_core_tool_from_dict( else: tool_constructor_params.append(p_schema) - transport = ToolboxTransport( - base_url=url, session=session - ) + transport = ToolboxTransport(base_url=url, session=session) return ToolboxCoreTool( transport=transport, name=name, From fcf7da3e3890dae46b72393cd462cffa41559c1d Mon Sep 17 00:00:00 2001 From: Twisha Bansal Date: Thu, 28 Aug 2025 14:58:19 +0530 Subject: [PATCH 16/23] remove redundant tests --- .../tests/test_async_tools.py | 36 --------------- .../tests/test_async_tools.py | 44 ------------------- 2 files changed, 80 deletions(-) diff --git a/packages/toolbox-langchain/tests/test_async_tools.py b/packages/toolbox-langchain/tests/test_async_tools.py index 26080ea1b..84aeab52e 100644 --- a/packages/toolbox-langchain/tests/test_async_tools.py +++ b/packages/toolbox-langchain/tests/test_async_tools.py @@ -289,42 +289,6 @@ async def test_toolbox_tool_call_with_auth_tokens(self, auth_toolbox_tool): headers={"test-auth-source_token": "test-token"}, ) - async def test_toolbox_tool_call_with_auth_tokens_insecure( - self, auth_toolbox_tool, auth_tool_schema_dict - ): - core_tool_of_auth_tool = auth_toolbox_tool._AsyncToolboxTool__core_tool - mock_session = ( - core_tool_of_auth_tool._ToolboxTool__transport._ToolboxTransport__session - ) - - with pytest.warns( - UserWarning, - match="Sending ID token over HTTP. User data may be exposed. Use HTTPS for secure communication.", - ): - insecure_core_tool = self._create_core_tool_from_dict( - session=mock_session, - name="test_tool", - schema_dict=auth_tool_schema_dict, - url="http://test-url", - ) - - insecure_auth_langchain_tool = AsyncToolboxTool(core_tool=insecure_core_tool) - - tool_with_getter = insecure_auth_langchain_tool.add_auth_token_getters( - {"test-auth-source": lambda: "test-token"} - ) - result = await tool_with_getter.ainvoke({"param2": 123}) - assert result == "test-result" - - modified_core_tool = tool_with_getter._AsyncToolboxTool__core_tool - session = modified_core_tool._ToolboxTool__transport._ToolboxTransport__session - assert modified_core_tool._ToolboxTool__transport.base_url == "http://test-url" - session.post.assert_called_once_with( - "http://test-url/api/tool/test_tool/invoke", - json={"param2": 123}, - headers={"test-auth-source_token": "test-token"}, - ) - async def test_toolbox_tool_call_with_invalid_input(self, toolbox_tool): with pytest.raises(ValidationError) as e: await toolbox_tool.ainvoke({"param1": 123, "param2": "invalid"}) diff --git a/packages/toolbox-llamaindex/tests/test_async_tools.py b/packages/toolbox-llamaindex/tests/test_async_tools.py index 6237a0c40..2230d51ec 100644 --- a/packages/toolbox-llamaindex/tests/test_async_tools.py +++ b/packages/toolbox-llamaindex/tests/test_async_tools.py @@ -306,50 +306,6 @@ async def test_toolbox_tool_call_with_auth_tokens(self, auth_toolbox_tool): headers={"test-auth-source_token": "test-token"}, ) - async def test_toolbox_tool_call_with_auth_tokens_insecure( - self, auth_toolbox_tool, auth_tool_schema_dict - ): - core_tool_of_auth_tool = auth_toolbox_tool._AsyncToolboxTool__core_tool - mock_session = ( - core_tool_of_auth_tool._ToolboxTool__transport._ToolboxTransport__session - ) - - with pytest.warns( - UserWarning, - match="Sending ID token over HTTP. User data may be exposed. Use HTTPS for secure communication.", - ): - insecure_core_tool = self._create_core_tool_from_dict( - session=mock_session, - name="test_tool", - schema_dict=auth_tool_schema_dict, - url="http://test-url", - ) - - insecure_auth_llamaindex_tool = AsyncToolboxTool(core_tool=insecure_core_tool) - - tool_with_getter = insecure_auth_llamaindex_tool.add_auth_token_getters( - {"test-auth-source": lambda: "test-token"} - ) - result = await tool_with_getter.acall(param2=123) - assert result == ToolOutput( - content="test-result", - tool_name="test_tool", - raw_input={"param2": 123}, - raw_output="test-result", - ) - - modified_core_tool = tool_with_getter._AsyncToolboxTool__core_tool - transport = modified_core_tool._ToolboxTool__transport - session = transport._ToolboxTransport__session - - assert transport.base_url == "http://test-url" - - session.post.assert_called_once_with( - "http://test-url/api/tool/test_tool/invoke", - json={"param2": 123}, - headers={"test-auth-source_token": "test-token"}, - ) - async def test_toolbox_tool_call_with_empty_input(self, toolbox_tool): with pytest.raises(TypeError) as e: await toolbox_tool.acall() From c65a94adc3af3e61e44d436c0b6e149b253a0269 Mon Sep 17 00:00:00 2001 From: Twisha Bansal Date: Mon, 1 Sep 2025 16:00:22 +0530 Subject: [PATCH 17/23] make invoke method return str --- packages/toolbox-core/src/toolbox_core/itransport.py | 2 +- packages/toolbox-core/src/toolbox_core/tool.py | 2 +- packages/toolbox-core/src/toolbox_core/toolbox_transport.py | 2 +- packages/toolbox-core/tests/test_toolbox_transport.py | 4 ++-- 4 files changed, 5 insertions(+), 5 deletions(-) diff --git a/packages/toolbox-core/src/toolbox_core/itransport.py b/packages/toolbox-core/src/toolbox_core/itransport.py index 7d4a87e5a..0b38beb57 100644 --- a/packages/toolbox-core/src/toolbox_core/itransport.py +++ b/packages/toolbox-core/src/toolbox_core/itransport.py @@ -48,7 +48,7 @@ async def tools_list( @abstractmethod async def tool_invoke( self, tool_name: str, arguments: dict, headers: Mapping[str, str] - ) -> dict: + ) -> str: """Invokes a specific tool on the server.""" pass diff --git a/packages/toolbox-core/src/toolbox_core/tool.py b/packages/toolbox-core/src/toolbox_core/tool.py index 5678da794..90e4616f6 100644 --- a/packages/toolbox-core/src/toolbox_core/tool.py +++ b/packages/toolbox-core/src/toolbox_core/tool.py @@ -290,7 +290,7 @@ async def __call__(self, *args: Any, **kwargs: Any) -> str: payload, headers, ) - return body.get("result", body) + return body def add_auth_token_getters( self, diff --git a/packages/toolbox-core/src/toolbox_core/toolbox_transport.py b/packages/toolbox-core/src/toolbox_core/toolbox_transport.py index 7fa6a006c..82d0c3153 100644 --- a/packages/toolbox-core/src/toolbox_core/toolbox_transport.py +++ b/packages/toolbox-core/src/toolbox_core/toolbox_transport.py @@ -88,7 +88,7 @@ async def tool_invoke( if not resp.ok: err = body.get("error", f"unexpected status from server: {resp.status}") raise Exception(err) - return body + return body.get("result") async def close(self): if self.__manage_session and not self.__session.closed: diff --git a/packages/toolbox-core/tests/test_toolbox_transport.py b/packages/toolbox-core/tests/test_toolbox_transport.py index a40c79c2a..23a373d95 100644 --- a/packages/toolbox-core/tests/test_toolbox_transport.py +++ b/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, Mapping, Optional, Union +from typing import AsyncGenerator, Mapping, Union, Optional from unittest.mock import AsyncMock import pytest @@ -135,7 +135,7 @@ async def test_tool_invoke_success(http_session: ClientSession): m.post(url, status=200, payload=response_payload) result = await transport.tool_invoke(TEST_TOOL_NAME, args, headers) - assert result == response_payload + assert result == "success" m.assert_called_once_with(url, method="POST", json=args, headers=headers) From e2269f89c878152aad9be096d6e3a3989261cad3 Mon Sep 17 00:00:00 2001 From: Twisha Bansal Date: Mon, 1 Sep 2025 16:00:40 +0530 Subject: [PATCH 18/23] lint --- packages/toolbox-core/tests/test_toolbox_transport.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/toolbox-core/tests/test_toolbox_transport.py b/packages/toolbox-core/tests/test_toolbox_transport.py index 23a373d95..2921e09f8 100644 --- a/packages/toolbox-core/tests/test_toolbox_transport.py +++ b/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, Mapping, Union, Optional +from typing import AsyncGenerator, Mapping, Optional, Union from unittest.mock import AsyncMock import pytest From 25802e2fa70734c3051a123f6eb3ba0219ca6ce0 Mon Sep 17 00:00:00 2001 From: Twisha Bansal Date: Mon, 1 Sep 2025 16:02:48 +0530 Subject: [PATCH 19/23] fix return type --- packages/toolbox-core/src/toolbox_core/toolbox_transport.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/toolbox-core/src/toolbox_core/toolbox_transport.py b/packages/toolbox-core/src/toolbox_core/toolbox_transport.py index 82d0c3153..593de0b2f 100644 --- a/packages/toolbox-core/src/toolbox_core/toolbox_transport.py +++ b/packages/toolbox-core/src/toolbox_core/toolbox_transport.py @@ -69,7 +69,7 @@ async def tools_list( async def tool_invoke( self, tool_name: str, arguments: dict, headers: Mapping[str, str] - ) -> dict: + ) -> str: # ID tokens contain sensitive user information (claims). Transmitting # these over HTTP exposes the data to interception and unauthorized # access. Always use HTTPS to ensure secure communication and protect From d41aed8d03f9b3d12737f39b523c1bfa882d409e Mon Sep 17 00:00:00 2001 From: Twisha Bansal Date: Mon, 1 Sep 2025 16:04:32 +0530 Subject: [PATCH 20/23] small refactor --- packages/toolbox-core/src/toolbox_core/tool.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/packages/toolbox-core/src/toolbox_core/tool.py b/packages/toolbox-core/src/toolbox_core/tool.py index 90e4616f6..70c48a44c 100644 --- a/packages/toolbox-core/src/toolbox_core/tool.py +++ b/packages/toolbox-core/src/toolbox_core/tool.py @@ -285,12 +285,11 @@ async def __call__(self, *args: Any, **kwargs: Any) -> str: token_getter ) - body = await self.__transport.tool_invoke( + return await self.__transport.tool_invoke( self.__name__, payload, headers, ) - return body def add_auth_token_getters( self, From 31809a1e41932e01b60bd919a4b18483d8aa53e5 Mon Sep 17 00:00:00 2001 From: Twisha Bansal Date: Thu, 4 Sep 2025 14:25:03 +0530 Subject: [PATCH 21/23] rename private method --- packages/toolbox-core/src/toolbox_core/toolbox_transport.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/packages/toolbox-core/src/toolbox_core/toolbox_transport.py b/packages/toolbox-core/src/toolbox_core/toolbox_transport.py index 593de0b2f..0f1e7e40d 100644 --- a/packages/toolbox-core/src/toolbox_core/toolbox_transport.py +++ b/packages/toolbox-core/src/toolbox_core/toolbox_transport.py @@ -40,7 +40,7 @@ def base_url(self) -> str: """The base URL for the transport.""" return self.__base_url - async def _get_manifest( + async def __get_manifest( self, url: str, headers: Optional[Mapping[str, str]] ) -> ManifestSchema: """Helper method to perform GET requests and parse the ManifestSchema.""" @@ -57,7 +57,7 @@ async def tool_get( self, tool_name: str, headers: Optional[Mapping[str, str]] = None ) -> ManifestSchema: url = f"{self.__base_url}/api/tool/{tool_name}" - return await self._get_manifest(url, headers) + return await self.__get_manifest(url, headers) async def tools_list( self, @@ -65,7 +65,7 @@ async def tools_list( headers: Optional[Mapping[str, str]] = None, ) -> ManifestSchema: url = f"{self.__base_url}/api/toolset/{toolset_name or ''}" - return await self._get_manifest(url, headers) + return await self.__get_manifest(url, headers) async def tool_invoke( self, tool_name: str, arguments: dict, headers: Mapping[str, str] From f43909ed8b073fe4be28615004edd9a1a76cc7bd Mon Sep 17 00:00:00 2001 From: Twisha Bansal Date: Thu, 9 Oct 2025 15:37:33 +0530 Subject: [PATCH 22/23] fix tests --- packages/toolbox-core/tests/test_tool.py | 151 ++--------------------- 1 file changed, 7 insertions(+), 144 deletions(-) diff --git a/packages/toolbox-core/tests/test_tool.py b/packages/toolbox-core/tests/test_tool.py index 922604473..ce8c7ca17 100644 --- a/packages/toolbox-core/tests/test_tool.py +++ b/packages/toolbox-core/tests/test_tool.py @@ -413,33 +413,6 @@ def test_tool_init_with_client_headers( assert tool_instance._ToolboxTool__client_headers == static_client_header -def test_tool_init_header_auth_conflict( - http_session, - sample_tool_auth_params, - sample_tool_description, - auth_getters, - auth_header_key, -): - """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) - - with pytest.raises( - ValueError, match=f"Client header\\(s\\) `{auth_header_key}` already registered" - ): - ToolboxTool( - transport=transport, - name="auth_conflict_tool", - description=sample_tool_description, - params=sample_tool_auth_params, - required_authn_params={}, - required_authz_tokens=[], - auth_service_token_getters=auth_getters, - bound_params={}, - client_headers=conflicting_client_header, - ) - - def test_tool_add_auth_token_getters_conflict_with_existing_client_header( http_session: ClientSession, sample_tool_params: list[ParameterSchema], @@ -493,44 +466,22 @@ async def test_auth_token_overrides_client_header( params=sample_tool_params, required_authn_params={}, required_authz_tokens=[], - auth_service_token_getters={}, + auth_service_token_getters={"test-auth": lambda: "value-from-auth-getter-123"}, bound_params={}, - client_headers={}, + client_headers={ + "test-auth_token": "value-from-client", + "X-Another-Header": "another-value", + }, ) - - auth_service_name = "test-auth" - auth_header_key = f"{auth_service_name}_token" - auth_token_value = "value-from-auth-getter-123" - auth_getters = {auth_service_name: lambda: auth_token_value} - tool_name = TEST_TOOL_NAME base_url = HTTPS_BASE_URL invoke_url = f"{base_url}/api/tool/{tool_name}/invoke" - client_headers = { - auth_header_key: "value-from-client", - "X-Another-Header": "another-value", - } - input_args = {"message": "test", "count": 1} mock_server_response = {"result": "Success"} with aioresponses() as m: m.post(invoke_url, status=200, payload=mock_server_response) - - tool_instance = ToolboxTool( - session=http_session, - base_url=base_url, - name=tool_name, - description=sample_tool_description, - params=sample_tool_params, - auth_service_token_getters=auth_getters, - client_headers=client_headers, - required_authn_params={}, - required_authz_tokens=[], - bound_params={}, - ) - # Call the tool await tool_instance(**input_args) @@ -539,7 +490,7 @@ async def test_auth_token_overrides_client_header( method="POST", json=input_args, headers={ - auth_header_key: auth_token_value, + "test-auth_token": "value-from-auth-getter-123", "X-Another-Header": "another-value", }, ) @@ -574,92 +525,4 @@ def test_add_auth_token_getter_unused_token( tool_instance.add_auth_token_getter( next(iter(unused_auth_getters)), unused_auth_getters[next(iter(unused_auth_getters))], - ) - - -def test_toolbox_tool_underscore_name_property(toolbox_tool: ToolboxTool): - """Tests the _name property.""" - assert toolbox_tool._name == TEST_TOOL_NAME - - -def test_toolbox_tool_underscore_description_property(toolbox_tool: ToolboxTool): - """Tests the _description property.""" - assert ( - toolbox_tool._description - == "A sample tool that processes a message and a count." - ) - - -def test_toolbox_tool_underscore_params_property( - toolbox_tool: ToolboxTool, sample_tool_params: list[ParameterSchema] -): - """Tests the _params property returns a deep copy.""" - params_copy = toolbox_tool._params - assert params_copy == sample_tool_params - assert ( - params_copy is not toolbox_tool._ToolboxTool__params - ) # Ensure it's a deepcopy - # Verify modifying the copy does not affect the original - params_copy.append( - ParameterSchema(name="new_param", type="integer", description="A new parameter") - ) - assert ( - len(toolbox_tool._ToolboxTool__params) == 2 - ) # Original should remain unchanged - - -def test_toolbox_tool_underscore_bound_params_property(toolbox_tool: ToolboxTool): - """Tests the _bound_params property returns an immutable MappingProxyType.""" - bound_params = toolbox_tool._bound_params - assert bound_params == {"fixed_param": "fixed_value"} - assert isinstance(bound_params, MappingProxyType) - # Verify immutability - with pytest.raises(TypeError): - bound_params["new_param"] = "new_value" - - -def test_toolbox_tool_underscore_required_authn_params_property( - toolbox_tool: ToolboxTool, -): - """Tests the _required_authn_params property returns an immutable MappingProxyType.""" - required_authn_params = toolbox_tool._required_authn_params - assert required_authn_params == {"message": ["service_a"]} - assert isinstance(required_authn_params, MappingProxyType) - # Verify immutability - with pytest.raises(TypeError): - required_authn_params["new_param"] = ["new_service"] - - -def test_toolbox_tool_underscore_required_authz_tokens_property( - toolbox_tool: ToolboxTool, -): - """Tests the _required_authz_tokens property returns an immutable MappingProxyType.""" - required_authz_tokens = toolbox_tool._required_authz_tokens - assert required_authz_tokens == ("service_b",) - assert isinstance(required_authz_tokens, tuple) - # Verify immutability - with pytest.raises(TypeError): - required_authz_tokens[0] = "new_service" - - -def test_toolbox_tool_underscore_auth_service_token_getters_property( - toolbox_tool: ToolboxTool, -): - """Tests the _auth_service_token_getters property returns an immutable MappingProxyType.""" - auth_getters = toolbox_tool._auth_service_token_getters - assert "service_x" in auth_getters - assert auth_getters["service_x"]() == "token_x" - assert isinstance(auth_getters, MappingProxyType) - # Verify immutability - with pytest.raises(TypeError): - auth_getters["new_service"] = lambda: "new_token" - - -def test_toolbox_tool_underscore_client_headers_property(toolbox_tool: ToolboxTool): - """Tests the _client_headers property returns an immutable MappingProxyType.""" - client_headers = toolbox_tool._client_headers - assert client_headers == {"X-Test-Client": "client_header_value"} - assert isinstance(client_headers, MappingProxyType) - # Verify immutability - with pytest.raises(TypeError): - client_headers["new_header"] = "new_value" + ) \ No newline at end of file From f07d4b34030cf2ab5d45e0850f199c2598ade3e9 Mon Sep 17 00:00:00 2001 From: Twisha Bansal Date: Thu, 9 Oct 2025 15:39:07 +0530 Subject: [PATCH 23/23] lint --- packages/toolbox-core/tests/test_tool.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/toolbox-core/tests/test_tool.py b/packages/toolbox-core/tests/test_tool.py index ce8c7ca17..7708c9381 100644 --- a/packages/toolbox-core/tests/test_tool.py +++ b/packages/toolbox-core/tests/test_tool.py @@ -525,4 +525,4 @@ def test_add_auth_token_getter_unused_token( tool_instance.add_auth_token_getter( next(iter(unused_auth_getters)), unused_auth_getters[next(iter(unused_auth_getters))], - ) \ No newline at end of file + )