Skip to content

Commit 394fa4d

Browse files
authored
fix(toolbox-core)!: Simplify add_headers to make it sync in async client (#231)
* chore: Add unit test cases * chore: Delint * feat: Warn on insecure tool invocation with authentication This change introduces a warning that is displayed immediately before a tool invocation if: 1. The invocation includes an authentication header. 2. The connection is being made over non-secure HTTP. > [!IMPORTANT] The purpose of this warning is to alert the user to the security risk of sending credentials over an unencrypted channel and to encourage the use of HTTPS. * fix!: Warn about https only during tool initialization * chore: Simplify add_headers to make it sync in async client * chore: Fix unit tests
1 parent 9427ea8 commit 394fa4d

File tree

4 files changed

+36
-23
lines changed

4 files changed

+36
-23
lines changed

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -308,11 +308,11 @@ async def load_toolset(
308308

309309
return tools
310310

311-
async def add_headers(
311+
def add_headers(
312312
self, headers: Mapping[str, Union[Callable, Coroutine, str]]
313313
) -> None:
314314
"""
315-
Asynchronously Add headers to be included in each request sent through this client.
315+
Add headers to be included in each request sent through this client.
316316
317317
Args:
318318
headers: Headers to include in each request sent through this client.

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

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -144,7 +144,7 @@ def load_toolset(
144144
if not self.__loop or not self.__thread:
145145
raise ValueError("Background loop or thread cannot be None.")
146146

147-
async_tools = asyncio.run_coroutine_threadsafe(coro, self.__loop).result() # type: ignore
147+
async_tools = asyncio.run_coroutine_threadsafe(coro, self.__loop).result()
148148
return [
149149
ToolboxSyncTool(async_tool, self.__loop, self.__thread)
150150
for async_tool in async_tools
@@ -154,18 +154,15 @@ def add_headers(
154154
self, headers: Mapping[str, Union[Callable, Coroutine, str]]
155155
) -> None:
156156
"""
157-
Synchronously Add headers to be included in each request sent through this client.
157+
Add headers to be included in each request sent through this client.
158158
159159
Args:
160160
headers: Headers to include in each request sent through this client.
161161
162162
Raises:
163163
ValueError: If any of the headers are already registered in the client.
164164
"""
165-
coro = self.__async_client.add_headers(headers)
166-
167-
# We have already created a new loop in the init method in case it does not already exist
168-
asyncio.run_coroutine_threadsafe(coro, self.__loop).result() # type: ignore
165+
self.__async_client.add_headers(headers)
169166

170167
def __enter__(self):
171168
"""Enter the runtime context related to this client instance."""

packages/toolbox-core/tests/test_client.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1443,7 +1443,7 @@ async def test_add_headers_success(
14431443
)
14441444

14451445
async with ToolboxClient(TEST_BASE_URL) as client:
1446-
await client.add_headers(static_header)
1446+
client.add_headers(static_header)
14471447
assert client._ToolboxClient__client_headers == static_header
14481448

14491449
tool = await client.load_tool(tool_name)

packages/toolbox-core/tests/test_sync_client.py

Lines changed: 30 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515

1616
import inspect
1717
from typing import Any, Callable, Mapping, Optional
18-
from unittest.mock import AsyncMock, patch
18+
from unittest.mock import AsyncMock, Mock, patch
1919

2020
import pytest
2121
from aioresponses import CallbackResult, aioresponses
@@ -44,8 +44,12 @@ def sync_client_environment():
4444
# This ensures any client created will start a new loop/thread.
4545

4646
# Ensure no loop/thread is running from a previous misbehaving test or setup
47-
assert original_loop is None or not original_loop.is_running()
48-
assert original_thread is None or not original_thread.is_alive()
47+
if original_loop and original_loop.is_running():
48+
original_loop.call_soon_threadsafe(original_loop.stop)
49+
if original_thread and original_thread.is_alive():
50+
original_thread.join(timeout=5)
51+
ToolboxSyncClient._ToolboxSyncClient__loop = None
52+
ToolboxSyncClient._ToolboxSyncClient__thread = None
4953

5054
ToolboxSyncClient._ToolboxSyncClient__loop = None
5155
ToolboxSyncClient._ToolboxSyncClient__thread = None
@@ -408,20 +412,32 @@ def post_callback(url, **kwargs):
408412
result = tool(param1="test")
409413
assert result == expected_payload["result"]
410414

411-
@pytest.mark.usefixtures("sync_client_environment")
412415
def test_sync_add_headers_duplicate_fail(self):
413-
"""
414-
Tests that adding a duplicate header via add_headers raises ValueError.
415-
Manually create client to control initial headers.
416-
"""
416+
"""Tests that adding a duplicate header via add_headers raises ValueError (from async client)."""
417417
initial_headers = {"X-Initial-Header": "initial_value"}
418+
mock_async_client = AsyncMock(spec=ToolboxClient)
418419

419-
with ToolboxSyncClient(TEST_BASE_URL, client_headers=initial_headers) as client:
420-
with pytest.raises(
421-
ValueError,
422-
match="Client header\\(s\\) `X-Initial-Header` already registered",
423-
):
424-
client.add_headers({"X-Initial-Header": "another_value"})
420+
# Configure add_headers to simulate the ValueError from ToolboxClient
421+
def mock_add_headers(headers):
422+
# Simulate ToolboxClient's check
423+
if "X-Initial-Header" in headers:
424+
raise ValueError(
425+
"Client header(s) `X-Initial-Header` already registered"
426+
)
427+
428+
mock_async_client.add_headers = Mock(side_effect=mock_add_headers)
429+
430+
with patch(
431+
"toolbox_core.sync_client.ToolboxClient", return_value=mock_async_client
432+
):
433+
with ToolboxSyncClient(
434+
TEST_BASE_URL, client_headers=initial_headers
435+
) as client:
436+
with pytest.raises(
437+
ValueError,
438+
match="Client header\\(s\\) `X-Initial-Header` already registered",
439+
):
440+
client.add_headers({"X-Initial-Header": "another_value"})
425441

426442

427443
class TestSyncAuth:

0 commit comments

Comments
 (0)