diff --git a/packages/toolbox-adk/integration.cloudbuild.yaml b/packages/toolbox-adk/integration.cloudbuild.yaml index ff1b1c1f5..3bdf3fa4a 100644 --- a/packages/toolbox-adk/integration.cloudbuild.yaml +++ b/packages/toolbox-adk/integration.cloudbuild.yaml @@ -50,5 +50,5 @@ options: substitutions: _VERSION: '3.13' # Default values (can be overridden by triggers) - _TOOLBOX_VERSION: '1.4.0' + _TOOLBOX_VERSION: 'v1.6.0' _TOOLBOX_MANIFEST_VERSION: '34' diff --git a/packages/toolbox-adk/pyproject.toml b/packages/toolbox-adk/pyproject.toml index 3ad5c35e4..52072abf8 100644 --- a/packages/toolbox-adk/pyproject.toml +++ b/packages/toolbox-adk/pyproject.toml @@ -40,6 +40,7 @@ test = [ "pytest==9.0.3", "pytest-asyncio==1.4.0", "pytest-cov==7.1.0", + "numpy<2.2.0", "pytest-mock==3.15.1" ] diff --git a/packages/toolbox-adk/src/toolbox_adk/tool.py b/packages/toolbox-adk/src/toolbox_adk/tool.py index db41f81b9..0d94821b6 100644 --- a/packages/toolbox-adk/src/toolbox_adk/tool.py +++ b/packages/toolbox-adk/src/toolbox_adk/tool.py @@ -17,11 +17,7 @@ from typing import Any, Awaitable, Callable, Dict, Mapping, Optional import toolbox_core -from fastapi.openapi.models import ( - OAuth2, - OAuthFlowAuthorizationCode, - OAuthFlows, -) +from fastapi.openapi.models import OAuth2, OAuthFlowAuthorizationCode, OAuthFlows from google.adk.auth.auth_credential import ( AuthCredential, AuthCredentialTypes, diff --git a/packages/toolbox-adk/tests/integration/conftest.py b/packages/toolbox-adk/tests/integration/conftest.py index 7a0346c3a..20aea226d 100644 --- a/packages/toolbox-adk/tests/integration/conftest.py +++ b/packages/toolbox-adk/tests/integration/conftest.py @@ -25,7 +25,7 @@ from typing import Generator import google -import pytest_asyncio +import pytest from google.auth import compute_engine from google.cloud import secretmanager, storage @@ -75,7 +75,8 @@ def get_toolbox_binary_url(toolbox_version: str) -> str: arch = ( "arm64" if os_system == "darwin" and platform.machine() == "arm64" else "amd64" ) - return f"v{toolbox_version}/{os_system}/{arch}/toolbox" + ext = ".exe" if os_system == "windows" else "" + return f"{toolbox_version}/{os_system}/{arch}/toolbox{ext}" def get_auth_token(client_id: str) -> str: @@ -92,17 +93,17 @@ def get_auth_token(client_id: str) -> str: #### Define Fixtures -@pytest_asyncio.fixture(scope="session") +@pytest.fixture(scope="session") def project_id() -> str: return get_env_var("GOOGLE_CLOUD_PROJECT") -@pytest_asyncio.fixture(scope="session") +@pytest.fixture(scope="session") def toolbox_version() -> str: return get_env_var("TOOLBOX_VERSION") -@pytest_asyncio.fixture(scope="session") +@pytest.fixture(scope="session") def tools_file_path(project_id: str) -> Generator[str]: """Provides a temporary file path containing the tools manifest.""" if os.environ.get("TEST_MOCK_GCP"): @@ -122,7 +123,7 @@ def tools_file_path(project_id: str) -> Generator[str]: os.remove(tools_file_path) -@pytest_asyncio.fixture(scope="session") +@pytest.fixture(scope="session") def auth_token1(project_id: str) -> str: if os.environ.get("TEST_MOCK_GCP"): return "mock-token-1" @@ -132,7 +133,7 @@ def auth_token1(project_id: str) -> str: return get_auth_token(client_id) -@pytest_asyncio.fixture(scope="session") +@pytest.fixture(scope="session") def auth_token2(project_id: str) -> str: if os.environ.get("TEST_MOCK_GCP"): return "mock-token-2" @@ -142,7 +143,7 @@ def auth_token2(project_id: str) -> str: return get_auth_token(client_id) -@pytest_asyncio.fixture(scope="session") +@pytest.fixture(scope="session") def toolbox_server(toolbox_version: str, tools_file_path: str) -> Generator[None]: """Starts the toolbox server as a subprocess.""" if os.environ.get("TEST_MOCK_GCP"): @@ -151,7 +152,12 @@ def toolbox_server(toolbox_version: str, tools_file_path: str) -> Generator[None print("Downloading toolbox binary from gcs bucket...") source_blob_name = get_toolbox_binary_url(toolbox_version) - download_blob("mcp-toolbox-for-databases", source_blob_name, "toolbox") + bucket_name = ( + "mcp-toolbox-for-databases-dev" + if toolbox_version in ("main", "mcp-v202606") + else "mcp-toolbox-for-databases" + ) + download_blob(bucket_name, source_blob_name, "toolbox") print("Toolbox binary downloaded successfully.") try: @@ -159,20 +165,30 @@ def toolbox_server(toolbox_version: str, tools_file_path: str) -> Generator[None # Make toolbox executable os.chmod("toolbox", 0o700) # Run toolbox binary - toolbox_server = subprocess.Popen( - ["./toolbox", "--tools-file", tools_file_path] + toolbox_server_1 = subprocess.Popen( + ["./toolbox", "--port", "5000", "--tools-file", tools_file_path] + ) + toolbox_server_2 = subprocess.Popen( + [ + "./toolbox", + "--port", + "5001", + "--tools-file", + tools_file_path, + "--enable-draft-specs", + ] ) # Wait for server to start # Retry logic with a timeout for _ in range(5): # retries time.sleep(2) - print("Checking if toolbox is successfully started...") - if toolbox_server.poll() is None: - print("Toolbox server started successfully.") + print("Checking if both toolbox servers are successfully started...") + if toolbox_server_1.poll() is None and toolbox_server_2.poll() is None: + print("Toolbox servers started successfully.") break else: - raise RuntimeError("Toolbox server failed to start after 5 retries.") + raise RuntimeError("Toolbox servers failed to start after 5 retries.") except subprocess.CalledProcessError as e: print(e.stderr.decode("utf-8")) print(e.stdout.decode("utf-8")) @@ -180,5 +196,31 @@ def toolbox_server(toolbox_version: str, tools_file_path: str) -> Generator[None yield # Clean up toolbox server - toolbox_server.terminate() - toolbox_server.wait(timeout=5) + toolbox_server_1.terminate() + toolbox_server_2.terminate() + toolbox_server_1.wait(timeout=5) + toolbox_server_2.wait(timeout=5) + + +@pytest.fixture( + params=["http://localhost:5000", "http://localhost:5001"], scope="session" +) +def toolbox_server_url(request) -> str: + return request.param + + +@pytest.fixture(autouse=True) +def patch_toolbox_client_url(toolbox_server_url): + from toolbox_adk.toolset import ToolboxToolset + + original_init = ToolboxToolset.__init__ + + def new_init(self, server_url="http://localhost:5000", *args, **kwargs): + if server_url == "http://localhost:5000": + server_url = toolbox_server_url + original_init(self, server_url, *args, **kwargs) + + from unittest.mock import patch + + with patch.object(ToolboxToolset, "__init__", new_init, create=True): + yield diff --git a/packages/toolbox-adk/tests/unit/test_credentials.py b/packages/toolbox-adk/tests/unit/test_credentials.py index 749d28667..9e94d5638 100644 --- a/packages/toolbox-adk/tests/unit/test_credentials.py +++ b/packages/toolbox-adk/tests/unit/test_credentials.py @@ -109,10 +109,7 @@ def test_from_adk_credentials_http_bearer(self): def test_from_adk_credentials_api_key(self): from fastapi.openapi.models import APIKey, APIKeyIn - from google.adk.auth.auth_credential import ( - AuthCredential, - AuthCredentialTypes, - ) + from google.adk.auth.auth_credential import AuthCredential, AuthCredentialTypes auth_credential = AuthCredential( auth_type=AuthCredentialTypes.API_KEY, api_key="abc" @@ -129,10 +126,7 @@ def test_from_adk_credentials_api_key(self): def test_from_adk_credentials_api_key_default_location(self): from fastapi.openapi.models import APIKey - from google.adk.auth.auth_credential import ( - AuthCredential, - AuthCredentialTypes, - ) + from google.adk.auth.auth_credential import AuthCredential, AuthCredentialTypes auth_credential = AuthCredential( auth_type=AuthCredentialTypes.API_KEY, api_key="abc" @@ -157,10 +151,7 @@ class MockScheme: def test_from_adk_credentials_api_key_query_fail(self): import pytest from fastapi.openapi.models import APIKey, APIKeyIn - from google.adk.auth.auth_credential import ( - AuthCredential, - AuthCredentialTypes, - ) + from google.adk.auth.auth_credential import AuthCredential, AuthCredentialTypes cred = AuthCredential(auth_type=AuthCredentialTypes.API_KEY, api_key="abc") scheme = APIKey(type="apiKey", name="key", **{"in": APIKeyIn.query}) @@ -172,10 +163,7 @@ def test_from_adk_credentials_api_key_query_fail(self): def test_from_adk_credentials_api_key_no_scheme_raises(self): import pytest - from google.adk.auth.auth_credential import ( - AuthCredential, - AuthCredentialTypes, - ) + from google.adk.auth.auth_credential import AuthCredential, AuthCredentialTypes auth_credential = AuthCredential( auth_type=AuthCredentialTypes.API_KEY, api_key="my-key" @@ -187,10 +175,7 @@ def test_from_adk_credentials_api_key_no_scheme_raises(self): def test_from_adk_credentials_unsupported(self): import pytest - from google.adk.auth.auth_credential import ( - AuthCredential, - AuthCredentialTypes, - ) + from google.adk.auth.auth_credential import AuthCredential, AuthCredentialTypes auth_credential = AuthCredential( auth_type=AuthCredentialTypes.OAUTH2 diff --git a/packages/toolbox-core/integration.cloudbuild.yaml b/packages/toolbox-core/integration.cloudbuild.yaml index de5f69465..a384d4de7 100644 --- a/packages/toolbox-core/integration.cloudbuild.yaml +++ b/packages/toolbox-core/integration.cloudbuild.yaml @@ -45,5 +45,5 @@ options: logging: CLOUD_LOGGING_ONLY substitutions: _VERSION: '3.13' - _TOOLBOX_VERSION: '1.4.0' + _TOOLBOX_VERSION: 'v1.6.0' _TOOLBOX_MANIFEST_VERSION: '34' diff --git a/packages/toolbox-core/pyproject.toml b/packages/toolbox-core/pyproject.toml index 6c2bf1a6d..21237a970 100644 --- a/packages/toolbox-core/pyproject.toml +++ b/packages/toolbox-core/pyproject.toml @@ -56,6 +56,7 @@ test = [ "pytest-aioresponses==0.3.0", "pytest-asyncio==1.4.0", "pytest-cov==7.1.0", + "numpy<2.2.0", "pytest-mock==3.15.1", "google-cloud-secret-manager==2.28.0", "google-cloud-storage==3.10.1", diff --git a/packages/toolbox-core/src/toolbox_core/client.py b/packages/toolbox-core/src/toolbox_core/client.py index 641e80599..b2224eedc 100644 --- a/packages/toolbox-core/src/toolbox_core/client.py +++ b/packages/toolbox-core/src/toolbox_core/client.py @@ -63,7 +63,7 @@ def __init__( def _create_transport(self, protocol: Protocol) -> ITransport: match protocol: - case Protocol.MCP_LATEST: + case Protocol.MCP_DRAFT: return McpHttpTransportV20260618( self._url, self._session, diff --git a/packages/toolbox-core/src/toolbox_core/mcp_transport/v20241105/mcp.py b/packages/toolbox-core/src/toolbox_core/mcp_transport/v20241105/mcp.py index ec6a2475d..7f03dcee0 100644 --- a/packages/toolbox-core/src/toolbox_core/mcp_transport/v20241105/mcp.py +++ b/packages/toolbox-core/src/toolbox_core/mcp_transport/v20241105/mcp.py @@ -18,6 +18,7 @@ from pydantic import BaseModel from ... import version +from ...exceptions import ProtocolNegotiationError from ...protocol import ManifestSchema, TelemetryAttributes from .. import telemetry from ..transport_base import _McpHttpTransportBase @@ -131,9 +132,7 @@ async def _initialize_session( self._server_version = result.serverInfo.version if result.protocolVersion != self._protocol_version: - raise RuntimeError( - f"MCP version mismatch: client does not support server version {result.protocolVersion}" - ) + raise ProtocolNegotiationError(result.protocolVersion) if not result.capabilities.tools: if self._manage_session: diff --git a/packages/toolbox-core/src/toolbox_core/mcp_transport/v20250326/mcp.py b/packages/toolbox-core/src/toolbox_core/mcp_transport/v20250326/mcp.py index 0194bd282..a28f9db0e 100644 --- a/packages/toolbox-core/src/toolbox_core/mcp_transport/v20250326/mcp.py +++ b/packages/toolbox-core/src/toolbox_core/mcp_transport/v20250326/mcp.py @@ -18,6 +18,7 @@ from pydantic import BaseModel from ... import version +from ...exceptions import ProtocolNegotiationError from ...protocol import ManifestSchema, TelemetryAttributes from .. import telemetry from ..transport_base import _McpHttpTransportBase @@ -147,10 +148,7 @@ async def _initialize_session( self._server_version = result.serverInfo.version if result.protocolVersion != self._protocol_version: - raise RuntimeError( - "MCP version mismatch: client does not support server version" - f" {result.protocolVersion}" - ) + raise ProtocolNegotiationError(result.protocolVersion) if not result.capabilities.tools: if self._manage_session: 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 d7a626ed1..11a471ecf 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 @@ -18,6 +18,7 @@ from pydantic import BaseModel from ... import version +from ...exceptions import ProtocolNegotiationError from ...protocol import ManifestSchema, TelemetryAttributes from .. import telemetry from ..transport_base import _McpHttpTransportBase @@ -138,10 +139,7 @@ async def _initialize_session( self._server_version = result.serverInfo.version if result.protocolVersion != self._protocol_version: - raise RuntimeError( - "MCP version mismatch: client does not support server version" - f" {result.protocolVersion}" - ) + raise ProtocolNegotiationError(result.protocolVersion) if not result.capabilities.tools: if self._manage_session: 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 999301552..1bf741e76 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 @@ -18,6 +18,7 @@ from pydantic import BaseModel from ... import version +from ...exceptions import ProtocolNegotiationError from ...protocol import ManifestSchema, TelemetryAttributes from .. import telemetry from ..transport_base import _McpHttpTransportBase @@ -138,10 +139,7 @@ async def _initialize_session( self._server_version = result.serverInfo.version if result.protocolVersion != self._protocol_version: - raise RuntimeError( - "MCP version mismatch: client does not support server version" - f" {result.protocolVersion}" - ) + raise ProtocolNegotiationError(result.protocolVersion) if not result.capabilities.tools: if self._manage_session: diff --git a/packages/toolbox-core/src/toolbox_core/protocol.py b/packages/toolbox-core/src/toolbox_core/protocol.py index 57d12fb33..350e0251f 100644 --- a/packages/toolbox-core/src/toolbox_core/protocol.py +++ b/packages/toolbox-core/src/toolbox_core/protocol.py @@ -51,14 +51,17 @@ class Protocol(str, Enum): MCP_v20250326 = "2025-03-26" MCP_v20241105 = "2024-11-05" MCP_v20251125 = "2025-11-25" - MCP = MCP_v20250618 - MCP_LATEST = "DRAFT-2026-v1" + MCP_v2026_DRAFT = "DRAFT-2026-v1" + + MCP = MCP_v20251125 + MCP_LATEST = MCP_v20251125 + MCP_DRAFT = MCP_v2026_DRAFT @staticmethod def get_supported_mcp_versions() -> list[str]: """Returns a list of supported MCP protocol versions.""" return [ - Protocol.MCP_LATEST.value, + Protocol.MCP_DRAFT.value, Protocol.MCP_v20251125.value, Protocol.MCP_v20250618.value, Protocol.MCP_v20250326.value, diff --git a/packages/toolbox-core/tests/conftest.py b/packages/toolbox-core/tests/conftest.py index 338d031c0..2e500752e 100644 --- a/packages/toolbox-core/tests/conftest.py +++ b/packages/toolbox-core/tests/conftest.py @@ -25,7 +25,7 @@ from typing import Generator import google -import pytest_asyncio +import pytest from google.auth import compute_engine from google.cloud import secretmanager, storage @@ -75,7 +75,8 @@ def get_toolbox_binary_url(toolbox_version: str) -> str: arch = ( "arm64" if os_system == "darwin" and platform.machine() == "arm64" else "amd64" ) - return f"v{toolbox_version}/{os_system}/{arch}/toolbox" + ext = ".exe" if os_system == "windows" else "" + return f"{toolbox_version}/{os_system}/{arch}/toolbox{ext}" def get_auth_token(client_id: str) -> str: @@ -92,17 +93,17 @@ def get_auth_token(client_id: str) -> str: #### Define Fixtures -@pytest_asyncio.fixture(scope="session") +@pytest.fixture(scope="session") def project_id() -> str: return get_env_var("GOOGLE_CLOUD_PROJECT") -@pytest_asyncio.fixture(scope="session") +@pytest.fixture(scope="session") def toolbox_version() -> str: return get_env_var("TOOLBOX_VERSION") -@pytest_asyncio.fixture(scope="session") +@pytest.fixture(scope="session") def tools_file_path(project_id: str) -> Generator[str]: """Provides a temporary file path containing the tools manifest.""" tools_manifest = access_secret_version( @@ -115,7 +116,7 @@ def tools_file_path(project_id: str) -> Generator[str]: os.remove(tools_file_path) -@pytest_asyncio.fixture(scope="session") +@pytest.fixture(scope="session") def auth_token1(project_id: str) -> str: client_id = access_secret_version( project_id=project_id, secret_id="sdk_testing_client1" @@ -123,7 +124,7 @@ def auth_token1(project_id: str) -> str: return get_auth_token(client_id) -@pytest_asyncio.fixture(scope="session") +@pytest.fixture(scope="session") def auth_token2(project_id: str) -> str: client_id = access_secret_version( project_id=project_id, secret_id="sdk_testing_client2" @@ -131,32 +132,47 @@ def auth_token2(project_id: str) -> str: return get_auth_token(client_id) -@pytest_asyncio.fixture(scope="session") +@pytest.fixture(scope="session") def toolbox_server(toolbox_version: str, tools_file_path: str) -> Generator[None]: """Starts the toolbox server as a subprocess.""" print("Downloading toolbox binary from gcs bucket...") source_blob_name = get_toolbox_binary_url(toolbox_version) - download_blob("mcp-toolbox-for-databases", source_blob_name, "toolbox") + bucket_name = ( + "mcp-toolbox-for-databases-dev" + if toolbox_version in ("main", "mcp-v202606") + else "mcp-toolbox-for-databases" + ) + download_blob(bucket_name, source_blob_name, "toolbox") print("Toolbox binary downloaded successfully.") try: print("Opening toolbox server process...") # Make toolbox executable os.chmod("toolbox", 0o700) # Run toolbox binary - toolbox_server = subprocess.Popen( - ["./toolbox", "--tools-file", tools_file_path] + toolbox_server_1 = subprocess.Popen( + ["./toolbox", "--port", "5000", "--tools-file", tools_file_path] + ) + toolbox_server_2 = subprocess.Popen( + [ + "./toolbox", + "--port", + "5001", + "--tools-file", + tools_file_path, + "--enable-draft-specs", + ] ) # Wait for server to start # Retry logic with a timeout for _ in range(5): # retries time.sleep(2) - print("Checking if toolbox is successfully started...") - if toolbox_server.poll() is None: - print("Toolbox server started successfully.") + print("Checking if both toolbox servers are successfully started...") + if toolbox_server_1.poll() is None and toolbox_server_2.poll() is None: + print("Toolbox servers started successfully.") break else: - raise RuntimeError("Toolbox server failed to start after 5 retries.") + raise RuntimeError("Toolbox servers failed to start after 5 retries.") except subprocess.CalledProcessError as e: print(e.stderr.decode("utf-8")) print(e.stdout.decode("utf-8")) @@ -164,5 +180,39 @@ def toolbox_server(toolbox_version: str, tools_file_path: str) -> Generator[None yield # Clean up toolbox server - toolbox_server.terminate() - toolbox_server.wait(timeout=5) + toolbox_server_1.terminate() + toolbox_server_2.terminate() + toolbox_server_1.wait(timeout=5) + toolbox_server_2.wait(timeout=5) + + +@pytest.fixture( + params=["http://localhost:5000", "http://localhost:5001"], scope="session" +) +def toolbox_server_url(request) -> str: + return request.param + + +@pytest.fixture(autouse=True) +def patch_toolbox_client_url(toolbox_server_url): + from toolbox_core.client import ToolboxClient + from toolbox_core.sync_client import ToolboxSyncClient + + original_init = ToolboxClient.__init__ + original_sync_init = ToolboxSyncClient.__init__ + + def new_init(self, url="http://localhost:5000", *args, **kwargs): + if url == "http://localhost:5000": + url = toolbox_server_url + original_init(self, url, *args, **kwargs) + + def new_sync_init(self, url="http://localhost:5000", *args, **kwargs): + if url == "http://localhost:5000": + url = toolbox_server_url + original_sync_init(self, url, *args, **kwargs) + + from unittest.mock import patch + + with patch.object(ToolboxClient, "__init__", new_init, create=True): + with patch.object(ToolboxSyncClient, "__init__", new_sync_init, create=True): + yield diff --git a/packages/toolbox-core/tests/mcp_transport/test_v20241105.py b/packages/toolbox-core/tests/mcp_transport/test_v20241105.py index 1ce2f39bb..eadb0e264 100644 --- a/packages/toolbox-core/tests/mcp_transport/test_v20241105.py +++ b/packages/toolbox-core/tests/mcp_transport/test_v20241105.py @@ -18,6 +18,7 @@ import pytest_asyncio from aiohttp import ClientSession +from toolbox_core.exceptions import ProtocolNegotiationError from toolbox_core.mcp_transport.v20241105 import types from toolbox_core.mcp_transport.v20241105.mcp import McpHttpTransportV20241105 from toolbox_core.protocol import ManifestSchema, Protocol @@ -248,7 +249,7 @@ async def test_initialize_session_protocol_mismatch(self, transport, mocker): ), ) - with pytest.raises(RuntimeError, match="MCP version mismatch"): + with pytest.raises(ProtocolNegotiationError): await transport._initialize_session() async def test_initialize_session_missing_tools_capability(self, transport, mocker): diff --git a/packages/toolbox-core/tests/mcp_transport/test_v20250618.py b/packages/toolbox-core/tests/mcp_transport/test_v20250618.py index fd2e50f7c..b5a1a0a94 100644 --- a/packages/toolbox-core/tests/mcp_transport/test_v20250618.py +++ b/packages/toolbox-core/tests/mcp_transport/test_v20250618.py @@ -18,6 +18,7 @@ import pytest_asyncio from aiohttp import ClientSession +from toolbox_core.exceptions import ProtocolNegotiationError from toolbox_core.mcp_transport.v20250618 import types from toolbox_core.mcp_transport.v20250618.mcp import McpHttpTransportV20250618 from toolbox_core.protocol import ManifestSchema, Protocol, TelemetryAttributes @@ -256,7 +257,7 @@ async def test_initialize_session_protocol_mismatch(self, transport, mocker): ), ) - with pytest.raises(RuntimeError, match="MCP version mismatch"): + with pytest.raises(ProtocolNegotiationError): await transport._initialize_session() async def test_initialize_session_missing_tools_capability(self, transport, mocker): diff --git a/packages/toolbox-core/tests/mcp_transport/test_v20251125.py b/packages/toolbox-core/tests/mcp_transport/test_v20251125.py index 9942041aa..f575b81bc 100644 --- a/packages/toolbox-core/tests/mcp_transport/test_v20251125.py +++ b/packages/toolbox-core/tests/mcp_transport/test_v20251125.py @@ -18,6 +18,7 @@ import pytest_asyncio from aiohttp import ClientSession +from toolbox_core.exceptions import ProtocolNegotiationError from toolbox_core.mcp_transport.v20251125 import types from toolbox_core.mcp_transport.v20251125.mcp import McpHttpTransportV20251125 from toolbox_core.protocol import ManifestSchema, Protocol @@ -256,7 +257,7 @@ async def test_initialize_session_protocol_mismatch(self, transport, mocker): ), ) - with pytest.raises(RuntimeError, match="MCP version mismatch"): + with pytest.raises(ProtocolNegotiationError): await transport._initialize_session() async def test_initialize_session_missing_tools_capability(self, transport, mocker): diff --git a/packages/toolbox-core/tests/mcp_transport/test_v20260618.py b/packages/toolbox-core/tests/mcp_transport/test_v20260618.py index 40d1842e7..5c751320f 100644 --- a/packages/toolbox-core/tests/mcp_transport/test_v20260618.py +++ b/packages/toolbox-core/tests/mcp_transport/test_v20260618.py @@ -70,7 +70,7 @@ async def transport(request, mocker): transport = McpHttpTransportV20260618( "http://fake-server.com", session=mock_session, - protocol=Protocol.MCP_LATEST, + protocol=Protocol.MCP_DRAFT, telemetry_enabled=request.param, ) yield transport diff --git a/packages/toolbox-core/tests/test_client.py b/packages/toolbox-core/tests/test_client.py index 20690eb46..e28bb03a1 100644 --- a/packages/toolbox-core/tests/test_client.py +++ b/packages/toolbox-core/tests/test_client.py @@ -276,7 +276,7 @@ async def test_load_tool_protocol_fallback_success(test_tool_str): mock_2025.tool_invoke.return_value = "ok_from_fallback" mock_2025_cls.return_value = mock_2025 - async with ToolboxClient(TEST_BASE_URL, protocol=Protocol.MCP_LATEST) as client: + async with ToolboxClient(TEST_BASE_URL, protocol=Protocol.MCP_DRAFT) as client: # This should trigger the fallback loaded_tool = await client.load_tool(TOOL_NAME) @@ -320,7 +320,7 @@ async def test_load_tool_protocol_fallback_infinite_loop_prevention(test_tool_st mock_2025.tool_get.side_effect = ProtocolNegotiationError("2024-11-05") mock_2025_cls.return_value = mock_2025 - async with ToolboxClient(TEST_BASE_URL, protocol=Protocol.MCP_LATEST) as client: + async with ToolboxClient(TEST_BASE_URL, protocol=Protocol.MCP_DRAFT) as client: with pytest.raises( ProtocolNegotiationError, match="Server requires protocol fallback to 2024-11-05", @@ -810,7 +810,7 @@ async def test_client_init_with_client_info(): def test_toolbox_client_no_warning_on_mcp(): """Test that initializing ToolboxClient with Protocol.MCP issues NO DeprecationWarning.""" # Mock the transport to avoid actual connection attempts or MCP version warnings - with patch("toolbox_core.client.McpHttpTransportV20250618") as mock_transport: + with patch("toolbox_core.client.McpHttpTransportV20251125") as mock_transport: with warnings.catch_warnings(record=True) as w: warnings.simplefilter("always") diff --git a/packages/toolbox-core/tests/test_e2e_mcp.py b/packages/toolbox-core/tests/test_e2e_mcp.py index c35695ffa..7ffb3ddfa 100644 --- a/packages/toolbox-core/tests/test_e2e_mcp.py +++ b/packages/toolbox-core/tests/test_e2e_mcp.py @@ -24,11 +24,9 @@ from toolbox_core.tool import ToolboxTool -# TODO: Include draft versions in E2E integration tests once the server -# supports SEP-2575 (stateless MCP / Request-Metadata). @pytest_asyncio.fixture( scope="function", - params=[v for v in Protocol.get_supported_mcp_versions() if "DRAFT" not in v], + params=[v for v in Protocol.get_supported_mcp_versions()], ) async def toolbox(request): """Creates a ToolboxClient instance shared by all tests in this module.""" @@ -100,20 +98,27 @@ async def test_run_tool_missing_params(self, get_n_rows_tool: ToolboxTool): with pytest.raises(TypeError, match="missing a required argument: 'num_rows'"): await get_n_rows_tool() - async def test_protocol_fallback_e2e(self): - """Tests that a client using MCP_LATEST can fallback to an older protocol against a server that doesn't support the latest version.""" - # The E2E server currently does not support DRAFT 2026, so this will trigger a fallback. + 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. async with ToolboxClient( - "http://localhost:5000", protocol=Protocol.MCP_LATEST + 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 # Verify that fallback occurred by checking the transport's final protocol version - assert ( - client._ToolboxClient__transport._protocol_version - != Protocol.MCP_LATEST.value - ) + if "5001" in toolbox_server_url: + assert ( + client._ToolboxClient__transport._protocol_version + == Protocol.MCP_DRAFT.value + ) + else: + assert ( + client._ToolboxClient__transport._protocol_version + != Protocol.MCP_DRAFT.value + ) async def test_run_tool_wrong_param_type(self, get_n_rows_tool: ToolboxTool): """Invoke a tool with wrong param type.""" @@ -463,3 +468,25 @@ async def test_run_tool_with_wrong_map_value_type(self, toolbox: ToolboxClient): execution_context={"env": "staging"}, user_scores={"user4": "not-an-integer"}, ) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("toolbox_server") +async def test_mcp_default_protocol(): + """Verify that omitting the protocol argument defaults correctly and works.""" + async with ToolboxClient("http://localhost:5000") 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_draft_fallback(): + """Verify that explicitly using MCP_DRAFT against a server that doesn't support it falls back successfully.""" + async with ToolboxClient( + "http://localhost:5000", protocol=Protocol.MCP_DRAFT + ) as client: + tool = await client.load_tool("get-n-rows") + response = await tool(num_rows="1") + assert "row1" in response diff --git a/packages/toolbox-langchain/integration.cloudbuild.yaml b/packages/toolbox-langchain/integration.cloudbuild.yaml index 140803960..ccadbb557 100644 --- a/packages/toolbox-langchain/integration.cloudbuild.yaml +++ b/packages/toolbox-langchain/integration.cloudbuild.yaml @@ -49,5 +49,5 @@ options: logging: CLOUD_LOGGING_ONLY substitutions: _VERSION: '3.13' - _TOOLBOX_VERSION: '1.4.0' + _TOOLBOX_VERSION: 'v1.6.0' _TOOLBOX_MANIFEST_VERSION: '34' diff --git a/packages/toolbox-langchain/pyproject.toml b/packages/toolbox-langchain/pyproject.toml index 247fdbfab..d8778d62d 100644 --- a/packages/toolbox-langchain/pyproject.toml +++ b/packages/toolbox-langchain/pyproject.toml @@ -50,6 +50,7 @@ test = [ "pytest-asyncio==1.4.0", "pytest==9.0.3", "pytest-cov==7.1.0", + "numpy<2.2.0", "Pillow==12.2.0; python_version >= '3.10'", "google-cloud-secret-manager==2.28.0", "google-cloud-storage==3.10.1", diff --git a/packages/toolbox-langchain/tests/conftest.py b/packages/toolbox-langchain/tests/conftest.py index a1c87a7d8..7fe322b6d 100644 --- a/packages/toolbox-langchain/tests/conftest.py +++ b/packages/toolbox-langchain/tests/conftest.py @@ -25,7 +25,7 @@ from typing import Generator import google -import pytest_asyncio +import pytest from google.auth import compute_engine from google.cloud import secretmanager, storage @@ -75,7 +75,8 @@ def get_toolbox_binary_url(toolbox_version: str) -> str: arch = ( "arm64" if os_system == "darwin" and platform.machine() == "arm64" else "amd64" ) - return f"v{toolbox_version}/{os_system}/{arch}/toolbox" + ext = ".exe" if os_system == "windows" else "" + return f"{toolbox_version}/{os_system}/{arch}/toolbox{ext}" def get_auth_token(client_id: str) -> str: @@ -92,17 +93,17 @@ def get_auth_token(client_id: str) -> str: #### Define Fixtures -@pytest_asyncio.fixture(scope="session") +@pytest.fixture(scope="session") def project_id() -> str: return get_env_var("GOOGLE_CLOUD_PROJECT") -@pytest_asyncio.fixture(scope="session") +@pytest.fixture(scope="session") def toolbox_version() -> str: return get_env_var("TOOLBOX_VERSION") -@pytest_asyncio.fixture(scope="session") +@pytest.fixture(scope="session") def tools_file_path(project_id: str) -> Generator[str]: """Provides a temporary file path containing the tools manifest.""" tools_manifest = access_secret_version( @@ -115,7 +116,7 @@ def tools_file_path(project_id: str) -> Generator[str]: os.remove(tools_file_path) -@pytest_asyncio.fixture(scope="session") +@pytest.fixture(scope="session") def auth_token1(project_id: str) -> str: client_id = access_secret_version( project_id=project_id, secret_id="sdk_testing_client1" @@ -123,7 +124,7 @@ def auth_token1(project_id: str) -> str: return get_auth_token(client_id) -@pytest_asyncio.fixture(scope="session") +@pytest.fixture(scope="session") def auth_token2(project_id: str) -> str: client_id = access_secret_version( project_id=project_id, secret_id="sdk_testing_client2" @@ -131,32 +132,47 @@ def auth_token2(project_id: str) -> str: return get_auth_token(client_id) -@pytest_asyncio.fixture(scope="session") +@pytest.fixture(scope="session") def toolbox_server(toolbox_version: str, tools_file_path: str) -> Generator[None]: """Starts the toolbox server as a subprocess.""" print("Downloading toolbox binary from gcs bucket...") source_blob_name = get_toolbox_binary_url(toolbox_version) - download_blob("mcp-toolbox-for-databases", source_blob_name, "toolbox") + bucket_name = ( + "mcp-toolbox-for-databases-dev" + if toolbox_version in ("main", "mcp-v202606") + else "mcp-toolbox-for-databases" + ) + download_blob(bucket_name, source_blob_name, "toolbox") print("Toolbox binary downloaded successfully.") try: print("Opening toolbox server process...") # Make toolbox executable os.chmod("toolbox", 0o700) # Run toolbox binary - toolbox_server = subprocess.Popen( - ["./toolbox", "--tools-file", tools_file_path] + toolbox_server_1 = subprocess.Popen( + ["./toolbox", "--port", "5000", "--tools-file", tools_file_path] + ) + toolbox_server_2 = subprocess.Popen( + [ + "./toolbox", + "--port", + "5001", + "--tools-file", + tools_file_path, + "--enable-draft-specs", + ] ) # Wait for server to start # Retry logic with a timeout for _ in range(5): # retries - time.sleep(2) - print("Checking if toolbox is successfully started...") - if toolbox_server.poll() is None: - print("Toolbox server started successfully.") + time.sleep(4) + print("Checking if both toolbox servers are successfully started...") + if toolbox_server_1.poll() is None and toolbox_server_2.poll() is None: + print("Toolbox servers started successfully.") break else: - raise RuntimeError("Toolbox server failed to start after 5 retries.") + raise RuntimeError("Toolbox servers failed to start after 5 retries.") except subprocess.CalledProcessError as e: print(e.stderr.decode("utf-8")) print(e.stdout.decode("utf-8")) @@ -164,5 +180,31 @@ def toolbox_server(toolbox_version: str, tools_file_path: str) -> Generator[None yield # Clean up toolbox server - toolbox_server.terminate() - toolbox_server.wait() + toolbox_server_1.terminate() + toolbox_server_2.terminate() + toolbox_server_1.wait() + toolbox_server_2.wait() + + +@pytest.fixture( + params=["http://localhost:5000", "http://localhost:5001"], scope="session" +) +def toolbox_server_url(request) -> str: + return request.param + + +@pytest.fixture(autouse=True) +def patch_toolbox_client_url(toolbox_server_url): + from toolbox_core.client import ToolboxClient + + original_init = ToolboxClient.__init__ + + def new_init(self, url="http://localhost:5000", *args, **kwargs): + if url == "http://localhost:5000": + url = toolbox_server_url + original_init(self, url, *args, **kwargs) + + from unittest.mock import patch + + with patch.object(ToolboxClient, "__init__", new_init, create=True): + yield diff --git a/packages/toolbox-llamaindex/integration.cloudbuild.yaml b/packages/toolbox-llamaindex/integration.cloudbuild.yaml index 028328fb7..2054900fa 100644 --- a/packages/toolbox-llamaindex/integration.cloudbuild.yaml +++ b/packages/toolbox-llamaindex/integration.cloudbuild.yaml @@ -49,5 +49,5 @@ options: logging: CLOUD_LOGGING_ONLY substitutions: _VERSION: '3.13' - _TOOLBOX_VERSION: '1.4.0' + _TOOLBOX_VERSION: 'v1.6.0' _TOOLBOX_MANIFEST_VERSION: '34' diff --git a/packages/toolbox-llamaindex/pyproject.toml b/packages/toolbox-llamaindex/pyproject.toml index ec8107ead..08b64b3d8 100644 --- a/packages/toolbox-llamaindex/pyproject.toml +++ b/packages/toolbox-llamaindex/pyproject.toml @@ -50,6 +50,7 @@ test = [ "pytest-asyncio==1.4.0", "pytest==9.0.3", "pytest-cov==7.1.0", + "numpy<2.2.0", "Pillow==12.2.0; python_version >= '3.10'", "google-cloud-secret-manager==2.28.0", "google-cloud-storage==3.10.1", diff --git a/packages/toolbox-llamaindex/tests/conftest.py b/packages/toolbox-llamaindex/tests/conftest.py index c8a8fa5f9..e5bd6f771 100644 --- a/packages/toolbox-llamaindex/tests/conftest.py +++ b/packages/toolbox-llamaindex/tests/conftest.py @@ -25,7 +25,7 @@ from typing import Generator import google -import pytest_asyncio +import pytest from google.auth import compute_engine from google.cloud import secretmanager, storage @@ -75,7 +75,8 @@ def get_toolbox_binary_url(toolbox_version: str) -> str: arch = ( "arm64" if os_system == "darwin" and platform.machine() == "arm64" else "amd64" ) - return f"v{toolbox_version}/{os_system}/{arch}/toolbox" + ext = ".exe" if os_system == "windows" else "" + return f"{toolbox_version}/{os_system}/{arch}/toolbox{ext}" def get_auth_token(client_id: str) -> str: @@ -92,17 +93,17 @@ def get_auth_token(client_id: str) -> str: #### Define Fixtures -@pytest_asyncio.fixture(scope="session") +@pytest.fixture(scope="session") def project_id() -> str: return get_env_var("GOOGLE_CLOUD_PROJECT") -@pytest_asyncio.fixture(scope="session") +@pytest.fixture(scope="session") def toolbox_version() -> str: return get_env_var("TOOLBOX_VERSION") -@pytest_asyncio.fixture(scope="session") +@pytest.fixture(scope="session") def tools_file_path(project_id: str) -> Generator[str]: """Provides a temporary file path containing the tools manifest.""" tools_manifest = access_secret_version( @@ -115,7 +116,7 @@ def tools_file_path(project_id: str) -> Generator[str]: os.remove(tools_file_path) -@pytest_asyncio.fixture(scope="session") +@pytest.fixture(scope="session") def auth_token1(project_id: str) -> str: client_id = access_secret_version( project_id=project_id, secret_id="sdk_testing_client1" @@ -123,7 +124,7 @@ def auth_token1(project_id: str) -> str: return get_auth_token(client_id) -@pytest_asyncio.fixture(scope="session") +@pytest.fixture(scope="session") def auth_token2(project_id: str) -> str: client_id = access_secret_version( project_id=project_id, secret_id="sdk_testing_client2" @@ -131,32 +132,47 @@ def auth_token2(project_id: str) -> str: return get_auth_token(client_id) -@pytest_asyncio.fixture(scope="session") +@pytest.fixture(scope="session") def toolbox_server(toolbox_version: str, tools_file_path: str) -> Generator[None]: """Starts the toolbox server as a subprocess.""" print("Downloading toolbox binary from gcs bucket...") source_blob_name = get_toolbox_binary_url(toolbox_version) - download_blob("mcp-toolbox-for-databases", source_blob_name, "toolbox") + bucket_name = ( + "mcp-toolbox-for-databases-dev" + if toolbox_version in ("main", "mcp-v202606") + else "mcp-toolbox-for-databases" + ) + download_blob(bucket_name, source_blob_name, "toolbox") print("Toolbox binary downloaded successfully.") try: print("Opening toolbox server process...") # Make toolbox executable os.chmod("toolbox", 0o700) # Run toolbox binary - toolbox_server = subprocess.Popen( - ["./toolbox", "--tools-file", tools_file_path] + toolbox_server_1 = subprocess.Popen( + ["./toolbox", "--port", "5000", "--tools-file", tools_file_path] + ) + toolbox_server_2 = subprocess.Popen( + [ + "./toolbox", + "--port", + "5001", + "--tools-file", + tools_file_path, + "--enable-draft-specs", + ] ) # Wait for server to start # Retry logic with a timeout for _ in range(5): # retries time.sleep(4) - print("Checking if toolbox is successfully started...") - if toolbox_server.poll() is None: - print("Toolbox server started successfully.") + print("Checking if both toolbox servers are successfully started...") + if toolbox_server_1.poll() is None and toolbox_server_2.poll() is None: + print("Toolbox servers started successfully.") break else: - raise RuntimeError("Toolbox server failed to start after 5 retries.") + raise RuntimeError("Toolbox servers failed to start after 5 retries.") except subprocess.CalledProcessError as e: print(e.stderr.decode("utf-8")) print(e.stdout.decode("utf-8")) @@ -164,5 +180,31 @@ def toolbox_server(toolbox_version: str, tools_file_path: str) -> Generator[None yield # Clean up toolbox server - toolbox_server.terminate() - toolbox_server.wait() + toolbox_server_1.terminate() + toolbox_server_2.terminate() + toolbox_server_1.wait() + toolbox_server_2.wait() + + +@pytest.fixture( + params=["http://localhost:5000", "http://localhost:5001"], scope="session" +) +def toolbox_server_url(request) -> str: + return request.param + + +@pytest.fixture(autouse=True) +def patch_toolbox_client_url(toolbox_server_url): + from toolbox_core.client import ToolboxClient + + original_init = ToolboxClient.__init__ + + def new_init(self, url="http://localhost:5000", *args, **kwargs): + if url == "http://localhost:5000": + url = toolbox_server_url + original_init(self, url, *args, **kwargs) + + from unittest.mock import patch + + with patch.object(ToolboxClient, "__init__", new_init, create=True): + yield