move warning to diff file

Author: twishabansal

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__

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,

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]}"

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