googleapis · anubhav756 · May 16, 2025 · May 6, 2025 · May 6, 2025 · May 6, 2025
@@ -308,11 +308,11 @@ async def load_toolset(
 
         return tools
 
-    async def add_headers(
+    def add_headers(
         self, headers: Mapping[str, Union[Callable, Coroutine, str]]
     ) -> None:
         """
-        Asynchronously Add headers to be included in each request sent through this client.
+        Add headers to be included in each request sent through this client.
 
         Args:
             headers: Headers to include in each request sent through this client.

@@ -14,13 +14,13 @@
 
 
 import asyncio
+from asyncio import AbstractEventLoop
 from threading import Thread
-from typing import Any, Callable, Coroutine, Mapping, Optional, TypeVar, Union
+from typing import Any, Callable, Coroutine, Mapping, Optional, Union
 
 from .client import ToolboxClient
 from .sync_tool import ToolboxSyncTool
-
-T = TypeVar("T")
+from concurrent.futures import Future
 
 
 class ToolboxSyncClient:
@@ -31,7 +31,7 @@ class ToolboxSyncClient:
     service endpoint.
     """
 
-    __loop: Optional[asyncio.AbstractEventLoop] = None
+    __loop: Optional[AbstractEventLoop] = None
     __thread: Optional[Thread] = None
 
     def __init__(
@@ -58,11 +58,22 @@ def __init__(
         async def create_client():
             return ToolboxClient(url, client_headers=client_headers)
 
-        # Ignoring type since we're already checking the existence of a loop above.
         self.__async_client = asyncio.run_coroutine_threadsafe(
             create_client(), self.__class__.__loop
         ).result()
 
+    @property
+    def _async_client(self) -> ToolboxClient:
+        return self.__async_client
+
+    @property
+    def _loop(self) -> Optional[AbstractEventLoop]:
+        return self.__class__.__loop
+
+    @property
+    def _thread(self) -> Optional[Thread]:
+        return self.__class__.__thread
+
     def close(self):
         """
         Synchronously closes the underlying client session. Doing so will cause
@@ -75,6 +86,46 @@ def close(self):
         coro = self.__async_client.close()
         asyncio.run_coroutine_threadsafe(coro, self.__loop).result()
 
+    def load_tool_future(
+        self,
+        name: str,
+        auth_token_getters: dict[str, Callable[[], str]] = {},
+        bound_params: Mapping[str, Union[Callable[[], Any], Any]] = {},
+    ) -> Future[ToolboxSyncTool]:
+        """
+        Returns a future that loads a tool from the server.
+        """
+        if not self.__loop or not self.__thread:
+            raise ValueError("Background loop or thread cannot be None.")
+
+        async def async_worker() -> ToolboxSyncTool:
+            async_tool = await self.__async_client.load_tool(name, auth_token_getters, bound_params)
+            return ToolboxSyncTool(async_tool, self.__loop, self.__thread)
+        return asyncio.run_coroutine_threadsafe(async_worker(), self.__loop)
+
+    def load_toolset_future(
+        self,
+        name: Optional[str] = None,
+        auth_token_getters: dict[str, Callable[[], str]] = {},
+        bound_params: Mapping[str, Union[Callable[[], Any], Any]] = {},
+        strict: bool = False,
+    ) -> Future[list[ToolboxSyncTool]]:
+        """
+        Returns a future that fetches a toolset and loads all tools defined within it.
+        """
+        if not self.__loop or not self.__thread:
+            raise ValueError("Background loop or thread cannot be None.")
+
+        async def async_worker() -> list[ToolboxSyncTool]:
+            async_tools = await self.__async_client.load_toolset(
+                name, auth_token_getters, bound_params, strict
+            )
+            return [
+                ToolboxSyncTool(async_tool, self.__loop, self.__thread)
+                for async_tool in async_tools
+            ]
+        return asyncio.run_coroutine_threadsafe(async_worker(), self.__loop)
+
     def load_tool(
         self,
         name: str,
@@ -100,61 +151,52 @@ def load_tool(
                 for execution. The specific arguments and behavior of the callable
                 depend on the tool itself.
         """
-        coro = self.__async_client.load_tool(name, auth_token_getters, bound_params)
-
-        if not self.__loop or not self.__thread:
-            raise ValueError("Background loop or thread cannot be None.")
-
-        async_tool = asyncio.run_coroutine_threadsafe(coro, self.__loop).result()
-        return ToolboxSyncTool(async_tool, self.__loop, self.__thread)
+        return self.load_tool_future(name, auth_token_getters, bound_params).result()
 
     def load_toolset(
         self,
-        name: str,
+        name: Optional[str] = None,
         auth_token_getters: dict[str, Callable[[], str]] = {},
         bound_params: Mapping[str, Union[Callable[[], Any], Any]] = {},
+        strict: bool = False,
     ) -> list[ToolboxSyncTool]:
         """
         Synchronously fetches a toolset and loads all tools defined within it.
 
         Args:
-            name: Name of the toolset to load tools.
+            name: Name of the toolset to load. If None, loads the default toolset.
             auth_token_getters: A mapping of authentication service names to
                 callables that return the corresponding authentication token.
             bound_params: A mapping of parameter names to bind to specific values or
                 callables that are called to produce values as needed.
+            strict: If True, raises an error if *any* loaded tool instance fails
+                to utilize at least one provided parameter or auth token (if any
+                provided). If False (default), raises an error only if a
+                user-provided parameter or auth token cannot be applied to *any*
+                loaded tool across the set.
 
         Returns:
             list[ToolboxSyncTool]: A list of callables, one for each tool defined
             in the toolset.
-        """
-        coro = self.__async_client.load_toolset(name, auth_token_getters, bound_params)
-
-        if not self.__loop or not self.__thread:
-            raise ValueError("Background loop or thread cannot be None.")
 
-        async_tools = asyncio.run_coroutine_threadsafe(coro, self.__loop).result()  # type: ignore
-        return [
-            ToolboxSyncTool(async_tool, self.__loop, self.__thread)
-            for async_tool in async_tools
-        ]
+        Raises:
+            ValueError: If validation fails based on the `strict` flag.
+        """
+        return self.load_toolset_future(name, auth_token_getters, bound_params, strict).result()
 
     def add_headers(
         self, headers: Mapping[str, Union[Callable, Coroutine, str]]
     ) -> None:
         """
-        Synchronously Add headers to be included in each request sent through this client.
+        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.
         """
-        coro = self.__async_client.add_headers(headers)
-
-        # We have already created a new loop in the init method in case it does not already exist
-        asyncio.run_coroutine_threadsafe(coro, self.__loop).result()  # type: ignore
+        self.__async_client.add_headers(headers)
 
     def __enter__(self):
         """Enter the runtime context related to this client instance."""

@@ -17,12 +17,11 @@
 from asyncio import AbstractEventLoop
 from inspect import Signature
 from threading import Thread
-from typing import Any, Callable, Coroutine, Mapping, Sequence, TypeVar, Union
+from typing import Any, Callable, Coroutine, Mapping, Sequence, Union
 
 from .protocol import ParameterSchema
 from .tool import ToolboxTool
-
-T = TypeVar("T")
+from concurrent.futures import Future
 
 
 class ToolboxSyncTool:
@@ -69,6 +68,18 @@ def __init__(
             f"{self.__class__.__qualname__}.{self.__async_tool.__name__}"
         )
 
+    @property
+    def _async_tool(self) -> ToolboxTool:
+        return self.__async_tool
+
+    @property
+    def _loop(self) -> AbstractEventLoop:
+        return self.__loop
+
+    @property
+    def _thread(self) -> Thread:
+        return self.__thread
+
     @property
     def __name__(self) -> str:
         return self.__async_tool.__name__
@@ -119,6 +130,13 @@ def _auth_service_token_getters(self) -> Mapping[str, Callable[[], str]]:
     def _client_headers(self) -> Mapping[str, Union[Callable, Coroutine, str]]:
         return self.__async_tool._client_headers
 
+    def call_future(self, *args: Any, **kwargs: Any) -> Future[str]:
+        """
+        Returns future that calls the remote tool with the provided arguments.
+        """
+        coro = self.__async_tool(*args, **kwargs)
+        return asyncio.run_coroutine_threadsafe(coro, self.__loop)
+
     def __call__(self, *args: Any, **kwargs: Any) -> str:
         """
         Synchronously calls the remote tool with the provided arguments.
@@ -133,8 +151,7 @@ def __call__(self, *args: Any, **kwargs: Any) -> str:
         Returns:
             The string result returned by the remote tool execution.
         """
-        coro = self.__async_tool(*args, **kwargs)
-        return asyncio.run_coroutine_threadsafe(coro, self.__loop).result()
+        return self.call_future(*args, **kwargs).result()
 
     def add_auth_token_getters(
         self,

@@ -19,6 +19,7 @@
 from warnings import warn
 
 from aiohttp import ClientSession
+from pydantic import BaseModel
 
 from .protocol import ParameterSchema
 from .utils import (
@@ -158,6 +159,10 @@ def _auth_service_token_getters(self) -> Mapping[str, Callable[[], str]]:
     def _client_headers(self) -> Mapping[str, Union[Callable, Coroutine, str]]:
         return MappingProxyType(self.__client_headers)
 
+    @property
+    def _pydantic_model(self) -> type[BaseModel]:
+        return self.__pydantic_model
+
     def __copy(
         self,
         session: Optional[ClientSession] = None,
@@ -239,7 +244,7 @@ async def __call__(self, *args: Any, **kwargs: Any) -> str:
             for s in self.__required_authn_params.values():
                 req_auth_services.update(s)
             req_auth_services.update(self.__required_authz_tokens)
-            raise ValueError(
+            raise PermissionError(
                 f"One or more of the following authn services are required to invoke this tool"
                 f": {','.join(req_auth_services)}"
             )

@@ -1443,7 +1443,7 @@ async def test_add_headers_success(
         )
 
         async with ToolboxClient(TEST_BASE_URL) as client:
-            await client.add_headers(static_header)
+            client.add_headers(static_header)
             assert client._ToolboxClient__client_headers == static_header
 
             tool = await client.load_tool(tool_name)

@@ -147,7 +147,7 @@ async def test_run_tool_no_auth(self, toolbox: ToolboxClient):
         """Tests running a tool requiring auth without providing auth."""
         tool = await toolbox.load_tool("get-row-by-id-auth")
         with pytest.raises(
-            Exception,
+            PermissionError,
             match="One or more of the following authn services are required to invoke this tool: my-test-auth",
         ):
             await tool(id="2")
@@ -188,7 +188,7 @@ async def test_run_tool_param_auth_no_auth(self, toolbox: ToolboxClient):
         """Tests running a tool with a param requiring auth, without auth."""
         tool = await toolbox.load_tool("get-row-by-email-auth")
         with pytest.raises(
-            ValueError,
+            PermissionError,
             match="One or more of the following authn services are required to invoke this tool: my-test-auth",
         ):
             await tool()