diff --git a/packages/toolbox-adk/tests/constants.py b/packages/toolbox-adk/tests/constants.py new file mode 100644 index 000000000..701f830e6 --- /dev/null +++ b/packages/toolbox-adk/tests/constants.py @@ -0,0 +1,16 @@ +# Copyright 2026 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. + +TOOLBOX_SERVER_URL_STABLE = "http://localhost:5000" +TOOLBOX_SERVER_URL_DRAFT = "http://localhost:5001" diff --git a/packages/toolbox-adk/tests/integration/conftest.py b/packages/toolbox-adk/tests/integration/conftest.py index 5c672547a..22254e4cb 100644 --- a/packages/toolbox-adk/tests/integration/conftest.py +++ b/packages/toolbox-adk/tests/integration/conftest.py @@ -29,8 +29,7 @@ from google.auth import compute_engine from google.cloud import secretmanager, storage -TOOLBOX_SERVER_URL_STABLE = "http://localhost:5000" -TOOLBOX_SERVER_URL_DRAFT = "http://localhost:5001" +from tests.constants import TOOLBOX_SERVER_URL_DRAFT, TOOLBOX_SERVER_URL_STABLE #### Define Utility Functions diff --git a/packages/toolbox-adk/tests/integration/test_integration.py b/packages/toolbox-adk/tests/integration/test_integration.py index 461b6133d..af1255de2 100644 --- a/packages/toolbox-adk/tests/integration/test_integration.py +++ b/packages/toolbox-adk/tests/integration/test_integration.py @@ -31,10 +31,9 @@ from pydantic import ValidationError from toolbox_core.protocol import Protocol +from tests.constants import TOOLBOX_SERVER_URL_STABLE from toolbox_adk import CredentialStrategy, ToolboxTool, ToolboxToolset -TOOLBOX_SERVER_URL = "http://localhost:5000" - pytestmark = pytest.mark.usefixtures("patch_toolbox_client_url") # Ensure TOOLBOX_VERSION is set for the fixture @@ -54,7 +53,7 @@ async def test_load_toolset_and_run(self): # Auth: TOOLBOX_IDENTITY for simplicity in this local test as we don't have ADK identity setup. toolset = ToolboxToolset( - server_url=TOOLBOX_SERVER_URL, + server_url=TOOLBOX_SERVER_URL_STABLE, toolset_name="my-toolset", credentials=CredentialStrategy.toolbox_identity(), ) @@ -88,7 +87,7 @@ async def test_load_toolset_and_run(self): async def test_load_toolset_with_default_protocol(self): """Test initializing toolset with default protocol (MCP).""" toolset = ToolboxToolset( - server_url=TOOLBOX_SERVER_URL, + server_url=TOOLBOX_SERVER_URL_STABLE, toolset_name="my-toolset", credentials=CredentialStrategy.toolbox_identity(), ) @@ -112,7 +111,7 @@ async def test_load_toolset_with_default_protocol(self): async def test_load_toolset_with_explicit_protocol(self): """Test initializing toolset with specific protocol (MCP_v20251125).""" toolset = ToolboxToolset( - server_url=TOOLBOX_SERVER_URL, + server_url=TOOLBOX_SERVER_URL_STABLE, toolset_name="my-toolset", credentials=CredentialStrategy.toolbox_identity(), protocol=Protocol.MCP_v20251125, @@ -136,7 +135,7 @@ async def test_load_toolset_with_explicit_protocol(self): async def test_partial_loading_by_names(self): toolset = ToolboxToolset( - server_url=TOOLBOX_SERVER_URL, + server_url=TOOLBOX_SERVER_URL_STABLE, tool_names=["get-n-rows"], credentials=CredentialStrategy.toolbox_identity(), ) @@ -156,7 +155,7 @@ async def test_partial_loading_by_names(self): async def test_bound_params_e2e(self): # Test binding param at toolset level toolset = ToolboxToolset( - server_url=TOOLBOX_SERVER_URL, + server_url=TOOLBOX_SERVER_URL_STABLE, tool_names=["get-n-rows"], bound_params={"num_rows": "2"}, credentials=CredentialStrategy.toolbox_identity(), @@ -172,7 +171,7 @@ async def test_bound_params_e2e(self): async def test_3lo_flow_simulation(self): toolset = ToolboxToolset( - server_url=TOOLBOX_SERVER_URL, + server_url=TOOLBOX_SERVER_URL_STABLE, # Load a specific tool that we know the arguments for tool_names=["get-n-rows"], credentials=CredentialStrategy.user_identity( @@ -269,7 +268,7 @@ async def test_3lo_flow_simulation(self): async def test_manual_token_integration(self): """Test the MANUAL_TOKEN strategy.""" toolset = ToolboxToolset( - server_url=TOOLBOX_SERVER_URL, + server_url=TOOLBOX_SERVER_URL_STABLE, toolset_name="my-toolset", credentials=CredentialStrategy.manual_token(token="fake-manual-token"), ) @@ -288,7 +287,7 @@ async def test_manual_credentials_integration(self): mock_creds.token = "fake-creds-token" toolset = ToolboxToolset( - server_url=TOOLBOX_SERVER_URL, + server_url=TOOLBOX_SERVER_URL_STABLE, toolset_name="my-toolset", credentials=CredentialStrategy.manual_credentials(credentials=mock_creds), ) @@ -301,7 +300,7 @@ async def test_manual_credentials_integration(self): async def test_api_key_integration(self): """Test the API_KEY strategy.""" toolset = ToolboxToolset( - server_url=TOOLBOX_SERVER_URL, + server_url=TOOLBOX_SERVER_URL_STABLE, toolset_name="my-toolset", credentials=CredentialStrategy.api_key(key="my-key", header_name="x-foo"), ) @@ -337,7 +336,7 @@ async def test_adk_integration_optional_params(self): # 3. Use in Toolset toolset = ToolboxToolset( - server_url=TOOLBOX_SERVER_URL, + server_url=TOOLBOX_SERVER_URL_STABLE, toolset_name="my-toolset", credentials=strategy, ) @@ -361,7 +360,7 @@ async def test_header_collision(self): creds_token = "Bearer strategy-token" toolset = ToolboxToolset( - server_url=TOOLBOX_SERVER_URL, + server_url=TOOLBOX_SERVER_URL_STABLE, toolset_name="my-toolset", additional_headers={"Authorization": manual_override}, credentials=CredentialStrategy.manual_token(token="strategy-token"), @@ -395,7 +394,7 @@ async def test_load_toolset_specific( ): """Load a specific toolset""" toolset = ToolboxToolset( - server_url=TOOLBOX_SERVER_URL, + server_url=TOOLBOX_SERVER_URL_STABLE, toolset_name=toolset_name, credentials=CredentialStrategy.toolbox_identity(), ) @@ -410,7 +409,7 @@ async def test_load_toolset_specific( async def test_load_toolset_default(self): """Load the default toolset, i.e. all tools.""" toolset = ToolboxToolset( - server_url=TOOLBOX_SERVER_URL, + server_url=TOOLBOX_SERVER_URL_STABLE, credentials=CredentialStrategy.toolbox_identity(), ) try: @@ -433,7 +432,7 @@ async def test_load_toolset_default(self): async def test_run_tool(self): """Invoke a tool.""" toolset = ToolboxToolset( - server_url=TOOLBOX_SERVER_URL, + server_url=TOOLBOX_SERVER_URL_STABLE, tool_names=["get-n-rows"], credentials=CredentialStrategy.toolbox_identity(), ) @@ -455,7 +454,7 @@ async def test_run_tool(self): async def test_run_tool_missing_params(self): """Invoke a tool with missing params.""" toolset = ToolboxToolset( - server_url=TOOLBOX_SERVER_URL, + server_url=TOOLBOX_SERVER_URL_STABLE, tool_names=["get-n-rows"], credentials=CredentialStrategy.toolbox_identity(), ) @@ -474,7 +473,7 @@ async def test_run_tool_missing_params(self): async def test_run_tool_wrong_param_type(self): """Invoke a tool with wrong param type.""" toolset = ToolboxToolset( - server_url=TOOLBOX_SERVER_URL, + server_url=TOOLBOX_SERVER_URL_STABLE, tool_names=["get-n-rows"], credentials=CredentialStrategy.toolbox_identity(), ) @@ -498,7 +497,7 @@ class TestBindParams: async def test_bind_params(self): """Bind a param to an existing tool.""" toolset = ToolboxToolset( - server_url=TOOLBOX_SERVER_URL, + server_url=TOOLBOX_SERVER_URL_STABLE, tool_names=["get-n-rows"], bound_params={"num_rows": "3"}, credentials=CredentialStrategy.toolbox_identity(), @@ -520,7 +519,7 @@ async def test_bind_params(self): async def test_bind_params_callable(self): """Bind a callable param to an existing tool.""" toolset = ToolboxToolset( - server_url=TOOLBOX_SERVER_URL, + server_url=TOOLBOX_SERVER_URL_STABLE, tool_names=["get-n-rows"], bound_params={"num_rows": lambda: "3"}, credentials=CredentialStrategy.toolbox_identity(), @@ -546,7 +545,7 @@ class TestAuth: async def test_run_tool_unauth_with_auth(self, auth_token2: str): """Tests running a tool that doesn't require auth, with auth provided.""" toolset = ToolboxToolset( - server_url=TOOLBOX_SERVER_URL, + server_url=TOOLBOX_SERVER_URL_STABLE, tool_names=["get-row-by-id"], auth_token_getters={"my-test-auth": lambda: auth_token2}, credentials=CredentialStrategy.toolbox_identity(), @@ -563,7 +562,7 @@ async def test_run_tool_unauth_with_auth(self, auth_token2: str): async def test_run_multiple_tools_unauth_with_auth(self, auth_token2: str): """Tests running multiple tools that don't require auth, verifying formatting of tool lists.""" toolset = ToolboxToolset( - server_url=TOOLBOX_SERVER_URL, + server_url=TOOLBOX_SERVER_URL_STABLE, tool_names=["get-row-by-id", "search-rows"], auth_token_getters={"my-test-auth": lambda: auth_token2}, credentials=CredentialStrategy.toolbox_identity(), @@ -580,7 +579,7 @@ async def test_run_multiple_tools_unauth_with_auth(self, auth_token2: str): async def test_run_multiple_tools_partial_auth_usage(self, auth_token2: str): """Tests that when some tokens are used and some aren't across diverse tools, only the truly unused tokens appear in the error.""" toolset = ToolboxToolset( - server_url=TOOLBOX_SERVER_URL, + server_url=TOOLBOX_SERVER_URL_STABLE, tool_names=[ "get-row-by-id-auth", "search-rows", @@ -605,7 +604,7 @@ async def test_run_tool_no_auth(self): """Tests running a tool requiring auth without providing auth.""" # Note: We load it without auth getters. Invocation should fail. toolset = ToolboxToolset( - server_url=TOOLBOX_SERVER_URL, + server_url=TOOLBOX_SERVER_URL_STABLE, tool_names=["get-row-by-id-auth"], credentials=CredentialStrategy.toolbox_identity(), ) @@ -625,7 +624,7 @@ async def test_run_tool_no_auth(self): async def test_run_tool_wrong_auth(self, auth_token2: str): """Tests running a tool with incorrect auth.""" toolset = ToolboxToolset( - server_url=TOOLBOX_SERVER_URL, + server_url=TOOLBOX_SERVER_URL_STABLE, tool_names=["get-row-by-id-auth"], auth_token_getters={"my-test-auth": lambda: auth_token2}, credentials=CredentialStrategy.toolbox_identity(), @@ -646,7 +645,7 @@ async def test_run_tool_wrong_auth(self, auth_token2: str): async def test_run_tool_auth(self, auth_token1: str): """Tests running a tool with correct auth.""" toolset = ToolboxToolset( - server_url=TOOLBOX_SERVER_URL, + server_url=TOOLBOX_SERVER_URL_STABLE, tool_names=["get-row-by-id-auth"], auth_token_getters={"my-test-auth": lambda: auth_token1}, credentials=CredentialStrategy.toolbox_identity(), @@ -668,7 +667,7 @@ async def get_token_asynchronously(): return auth_token1 toolset = ToolboxToolset( - server_url=TOOLBOX_SERVER_URL, + server_url=TOOLBOX_SERVER_URL_STABLE, tool_names=["get-row-by-id-auth"], auth_token_getters={"my-test-auth": get_token_asynchronously}, credentials=CredentialStrategy.toolbox_identity(), @@ -694,7 +693,7 @@ class TestOptionalParams: async def test_run_tool_with_optional_params_omitted(self): """Invoke a tool providing only the required parameter.""" toolset = ToolboxToolset( - server_url=TOOLBOX_SERVER_URL, + server_url=TOOLBOX_SERVER_URL_STABLE, tool_names=["search-rows"], credentials=CredentialStrategy.toolbox_identity(), ) @@ -713,7 +712,7 @@ async def test_run_tool_with_optional_params_omitted(self): async def test_run_tool_with_all_valid_params(self): """Invoke a tool providing all parameters.""" toolset = ToolboxToolset( - server_url=TOOLBOX_SERVER_URL, + server_url=TOOLBOX_SERVER_URL_STABLE, tool_names=["search-rows"], credentials=CredentialStrategy.toolbox_identity(), ) @@ -733,7 +732,7 @@ async def test_run_tool_with_all_valid_params(self): async def test_run_tool_with_missing_required_param(self): """Invoke a tool without its required parameter.""" toolset = ToolboxToolset( - server_url=TOOLBOX_SERVER_URL, + server_url=TOOLBOX_SERVER_URL_STABLE, tool_names=["search-rows"], credentials=CredentialStrategy.toolbox_identity(), ) @@ -758,7 +757,7 @@ class TestMapParams: async def test_run_tool_with_map_params(self): """Invoke a tool with valid map parameters.""" toolset = ToolboxToolset( - server_url=TOOLBOX_SERVER_URL, + server_url=TOOLBOX_SERVER_URL_STABLE, tool_names=["process-data"], credentials=CredentialStrategy.toolbox_identity(), ) @@ -789,7 +788,7 @@ async def test_run_tool_with_map_params(self): async def test_run_tool_with_wrong_map_value_type(self): """Invoke a tool with a map parameter having the wrong value type.""" toolset = ToolboxToolset( - server_url=TOOLBOX_SERVER_URL, + server_url=TOOLBOX_SERVER_URL_STABLE, tool_names=["process-data"], credentials=CredentialStrategy.toolbox_identity(), ) diff --git a/packages/toolbox-core/README.md b/packages/toolbox-core/README.md index 60daf7580..b7f80aa0a 100644 --- a/packages/toolbox-core/README.md +++ b/packages/toolbox-core/README.md @@ -63,6 +63,7 @@ The core package provides a framework-agnostic way to interact with your Toolbox - [Authenticating Tools](https://mcp-toolbox.dev/documentation/connect-to/toolbox-sdks/python-sdk/core/#authenticating-tools) - [Binding Parameter Values](https://mcp-toolbox.dev/documentation/connect-to/toolbox-sdks/python-sdk/core/#parameter-binding) - [OpenTelemetry](https://mcp-toolbox.dev/documentation/connect-to/toolbox-sdks/python-sdk/core/#opentelemetry) +- [Protocol Negotiation](https://mcp-toolbox.dev/documentation/connect-to/toolbox-sdks/python-sdk/core/#protocol-negotiation) # Contributing diff --git a/packages/toolbox-core/src/toolbox_core/client.py b/packages/toolbox-core/src/toolbox_core/client.py index b2224eedc..b5017a34c 100644 --- a/packages/toolbox-core/src/toolbox_core/client.py +++ b/packages/toolbox-core/src/toolbox_core/client.py @@ -53,12 +53,14 @@ def __init__( client_name: Optional[str], client_version: Optional[str], telemetry_enabled: bool, + supported_protocols: Optional[list[str]] = None, ): self._url = url self._session = session self._client_name = client_name self._client_version = client_version self._telemetry_enabled = telemetry_enabled + self._supported_protocols = supported_protocols self._active_transport = self._create_transport(protocol) def _create_transport(self, protocol: Protocol) -> ITransport: @@ -71,6 +73,7 @@ def _create_transport(self, protocol: Protocol) -> ITransport: self._client_name, self._client_version, telemetry_enabled=self._telemetry_enabled, + supported_protocols=self._supported_protocols, ) case Protocol.MCP_v20251125: return McpHttpTransportV20251125( @@ -80,6 +83,7 @@ def _create_transport(self, protocol: Protocol) -> ITransport: self._client_name, self._client_version, telemetry_enabled=self._telemetry_enabled, + supported_protocols=self._supported_protocols, ) case Protocol.MCP_v20250618: return McpHttpTransportV20250618( @@ -89,6 +93,7 @@ def _create_transport(self, protocol: Protocol) -> ITransport: self._client_name, self._client_version, telemetry_enabled=self._telemetry_enabled, + supported_protocols=self._supported_protocols, ) case Protocol.MCP_v20250326: return McpHttpTransportV20250326( @@ -98,6 +103,7 @@ def _create_transport(self, protocol: Protocol) -> ITransport: self._client_name, self._client_version, telemetry_enabled=self._telemetry_enabled, + supported_protocols=self._supported_protocols, ) case Protocol.MCP_v20241105: return McpHttpTransportV20241105( @@ -107,6 +113,7 @@ def _create_transport(self, protocol: Protocol) -> ITransport: self._client_name, self._client_version, telemetry_enabled=self._telemetry_enabled, + supported_protocols=self._supported_protocols, ) case _: raise ValueError(f"Unsupported MCP protocol version: {protocol}") @@ -126,7 +133,35 @@ async def _execute_with_fallback( try: return await getattr(self._active_transport, method_name)(*args, **kwargs) except ProtocolNegotiationError as e: - fallback_protocol = Protocol(e.negotiated_version) + server_version = e.negotiated_version + all_versions = Protocol.get_supported_mcp_versions() + + try: + server_idx = all_versions.index(server_version) + except ValueError: + raise RuntimeError( + f"Server returned unknown protocol version: {server_version}" + ) + + # Artificial Array: Server supports this version and all older ones + server_supported = all_versions[server_idx:] + + if self._supported_protocols: + client_supported = self._supported_protocols + mutually_supported = [ + v for v in client_supported if v in server_supported + ] + if mutually_supported: + fallback_protocol = Protocol(mutually_supported[0]) + else: + raise RuntimeError( + "No mutually supported protocol version. " + f"Client supports: {client_supported}, " + f"Server supports (and older): {server_version}" + ) + else: + fallback_protocol = Protocol(server_version) + logging.warning( f"Protocol fallback required. Switching from " f"{self._protocol_version} to {fallback_protocol.value}" @@ -166,7 +201,7 @@ def __init__( client_headers: Optional[ Mapping[str, Union[Callable[[], str], Callable[[], Awaitable[str]], str]] ] = None, - protocol: Protocol = Protocol.MCP, + protocol: Union[Protocol, list[Protocol], list[str]] = Protocol.MCP, client_name: Optional[str] = None, client_version: Optional[str] = None, telemetry_enabled: bool = False, @@ -182,19 +217,47 @@ def __init__( should typically be managed externally. client_headers: Headers to include in each request sent through this client. - protocol: The communication protocol to use. + protocol: The communication protocol to use. Can be a single version or a list of versions. + If no version is given, the latest version is tried for. + If a single version is given, then that version is tried for, followed by newest mutually supported version. + If a list of versions is provided (e.g. [Protocol.MCP_LATEST, "2025-11-25"]), negotiation is restricted to those versions, and the newest mutually supported version in that range is tried for. client_name: Optional client name for identification. client_version: Optional client version for identification. telemetry_enabled: Whether to enable OpenTelemetry tracing and metrics. (Default: False) """ + if isinstance(protocol, list): + if not protocol: + raise ValueError("protocol list cannot be empty") + user_protocols = [ + p.value if isinstance(p, Protocol) else str(p) for p in protocol + ] + + supported_mcp_versions = Protocol.get_supported_mcp_versions() + for p in user_protocols: + if p not in supported_mcp_versions: + raise ValueError( + f"Invalid protocol version '{p}'. Must be one of: {supported_mcp_versions}" + ) + + user_protocols_set = set(user_protocols) + # Intersect with the globally sorted list to strictly enforce newest-to-oldest ordering + supported_protocols = [ + v for v in supported_mcp_versions if v in user_protocols_set + ] + initial_protocol = Protocol(supported_protocols[0]) + else: + supported_protocols = None + initial_protocol = protocol + self.__transport = _McpTransportProxy( url, session, - protocol, + initial_protocol, client_name, client_version, telemetry_enabled, + supported_protocols, ) self.__client_headers = client_headers if client_headers is not None else {} diff --git a/packages/toolbox-core/src/toolbox_core/mcp_transport/transport_base.py b/packages/toolbox-core/src/toolbox_core/mcp_transport/transport_base.py index 735ae9066..9ffc2ebc8 100644 --- a/packages/toolbox-core/src/toolbox_core/mcp_transport/transport_base.py +++ b/packages/toolbox-core/src/toolbox_core/mcp_transport/transport_base.py @@ -42,6 +42,7 @@ def __init__( client_name: Optional[str] = None, client_version: Optional[str] = None, telemetry_enabled: bool = False, + supported_protocols: Optional[list[str]] = None, ): self._mcp_base_url = f"{base_url}/mcp/" self._protocol_version = protocol.value @@ -50,6 +51,9 @@ def __init__( self._client_name = client_name self._client_version = client_version self._telemetry_enabled = telemetry.resolve_telemetry_enabled(telemetry_enabled) + self._supported_protocols = ( + supported_protocols or Protocol.get_supported_mcp_versions() + ) self._tracer: Optional[telemetry.Tracer] = None self._operation_duration_histogram: Optional[telemetry.Histogram] = None diff --git a/packages/toolbox-core/src/toolbox_core/mcp_transport/v20250618/mcp.py b/packages/toolbox-core/src/toolbox_core/mcp_transport/v20250618/mcp.py index 2ea3ee636..35f759346 100644 --- a/packages/toolbox-core/src/toolbox_core/mcp_transport/v20250618/mcp.py +++ b/packages/toolbox-core/src/toolbox_core/mcp_transport/v20250618/mcp.py @@ -75,7 +75,7 @@ async def _send_request( err_val = json_resp["error"] if isinstance(err_val, dict) and err_val.get("code") == -32004: server_supported = err_val.get("data", {}).get("supported", []) - client_supported = Protocol.get_supported_mcp_versions() + client_supported = self._supported_protocols mutually_supported = [ v for v in client_supported if v in server_supported ] diff --git a/packages/toolbox-core/src/toolbox_core/mcp_transport/v20251125/mcp.py b/packages/toolbox-core/src/toolbox_core/mcp_transport/v20251125/mcp.py index 4a0e8d876..bf365dece 100644 --- a/packages/toolbox-core/src/toolbox_core/mcp_transport/v20251125/mcp.py +++ b/packages/toolbox-core/src/toolbox_core/mcp_transport/v20251125/mcp.py @@ -75,7 +75,7 @@ async def _send_request( err_val = json_resp["error"] if isinstance(err_val, dict) and err_val.get("code") == -32004: server_supported = err_val.get("data", {}).get("supported", []) - client_supported = Protocol.get_supported_mcp_versions() + client_supported = self._supported_protocols mutually_supported = [ v for v in client_supported if v in server_supported ] diff --git a/packages/toolbox-core/src/toolbox_core/mcp_transport/v20260618/mcp.py b/packages/toolbox-core/src/toolbox_core/mcp_transport/v20260618/mcp.py index 866249733..2955b72e7 100644 --- a/packages/toolbox-core/src/toolbox_core/mcp_transport/v20260618/mcp.py +++ b/packages/toolbox-core/src/toolbox_core/mcp_transport/v20260618/mcp.py @@ -88,7 +88,10 @@ async def _send_request( "supported", [] ) - client_supported = Protocol.get_supported_mcp_versions() + client_supported = ( + self._supported_protocols + or Protocol.get_supported_mcp_versions() + ) mutually_supported = [ v for v in client_supported if v in server_supported ] @@ -105,10 +108,28 @@ async def _send_request( isinstance(err_val, str) and "invalid protocol version" in err_val.lower() ): - # Legacy 2025-06-18 servers don't use the -32004 code or provide - # a supported versions list. They return this raw string error - # instead. We safely assume 2025-06-18 here. - raise ProtocolNegotiationError(Protocol.MCP_v20250618) + # Cascading Fallback: Legacy servers throw this string error. + # We pick the next version from the user's supported list. + client_supported = ( + self._supported_protocols + or Protocol.get_supported_mcp_versions() + ) + try: + current_idx = client_supported.index( + self._protocol_version + ) + if current_idx + 1 < len(client_supported): + raise ProtocolNegotiationError( + client_supported[current_idx + 1] + ) + else: + raise RuntimeError( + "Server threw 'invalid protocol version' but no fallback versions " + "remain in the user's supported protocols array." + ) + except ValueError: + # Current version not in list somehow, just fallback to highest stateful + raise ProtocolNegotiationError(Protocol.MCP_v20251125) except Exception as e: if isinstance(e, (RuntimeError, ProtocolNegotiationError)): raise e @@ -131,7 +152,10 @@ async def _send_request( err_val = json_resp["error"] if isinstance(err_val, dict) and err_val.get("code") == -32004: server_supported = err_val.get("data", {}).get("supported", []) - client_supported = Protocol.get_supported_mcp_versions() + client_supported = ( + self._supported_protocols + or Protocol.get_supported_mcp_versions() + ) mutually_supported = [ v for v in client_supported if v in server_supported ] diff --git a/packages/toolbox-core/src/toolbox_core/sync_client.py b/packages/toolbox-core/src/toolbox_core/sync_client.py index cad63b2a7..e438fd503 100644 --- a/packages/toolbox-core/src/toolbox_core/sync_client.py +++ b/packages/toolbox-core/src/toolbox_core/sync_client.py @@ -42,7 +42,7 @@ def __init__( client_headers: Optional[ Mapping[str, Union[Callable[[], str], Callable[[], Awaitable[str]], str]] ] = None, - protocol: Protocol = Protocol.MCP, + protocol: Union[Protocol, list[Protocol], list[str]] = Protocol.MCP, client_name: Optional[str] = None, client_version: Optional[str] = None, telemetry_enabled: bool = False, @@ -53,7 +53,10 @@ def __init__( Args: url: The base URL for the Toolbox service API (e.g., "http://localhost:5000"). client_headers: Headers to include in each request sent through this client. - protocol: The communication protocol to use. + protocol: The communication protocol to use. Can be a single version or a list of versions. + If no version is given, the latest version is tried for. + If a single version is given, then that version is tried for, followed by newest mutually supported version. + If a list of versions is provided (e.g. [Protocol.MCP_LATEST, "2025-11-25"]), negotiation is restricted to those versions, and the newest mutually supported version in that range is tried for. client_name: Optional client name for identification. client_version: Optional client version for identification. telemetry_enabled: Whether to enable OpenTelemetry tracing and metrics. (Default: False) diff --git a/packages/toolbox-core/tests/conftest.py b/packages/toolbox-core/tests/conftest.py index b78db9db6..b174dc258 100644 --- a/packages/toolbox-core/tests/conftest.py +++ b/packages/toolbox-core/tests/conftest.py @@ -29,6 +29,8 @@ from google.auth import compute_engine from google.cloud import secretmanager, storage +from tests.constants import TOOLBOX_SERVER_URL_DRAFT, TOOLBOX_SERVER_URL_STABLE + TOOLBOX_SERVER_URL_STABLE = "http://localhost:5000" TOOLBOX_SERVER_URL_DRAFT = "http://localhost:5001" diff --git a/packages/toolbox-core/tests/constants.py b/packages/toolbox-core/tests/constants.py new file mode 100644 index 000000000..701f830e6 --- /dev/null +++ b/packages/toolbox-core/tests/constants.py @@ -0,0 +1,16 @@ +# Copyright 2026 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. + +TOOLBOX_SERVER_URL_STABLE = "http://localhost:5000" +TOOLBOX_SERVER_URL_DRAFT = "http://localhost:5001" diff --git a/packages/toolbox-core/tests/test_client.py b/packages/toolbox-core/tests/test_client.py index e28bb03a1..09b56dffe 100644 --- a/packages/toolbox-core/tests/test_client.py +++ b/packages/toolbox-core/tests/test_client.py @@ -21,7 +21,9 @@ import pytest from aiohttp import web -from toolbox_core.client import ToolboxClient +from tests.constants import TOOLBOX_SERVER_URL_STABLE +from toolbox_core.client import ToolboxClient, _McpTransportProxy +from toolbox_core.exceptions import ProtocolNegotiationError from toolbox_core.itransport import ITransport from toolbox_core.protocol import ( ManifestSchema, @@ -814,7 +816,7 @@ def test_toolbox_client_no_warning_on_mcp(): with warnings.catch_warnings(record=True) as w: warnings.simplefilter("always") - client = ToolboxClient("http://localhost:5000", protocol=Protocol.MCP) + client = ToolboxClient(TOOLBOX_SERVER_URL_STABLE, protocol=Protocol.MCP) assert len(w) == 0 @@ -825,6 +827,137 @@ def test_toolbox_client_no_warning_on_explicit_mcp_version(): warnings.simplefilter("always") client = ToolboxClient( - "http://localhost:5000", protocol=Protocol.MCP_v20251125 + TOOLBOX_SERVER_URL_STABLE, protocol=Protocol.MCP_v20251125 ) assert len(w) == 0 + + +def test_toolbox_client_custom_protocols(): + """Test that custom protocols array is correctly parsed and sorted.""" + with patch("toolbox_core.client._McpTransportProxy") as mock_proxy: + client = ToolboxClient( + TOOLBOX_SERVER_URL_STABLE, + protocol=[Protocol.MCP_v20241105, Protocol.MCP_DRAFT, "2025-06-18"], + ) + mock_proxy.assert_called_once() + args, kwargs = mock_proxy.call_args + + # Check initial_protocol + assert args[2] == Protocol.MCP_DRAFT + # Check supported_protocols (must be sorted from newest to oldest) + assert args[6] == ["DRAFT-2026-v1", "2025-06-18", "2024-11-05"] + + +def test_toolbox_client_custom_protocols_invalid(): + """Test that custom protocols array raises error on invalid inputs.""" + + with pytest.raises(ValueError, match="protocol list cannot be empty"): + ToolboxClient(TOOLBOX_SERVER_URL_STABLE, protocol=[]) + + with pytest.raises(ValueError, match="Invalid protocol version 'invalid-version'"): + ToolboxClient(TOOLBOX_SERVER_URL_STABLE, protocol=["invalid-version"]) + + +@pytest.mark.asyncio +async def test_artificial_array(): + """The Artificial Array Test: simulate server returning 2025-03-26, should fallback to 2024-11-05.""" + proxy = _McpTransportProxy( + "http://mock", + None, + Protocol.MCP_DRAFT, + None, + None, + False, + [Protocol.MCP_DRAFT.value, Protocol.MCP_v20241105.value], + ) + + proxy._active_transport.tool_get = AsyncMock( + side_effect=ProtocolNegotiationError(Protocol.MCP_v20250326.value) + ) + + with patch.object(proxy, "_create_transport") as mock_create: + mock_new_transport = AsyncMock() + mock_new_transport.tool_get.return_value = "success" + mock_create.return_value = mock_new_transport + + res = await proxy.tool_get("mock") + + assert res == "success" + mock_create.assert_called_with(Protocol.MCP_v20241105) + + +@pytest.mark.asyncio +async def test_cascading_fallback(): + """The Cascading Fallback Test: simulate server stateless generic error which throws the next stateful version.""" + proxy = _McpTransportProxy( + "http://mock", + None, + Protocol.MCP_DRAFT, + None, + None, + False, + [Protocol.MCP_DRAFT.value, Protocol.MCP_v20251125.value], + ) + + proxy._active_transport.tool_get = AsyncMock( + side_effect=ProtocolNegotiationError(Protocol.MCP_v20251125.value) + ) + + with patch.object(proxy, "_create_transport") as mock_create: + mock_new_transport = AsyncMock() + mock_new_transport.tool_get.return_value = "success" + mock_create.return_value = mock_new_transport + + res = await proxy.tool_get("mock") + + assert res == "success" + mock_create.assert_called_with(Protocol.MCP_v20251125) + + +@pytest.mark.asyncio +async def test_strict_constraint(): + """The Strict Constraint Test: simulate legacy server returning an unsupported old version.""" + proxy = _McpTransportProxy( + "http://mock", + None, + Protocol.MCP_DRAFT, + None, + None, + False, + [Protocol.MCP_DRAFT.value, Protocol.MCP_v20251125.value], + ) + + proxy._active_transport.tool_get = AsyncMock( + side_effect=ProtocolNegotiationError(Protocol.MCP_v20241105.value) + ) + + with pytest.raises(RuntimeError, match="No mutually supported protocol version"): + await proxy.tool_get("mock") + + +@pytest.mark.asyncio +async def test_modern_smart_fallback(): + """The Modern Smart-Fallback Test: simulate modern payload correctly returning pre-intersected fallback.""" + proxy = _McpTransportProxy( + "http://mock", + None, + Protocol.MCP_DRAFT, + None, + None, + False, + [Protocol.MCP_DRAFT.value, Protocol.MCP_v20241105.value], + ) + + proxy._active_transport.tool_get = AsyncMock( + side_effect=ProtocolNegotiationError(Protocol.MCP_v20241105.value) + ) + + with patch.object(proxy, "_create_transport") as mock_create: + mock_new_transport = AsyncMock() + mock_new_transport.tool_get.return_value = "success" + mock_create.return_value = mock_new_transport + + res = await proxy.tool_get("mock") + + assert res == "success" + mock_create.assert_called_with(Protocol.MCP_v20241105) diff --git a/packages/toolbox-core/tests/test_e2e.py b/packages/toolbox-core/tests/test_e2e.py index 9f1f6b01f..920360812 100644 --- a/packages/toolbox-core/tests/test_e2e.py +++ b/packages/toolbox-core/tests/test_e2e.py @@ -19,6 +19,7 @@ import pytest_asyncio from pydantic import ValidationError +from tests.constants import TOOLBOX_SERVER_URL_STABLE from toolbox_core.client import ToolboxClient from toolbox_core.protocol import Protocol from toolbox_core.tool import ToolboxTool @@ -30,7 +31,7 @@ @pytest_asyncio.fixture(scope="function") async def toolbox(): """Creates a ToolboxClient instance shared by all tests in this module.""" - toolbox = ToolboxClient("http://localhost:5000", protocol=Protocol.MCP) + toolbox = ToolboxClient(TOOLBOX_SERVER_URL_STABLE, protocol=Protocol.MCP) try: yield toolbox finally: @@ -114,7 +115,7 @@ async def test_run_tool_wrong_param_type(self, get_n_rows_tool: ToolboxTool): async def test_load_and_run_tool_with_telemetry(self, telemetry_enabled: bool): """Load and invoke a tool with telemetry_enabled=True/False.""" async with ToolboxClient( - "http://localhost:5000", + TOOLBOX_SERVER_URL_STABLE, protocol=Protocol.MCP, telemetry_enabled=telemetry_enabled, ) as toolbox: diff --git a/packages/toolbox-core/tests/test_e2e_mcp.py b/packages/toolbox-core/tests/test_e2e_mcp.py index 890705da4..8e1b7c99d 100644 --- a/packages/toolbox-core/tests/test_e2e_mcp.py +++ b/packages/toolbox-core/tests/test_e2e_mcp.py @@ -19,6 +19,7 @@ import pytest_asyncio from pydantic import ValidationError +from tests.constants import TOOLBOX_SERVER_URL_DRAFT, TOOLBOX_SERVER_URL_STABLE from toolbox_core.client import ToolboxClient from toolbox_core.protocol import Protocol from toolbox_core.tool import ToolboxTool @@ -32,7 +33,7 @@ ) async def toolbox(request): """Creates a ToolboxClient instance shared by all tests in this module.""" - toolbox = ToolboxClient("http://localhost:5000", protocol=Protocol(request.param)) + toolbox = ToolboxClient(TOOLBOX_SERVER_URL_STABLE, protocol=Protocol(request.param)) try: yield toolbox finally: @@ -102,8 +103,8 @@ async def test_run_tool_missing_params(self, get_n_rows_tool: ToolboxTool): async def test_protocol_fallback_e2e(self, toolbox_server_url: str): """Tests that a client using MCP_DRAFT can fallback to an older protocol against a server that doesn't support the draft version.""" - # The E2E server currently does not support DRAFT 2026 on port 5000, so this will trigger a fallback. - # However, port 5001 does support DRAFT 2026. + # The stable server currently does not support DRAFT 2026, so this will trigger a fallback. + # However, the draft server does support DRAFT 2026. async with ToolboxClient( toolbox_server_url, protocol=Protocol.MCP_DRAFT ) as client: @@ -111,7 +112,7 @@ async def test_protocol_fallback_e2e(self, toolbox_server_url: str): response = await tool(num_rows="1") assert "row1" in response # Verify that fallback occurred by checking the transport's final protocol version - if "5001" in toolbox_server_url: + if toolbox_server_url == TOOLBOX_SERVER_URL_DRAFT: assert ( client._ToolboxClient__transport._protocol_version == Protocol.MCP_DRAFT.value @@ -474,9 +475,9 @@ async def test_run_tool_with_wrong_map_value_type(self, toolbox: ToolboxClient): @pytest.mark.asyncio @pytest.mark.usefixtures("toolbox_server") -async def test_mcp_default_protocol(): +async def test_mcp_default_protocol(toolbox_server_url: str): """Verify that omitting the protocol argument defaults correctly and works.""" - async with ToolboxClient("http://localhost:5000") as client: + async with ToolboxClient(toolbox_server_url) as client: tool = await client.load_tool("get-n-rows") response = await tool(num_rows="1") assert "row1" in response @@ -484,11 +485,53 @@ async def test_mcp_default_protocol(): @pytest.mark.asyncio @pytest.mark.usefixtures("toolbox_server") -async def test_mcp_draft_fallback(): +async def test_mcp_draft_fallback(toolbox_server_url: str): """Verify that explicitly using MCP_DRAFT against a server that doesn't support it falls back successfully.""" + async with ToolboxClient(toolbox_server_url, protocol=Protocol.MCP_DRAFT) as client: + tool = await client.load_tool("get-n-rows") + response = await tool(num_rows="1") + assert "row1" in response + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("toolbox_server") +async def test_mcp_latest_protocol(toolbox_server_url: str): + """Verify that explicitly using MCP_LATEST works successfully.""" + async with ToolboxClient( + toolbox_server_url, protocol=Protocol.MCP_LATEST + ) as client: + tool = await client.load_tool("get-n-rows") + response = await tool(num_rows="1") + assert "row1" in response + assert ( + client._ToolboxClient__transport._protocol_version + == Protocol.MCP_LATEST.value + ) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("toolbox_server") +async def test_mcp_custom_protocols_list(toolbox_server_url: str): + """Verify that passing a list of protocols with MCP_LATEST and MCP_DRAFT works successfully.""" async with ToolboxClient( - "http://localhost:5000", protocol=Protocol.MCP_DRAFT + toolbox_server_url, + protocol=[ + Protocol.MCP_v20241105, + Protocol.MCP_v20250326, + Protocol.MCP_LATEST, + Protocol.MCP_DRAFT, + ], ) as client: tool = await client.load_tool("get-n-rows") response = await tool(num_rows="1") assert "row1" in response + if toolbox_server_url == TOOLBOX_SERVER_URL_STABLE: + assert ( + client._ToolboxClient__transport._protocol_version + == Protocol.MCP_LATEST.value + ) + else: + assert ( + client._ToolboxClient__transport._protocol_version + == Protocol.MCP_DRAFT.value + ) diff --git a/packages/toolbox-core/tests/test_sync_e2e.py b/packages/toolbox-core/tests/test_sync_e2e.py index 38b90b131..172db9107 100644 --- a/packages/toolbox-core/tests/test_sync_e2e.py +++ b/packages/toolbox-core/tests/test_sync_e2e.py @@ -14,6 +14,7 @@ import pytest +from tests.constants import TOOLBOX_SERVER_URL_STABLE from toolbox_core.sync_client import ToolboxSyncClient from toolbox_core.sync_tool import ToolboxSyncTool @@ -24,7 +25,7 @@ @pytest.fixture(scope="module") def toolbox(): """Creates a ToolboxSyncClient instance shared by all tests in this module.""" - toolbox = ToolboxSyncClient("http://localhost:5000") + toolbox = ToolboxSyncClient(TOOLBOX_SERVER_URL_STABLE) try: yield toolbox finally: diff --git a/packages/toolbox-langchain/src/toolbox_langchain/async_client.py b/packages/toolbox-langchain/src/toolbox_langchain/async_client.py index 8d389d8c2..eae814e16 100644 --- a/packages/toolbox-langchain/src/toolbox_langchain/async_client.py +++ b/packages/toolbox-langchain/src/toolbox_langchain/async_client.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, Awaitable, Callable, Mapping, Optional, Union +from typing import Any, Awaitable, Callable, Mapping, Optional, Sequence, Union from warnings import warn from aiohttp import ClientSession @@ -35,7 +35,7 @@ def __init__( client_headers: Optional[ Mapping[str, Union[Callable[[], str], Callable[[], Awaitable[str]], str]] ] = None, - protocol: Protocol = Protocol.MCP, + protocol: Union[Protocol, list[Protocol], list[str]] = Protocol.MCP, telemetry_enabled: bool = False, ): """ diff --git a/packages/toolbox-langchain/src/toolbox_langchain/client.py b/packages/toolbox-langchain/src/toolbox_langchain/client.py index e0f607303..c35b61f7f 100644 --- a/packages/toolbox-langchain/src/toolbox_langchain/client.py +++ b/packages/toolbox-langchain/src/toolbox_langchain/client.py @@ -13,7 +13,7 @@ # limitations under the License. from asyncio import to_thread -from typing import Any, Awaitable, Callable, Mapping, Optional, Union +from typing import Any, Awaitable, Callable, Mapping, Optional, Sequence, Union from warnings import warn from toolbox_core.protocol import Protocol @@ -31,7 +31,7 @@ def __init__( client_headers: Optional[ Mapping[str, Union[Callable[[], str], Callable[[], Awaitable[str]], str]] ] = None, - protocol: Protocol = Protocol.MCP, + protocol: Union[Protocol, list[Protocol], list[str]] = Protocol.MCP, telemetry_enabled: bool = False, ) -> None: """ diff --git a/packages/toolbox-langchain/tests/conftest.py b/packages/toolbox-langchain/tests/conftest.py index 05d9a175a..58cf686d2 100644 --- a/packages/toolbox-langchain/tests/conftest.py +++ b/packages/toolbox-langchain/tests/conftest.py @@ -29,8 +29,7 @@ from google.auth import compute_engine from google.cloud import secretmanager, storage -TOOLBOX_SERVER_URL_STABLE = "http://localhost:5000" -TOOLBOX_SERVER_URL_DRAFT = "http://localhost:5001" +from tests.constants import TOOLBOX_SERVER_URL_DRAFT, TOOLBOX_SERVER_URL_STABLE #### Define Utility Functions diff --git a/packages/toolbox-langchain/tests/constants.py b/packages/toolbox-langchain/tests/constants.py new file mode 100644 index 000000000..701f830e6 --- /dev/null +++ b/packages/toolbox-langchain/tests/constants.py @@ -0,0 +1,16 @@ +# Copyright 2026 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. + +TOOLBOX_SERVER_URL_STABLE = "http://localhost:5000" +TOOLBOX_SERVER_URL_DRAFT = "http://localhost:5001" diff --git a/packages/toolbox-langchain/tests/test_e2e.py b/packages/toolbox-langchain/tests/test_e2e.py index 2a4a20bef..22f216d18 100644 --- a/packages/toolbox-langchain/tests/test_e2e.py +++ b/packages/toolbox-langchain/tests/test_e2e.py @@ -37,15 +37,13 @@ import pytest import pytest_asyncio from pydantic import ValidationError +from toolbox_core.protocol import Protocol +from tests.constants import TOOLBOX_SERVER_URL_STABLE from toolbox_langchain.client import ToolboxClient pytestmark = pytest.mark.usefixtures("patch_toolbox_client_url") -TOOLBOX_SERVER_URL = "http://localhost:5000" - -pytestmark = pytest.mark.usefixtures("patch_toolbox_client_url") - @pytest.mark.asyncio @pytest.mark.usefixtures("toolbox_server") @@ -53,7 +51,7 @@ class TestE2EClientAsync: @pytest.fixture(scope="function") def toolbox(self): """Provides a ToolboxClient instance for each test.""" - toolbox = ToolboxClient(TOOLBOX_SERVER_URL) + toolbox = ToolboxClient(TOOLBOX_SERVER_URL_STABLE) return toolbox @pytest_asyncio.fixture(scope="function") @@ -95,6 +93,14 @@ async def test_aload_toolset_all(self, toolbox): name = tool._ToolboxTool__core_tool.__name__ assert name in tool_names + async def test_aload_toolset_explicit_protocol(self): + toolbox = ToolboxClient( + TOOLBOX_SERVER_URL_STABLE, protocol=Protocol.MCP_v20251125 + ) + toolset = await toolbox.aload_toolset() + assert len(toolset) == 7 + toolbox.close() + async def test_run_tool_async(self, get_n_rows_tool): response = await get_n_rows_tool.ainvoke({"num_rows": "2"}) @@ -200,7 +206,7 @@ class TestE2EClientSync: @pytest.fixture(scope="session") def toolbox(self): """Provides a ToolboxClient instance for each test.""" - toolbox = ToolboxClient(TOOLBOX_SERVER_URL) + toolbox = ToolboxClient(TOOLBOX_SERVER_URL_STABLE) return toolbox @pytest.fixture(scope="function") @@ -242,6 +248,14 @@ def test_aload_toolset_all(self, toolbox): name = tool._ToolboxTool__core_tool.__name__ assert name in tool_names + def test_load_toolset_explicit_protocol(self): + toolbox = ToolboxClient( + TOOLBOX_SERVER_URL_STABLE, protocol=Protocol.MCP_v20251125 + ) + toolset = toolbox.load_toolset() + assert len(toolset) == 7 + toolbox.close() + @pytest.mark.asyncio async def test_run_tool_async(self, get_n_rows_tool): response = await get_n_rows_tool.ainvoke({"num_rows": "2"}) diff --git a/packages/toolbox-llamaindex/src/toolbox_llamaindex/async_client.py b/packages/toolbox-llamaindex/src/toolbox_llamaindex/async_client.py index e38563ee0..08ebc3ec5 100644 --- a/packages/toolbox-llamaindex/src/toolbox_llamaindex/async_client.py +++ b/packages/toolbox-llamaindex/src/toolbox_llamaindex/async_client.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, Awaitable, Callable, Mapping, Optional, Union +from typing import Any, Awaitable, Callable, Mapping, Optional, Sequence, Union from warnings import warn from aiohttp import ClientSession @@ -35,7 +35,7 @@ def __init__( client_headers: Optional[ Mapping[str, Union[Callable[[], str], Callable[[], Awaitable[str]], str]] ] = None, - protocol: Protocol = Protocol.MCP, + protocol: Union[Protocol, list[Protocol], list[str]] = Protocol.MCP, telemetry_enabled: bool = False, ): """ diff --git a/packages/toolbox-llamaindex/src/toolbox_llamaindex/client.py b/packages/toolbox-llamaindex/src/toolbox_llamaindex/client.py index 8cd7d16ce..b6c9d0db3 100644 --- a/packages/toolbox-llamaindex/src/toolbox_llamaindex/client.py +++ b/packages/toolbox-llamaindex/src/toolbox_llamaindex/client.py @@ -13,7 +13,7 @@ # limitations under the License. from asyncio import to_thread -from typing import Any, Awaitable, Callable, Mapping, Optional, Union +from typing import Any, Awaitable, Callable, Mapping, Optional, Sequence, Union from warnings import warn from toolbox_core.protocol import Protocol @@ -32,7 +32,7 @@ def __init__( client_headers: Optional[ Mapping[str, Union[Callable[[], str], Callable[[], Awaitable[str]], str]] ] = None, - protocol: Protocol = Protocol.MCP, + protocol: Union[Protocol, list[Protocol], list[str]] = Protocol.MCP, telemetry_enabled: bool = False, ) -> None: """ diff --git a/packages/toolbox-llamaindex/tests/conftest.py b/packages/toolbox-llamaindex/tests/conftest.py index fa9784f1f..c505933e3 100644 --- a/packages/toolbox-llamaindex/tests/conftest.py +++ b/packages/toolbox-llamaindex/tests/conftest.py @@ -29,8 +29,7 @@ from google.auth import compute_engine from google.cloud import secretmanager, storage -TOOLBOX_SERVER_URL_STABLE = "http://localhost:5000" -TOOLBOX_SERVER_URL_DRAFT = "http://localhost:5001" +from tests.constants import TOOLBOX_SERVER_URL_DRAFT, TOOLBOX_SERVER_URL_STABLE #### Define Utility Functions diff --git a/packages/toolbox-llamaindex/tests/constants.py b/packages/toolbox-llamaindex/tests/constants.py new file mode 100644 index 000000000..701f830e6 --- /dev/null +++ b/packages/toolbox-llamaindex/tests/constants.py @@ -0,0 +1,16 @@ +# Copyright 2026 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. + +TOOLBOX_SERVER_URL_STABLE = "http://localhost:5000" +TOOLBOX_SERVER_URL_DRAFT = "http://localhost:5001" diff --git a/packages/toolbox-llamaindex/tests/test_e2e.py b/packages/toolbox-llamaindex/tests/test_e2e.py index 4b0b2197c..225f66666 100644 --- a/packages/toolbox-llamaindex/tests/test_e2e.py +++ b/packages/toolbox-llamaindex/tests/test_e2e.py @@ -37,15 +37,13 @@ import pytest import pytest_asyncio from pydantic import ValidationError +from toolbox_core.protocol import Protocol +from tests.constants import TOOLBOX_SERVER_URL_STABLE from toolbox_llamaindex.client import ToolboxClient pytestmark = pytest.mark.usefixtures("patch_toolbox_client_url") -TOOLBOX_SERVER_URL = "http://localhost:5000" - -pytestmark = pytest.mark.usefixtures("patch_toolbox_client_url") - @pytest.mark.asyncio @pytest.mark.usefixtures("toolbox_server") @@ -53,7 +51,7 @@ class TestE2EClientAsync: @pytest.fixture(scope="function") def toolbox(self): """Provides a ToolboxClient instance for each test.""" - toolbox = ToolboxClient(TOOLBOX_SERVER_URL) + toolbox = ToolboxClient(TOOLBOX_SERVER_URL_STABLE) return toolbox @pytest_asyncio.fixture(scope="function") @@ -95,6 +93,14 @@ async def test_aload_toolset_all(self, toolbox): name = tool._ToolboxTool__core_tool.__name__ assert name in tool_names + async def test_aload_toolset_explicit_protocol(self): + toolbox = ToolboxClient( + TOOLBOX_SERVER_URL_STABLE, protocol=Protocol.MCP_v20251125 + ) + toolset = await toolbox.aload_toolset() + assert len(toolset) == 7 + toolbox.close() + async def test_run_tool_async(self, get_n_rows_tool): response = await get_n_rows_tool.acall(num_rows="2") @@ -200,7 +206,7 @@ class TestE2EClientSync: @pytest.fixture(scope="session") def toolbox(self): """Provides a ToolboxClient instance for each test.""" - toolbox = ToolboxClient(TOOLBOX_SERVER_URL) + toolbox = ToolboxClient(TOOLBOX_SERVER_URL_STABLE) return toolbox @pytest.fixture(scope="function") @@ -242,6 +248,14 @@ def test_aload_toolset_all(self, toolbox): name = tool._ToolboxTool__core_tool.__name__ assert name in tool_names + def test_load_toolset_explicit_protocol(self): + toolbox = ToolboxClient( + TOOLBOX_SERVER_URL_STABLE, protocol=Protocol.MCP_v20251125 + ) + toolset = toolbox.load_toolset() + assert len(toolset) == 7 + toolbox.close() + @pytest.mark.asyncio async def test_run_tool_async(self, get_n_rows_tool): response = await get_n_rows_tool.acall(num_rows="2")