deprecate instead of remove

Author: twishabansal

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

@@ -17,6 +17,7 @@
 from typing import Any, Awaitable, Callable, Mapping, Optional, Union
 
 from aiohttp import ClientSession
+from deprecated import deprecated
 
 from .protocol import ManifestSchema, ToolSchema
 from .tool import ToolboxTool
@@ -339,3 +340,26 @@ async def load_toolset(
             )
 
         return tools
+
+    @deprecated("Please add client level headers during client initialization.")
+    def add_headers(
+        self,
+        headers: Mapping[str, Union[Callable[[], str], Callable[[], Awaitable[str]]]],
+    ) -> None:
+        """
+        Add headers to be included in each request sent through this client.
+        Args:
+            headers: Headers to include in each request sent through this client.
+        Raises:
+            ValueError: If any of the headers are already registered in the client.
+        """
+        existing_headers = self.__client_headers.keys()
+        incoming_headers = headers.keys()
+        duplicates = existing_headers & incoming_headers
+        if duplicates:
+            raise ValueError(
+                f"Client header(s) `{', '.join(duplicates)}` already registered in the client."
+            )
+
+        merged_headers = {**self.__client_headers, **headers}
+        self.__client_headers = merged_headers

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

@@ -17,6 +17,8 @@
 from threading import Thread
 from typing import Any, Awaitable, Callable, Mapping, Optional, Union
 
+from deprecated import deprecated
+
 from .client import ToolboxClient
 from .sync_tool import ToolboxSyncTool
 
@@ -153,6 +155,24 @@ def load_toolset(
             for async_tool in async_tools
         ]
 
+    @deprecated("Please add client level headers during client initialization.")
+    def add_headers(
+        self,
+        headers: Mapping[
+            str, Union[Callable[[], str], Callable[[], Awaitable[str]], str]
+        ],
+    ) -> None:
+        """
+        Add headers to be included in each request sent through this client.
+
+        Args:
+            headers: Headers to include in each request sent through this client.
+
+        Raises:
+            ValueError: If any of the headers are already registered in the client.
+        """
+        self.__async_client.add_headers(headers)
+
     def __enter__(self):
         """Enter the runtime context related to this client instance."""
         return self

packages/toolbox-core/tests/test_client.py

@@ -1417,6 +1417,50 @@ async def test_load_toolset_with_headers(
             assert len(tools) == 1
             assert tools[0].__name__ == tool_name
 
+    @pytest.mark.asyncio
+    async def test_add_headers_success(
+        self, aioresponses, test_tool_str, static_header
+    ):
+        """Tests adding headers after client initialization."""
+        tool_name = "tool_after_add_headers"
+        manifest = ManifestSchema(
+            serverVersion="0.0.0", tools={tool_name: test_tool_str}
+        )
+        expected_payload = {"result": "added_ok"}
+
+        get_callback = self.create_callback_factory(
+            expected_header=static_header,
+            callback_payload=manifest.model_dump(),
+        )
+        aioresponses.get(f"{TEST_BASE_URL}/api/tool/{tool_name}", callback=get_callback)
+
+        post_callback = self.create_callback_factory(
+            expected_header=static_header,
+            callback_payload=expected_payload,
+        )
+        aioresponses.post(
+            f"{TEST_BASE_URL}/api/tool/{tool_name}/invoke", callback=post_callback
+        )
+
+        async with ToolboxClient(TEST_BASE_URL) as client:
+            client.add_headers(static_header)
+            assert client._ToolboxClient__client_headers == static_header
+
+            tool = await client.load_tool(tool_name)
+            result = await tool(param1="test")
+            assert result == expected_payload["result"]
+
+    @pytest.mark.asyncio
+    async def test_add_headers_duplicate_fail(self, static_header):
+        """Tests that adding a duplicate header via add_headers raises
+        ValueError."""
+        async with ToolboxClient(TEST_BASE_URL, client_headers=static_header) as client:
+            with pytest.raises(
+                ValueError,
+                match=f"Client header\\(s\\) `X-Static-Header` already registered",
+            ):
+                await client.add_headers(static_header)
+
     @pytest.mark.asyncio
     async def test_client_header_auth_token_conflict_fail(
         self, aioresponses, test_tool_auth

packages/toolbox-core/tests/test_sync_client.py

@@ -350,6 +350,72 @@ def test_load_tool_raises_if_loop_or_thread_none(self):
         # sync_client_environment will handle the final cleanup of original_class_loop/thread.
 
 
+class TestSyncClientHeaders:
+    """Additive tests for client header functionality specific to ToolboxSyncClient if any,
+    or counterparts to async client header tests."""
+
+    def test_sync_add_headers_success(
+        self, aioresponses, test_tool_str_schema, sync_client
+    ):
+        tool_name = "tool_after_add_headers_sync"
+        manifest = ManifestSchema(
+            serverVersion="0.0.0", tools={tool_name: test_tool_str_schema}
+        )
+        expected_payload = {"result": "added_sync_ok"}
+        headers_to_add = {"X-Custom-SyncHeader": "sync_value"}
+
+        def get_callback(url, **kwargs):
+            # The sync_client might have default headers. Check ours are present.
+            assert kwargs.get("headers") is not None
+            for key, value in headers_to_add.items():
+                assert kwargs["headers"].get(key) == value
+            return CallbackResult(status=200, payload=manifest.model_dump())
+
+        aioresponses.get(f"{TEST_BASE_URL}/api/tool/{tool_name}", callback=get_callback)
+
+        def post_callback(url, **kwargs):
+            assert kwargs.get("headers") is not None
+            for key, value in headers_to_add.items():
+                assert kwargs["headers"].get(key) == value
+            return CallbackResult(status=200, payload=expected_payload)
+
+        aioresponses.post(
+            f"{TEST_BASE_URL}/api/tool/{tool_name}/invoke", callback=post_callback
+        )
+
+        sync_client.add_headers(headers_to_add)
+        tool = sync_client.load_tool(tool_name)
+        result = tool(param1="test")
+        assert result == expected_payload["result"]
+
+    def test_sync_add_headers_duplicate_fail(self):
+        """Tests that adding a duplicate header via add_headers raises ValueError (from async client)."""
+        initial_headers = {"X-Initial-Header": "initial_value"}
+        mock_async_client = AsyncMock(spec=ToolboxClient)
+
+        # Configure add_headers to simulate the ValueError from ToolboxClient
+        def mock_add_headers(headers):
+            # Simulate ToolboxClient's check
+            if "X-Initial-Header" in headers:
+                raise ValueError(
+                    "Client header(s) `X-Initial-Header` already registered"
+                )
+
+        mock_async_client.add_headers = Mock(side_effect=mock_add_headers)
+
+        with patch(
+            "toolbox_core.sync_client.ToolboxClient", return_value=mock_async_client
+        ):
+            with ToolboxSyncClient(
+                TEST_BASE_URL, client_headers=initial_headers
+            ) as client:
+                with pytest.raises(
+                    ValueError,
+                    match="Client header\\(s\\) `X-Initial-Header` already registered",
+                ):
+                    client.add_headers({"X-Initial-Header": "another_value"})
+
+
 class TestSyncAuth:
     @pytest.fixture
     def expected_header_token(self):