fix: Correctly determine remaining required authz tokens

anubhav756 · anubhav756 · commit c9f8d205a87f · 2025-05-13T01:30:04.000+05:30
This PR fixes an issue where the system could inaccurately identify the authorization tokens still needed for tool invocation.

The `identify_required_authn_params` helper has been updated to leverage its new capability of recognizing all alternatives of required authorization tokens.

A new `ToolboxTool` member variable, `__required_authz_tokens`, now stores these alternatives. The tool invocation logic correctly uses this to check if any matching token has been provided.

This new member variable is also updated correctly by the remaining authz tokens while adding auth token getters, and validated right before tool invocation.
diff --git a/packages/toolbox-core/src/toolbox_core/client.py b/packages/toolbox-core/src/toolbox_core/client.py
@@ -79,10 +79,9 @@ def __parse_tool(
             else:  # regular parameter
                 params.append(p)
 
-        authn_params, _, used_auth_keys = identify_required_authn_params(
-            # TODO: Add schema.authRequired as second arg
+        authn_params, authz_tokens, used_auth_keys = identify_required_authn_params(
             authn_params,
-            [],
+            schema.authRequired,
             auth_token_getters.keys(),
         )
 
@@ -94,6 +93,7 @@ def __parse_tool(
             # create a read-only values to prevent mutation
             params=tuple(params),
             required_authn_params=types.MappingProxyType(authn_params),
+            required_authz_tokens=authz_tokens,
             auth_service_token_getters=types.MappingProxyType(auth_token_getters),
             bound_params=types.MappingProxyType(bound_params),
             client_headers=types.MappingProxyType(client_headers),
diff --git a/packages/toolbox-core/src/toolbox_core/tool.py b/packages/toolbox-core/src/toolbox_core/tool.py
@@ -49,6 +49,7 @@ def __init__(
         description: str,
         params: Sequence[ParameterSchema],
         required_authn_params: Mapping[str, list[str]],
+        required_authz_tokens: Sequence[str],
         auth_service_token_getters: Mapping[str, Callable[[], str]],
         bound_params: Mapping[str, Union[Callable[[], Any], Any]],
         client_headers: Mapping[str, Union[Callable, Coroutine, str]],
@@ -63,12 +64,14 @@ def __init__(
             name: The name of the remote tool.
             description: The description of the remote tool.
             params: The args of the tool.
-            required_authn_params: A map of required authenticated parameters to a list
-                of alternative services that can provide values for them.
-            auth_service_token_getters: A dict of authService -> token (or callables that
-                produce a token)
-            bound_params: A mapping of parameter names to bind to specific values or
-                callables that are called to produce values as needed.
+            required_authn_params: A map of required authenticated parameters to
+                a list of alternative services that can provide values for them.
+            required_authz_tokens: A sequence of alternative services for
+                providing authorization token for the tool invocation.
+            auth_service_token_getters: A dict of authService -> token (or
+                callables that produce a token)
+            bound_params: A mapping of parameter names to bind to specific
+                values or callables that are called to produce values as needed.
             client_headers: Client specific headers bound to the tool.
         """
         # used to invoke the toolbox API
@@ -106,6 +109,8 @@ def __init__(
 
         # map of parameter name to auth service required by it
         self.__required_authn_params = required_authn_params
+        # sequence of authorization tokens required by it
+        self.__required_authz_tokens = required_authz_tokens
         # map of authService -> token_getter
         self.__auth_service_token_getters = auth_service_token_getters
         # map of parameter name to value (or callable that produces that value)
@@ -149,6 +154,7 @@ def __copy(
         description: Optional[str] = None,
         params: Optional[Sequence[ParameterSchema]] = None,
         required_authn_params: Optional[Mapping[str, list[str]]] = None,
+        required_authz_tokens: Optional[Sequence[str]] = None,
         auth_service_token_getters: Optional[Mapping[str, Callable[[], str]]] = None,
         bound_params: Optional[Mapping[str, Union[Callable[[], Any], Any]]] = None,
         client_headers: Optional[Mapping[str, Union[Callable, Coroutine, str]]] = None,
@@ -162,12 +168,14 @@ def __copy(
             name: The name of the remote tool.
             description: The description of the remote tool.
             params: The args of the tool.
-            required_authn_params: A map of required authenticated parameters to a list
-                of alternative services that can provide values for them.
-            auth_service_token_getters: A dict of authService -> token (or callables
-                that produce a token)
-            bound_params: A mapping of parameter names to bind to specific values or
-                callables that are called to produce values as needed.
+            required_authn_params: A map of required authenticated parameters to
+                a list of alternative services that can provide values for them.
+            required_authz_tokens: A sequence of alternative services for
+                providing authorization token for the tool invocation.
+            auth_service_token_getters: A dict of authService -> token (or
+                callables that produce a token)
+            bound_params: A mapping of parameter names to bind to specific
+                values or callables that are called to produce values as needed.
             client_headers: Client specific headers bound to the tool.
         """
         check = lambda val, default: val if val is not None else default
@@ -180,6 +188,9 @@ def __copy(
             required_authn_params=check(
                 required_authn_params, self.__required_authn_params
             ),
+            required_authz_tokens=check(
+                required_authz_tokens, self.__required_authz_tokens
+            ),
             auth_service_token_getters=check(
                 auth_service_token_getters, self.__auth_service_token_getters
             ),
@@ -207,11 +218,15 @@ async def __call__(self, *args: Any, **kwargs: Any) -> str:
         """
 
         # check if any auth services need to be specified yet
-        if len(self.__required_authn_params) > 0:
+        if (
+            len(self.__required_authn_params) > 0
+            or len(self.__required_authz_tokens) > 0
+        ):
             # Gather all the required auth services into a set
             req_auth_services = set()
             for s in self.__required_authn_params.values():
                 req_auth_services.update(s)
+            req_auth_services.update(self.__required_authz_tokens)
             raise ValueError(
                 f"One or more of the following authn services are required to invoke this tool"
                 f": {','.join(req_auth_services)}"
@@ -297,18 +312,20 @@ def add_auth_token_getters(
             dict(self.__auth_service_token_getters, **auth_token_getters)
         )
         # create a read-only updated for params that are still required
-        new_req_authn_params = MappingProxyType(
+        new_req_authn_params, new_req_authz_tokens, used_auth_token_getters = (
             identify_required_authn_params(
-                # TODO: Add authRequired
                 self.__required_authn_params,
-                [],
+                self.__required_authz_tokens,
                 auth_token_getters.keys(),
-            )[0]
+            )
         )
 
+        # TODO: Add validation for used_auth_token_getters
+
         return self.__copy(
             auth_service_token_getters=new_getters,
-            required_authn_params=new_req_authn_params,
+            required_authn_params=MappingProxyType(new_req_authn_params),
+            required_authz_tokens=new_req_authz_tokens,
         )
 
     def bind_params(
diff --git a/packages/toolbox-core/src/toolbox_core/utils.py b/packages/toolbox-core/src/toolbox_core/utils.py
@@ -46,7 +46,7 @@ def create_func_docstring(description: str, params: Sequence[ParameterSchema]) -
 
 def identify_required_authn_params(
     req_authn_params: Mapping[str, list[str]],
-    req_authz_tokens: list[str],
+    req_authz_tokens: Sequence[str],
     auth_service_names: Iterable[str],
 ) -> tuple[dict[str, list[str]], list[str], set[str]]:
     """
@@ -100,7 +100,7 @@ def identify_required_authn_params(
     if matched_authz_services:
         used_services.update(matched_authz_services)
     else:
-        required_authz_tokens = req_authz_tokens
+        required_authz_tokens = list(req_authz_tokens)
 
     return required_authn_params, required_authz_tokens, used_services
 
diff --git a/packages/toolbox-core/tests/test_tool.py b/packages/toolbox-core/tests/test_tool.py
@@ -205,6 +205,7 @@ async def test_tool_creation_callable_and_run(
             description=sample_tool_description,
             params=sample_tool_params,
             required_authn_params={},
+            required_authz_tokens=[],
             auth_service_token_getters={},
             bound_params={},
             client_headers={},
@@ -250,6 +251,7 @@ async def test_tool_run_with_pydantic_validation_error(
             description=sample_tool_description,
             params=sample_tool_params,
             required_authn_params={},
+            required_authz_tokens=[],
             auth_service_token_getters={},
             bound_params={},
             client_headers={},
@@ -337,6 +339,7 @@ def test_tool_init_basic(http_session, sample_tool_params, sample_tool_descripti
         description=sample_tool_description,
         params=sample_tool_params,
         required_authn_params={},
+        required_authz_tokens=[],
         auth_service_token_getters={},
         bound_params={},
         client_headers={},
@@ -361,6 +364,7 @@ def test_tool_init_with_client_headers(
         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,
@@ -388,6 +392,7 @@ def test_tool_init_header_auth_conflict(
             description=sample_tool_description,
             params=sample_tool_auth_params,
             required_authn_params={},
+            required_authz_tokens=[],
             auth_service_token_getters=auth_getters,
             bound_params={},
             client_headers=conflicting_client_header,
@@ -410,6 +415,7 @@ def test_tool_add_auth_token_getters_conflict_with_existing_client_header(
         description=sample_tool_description,
         params=sample_tool_params,
         required_authn_params={},
+        required_authz_tokens=[],
         auth_service_token_getters={},
         bound_params={},
         client_headers={
