fix: Validate actions available to an entity in check_connected_account (#1095)

Co-authored-by: Devanshusisodiya <[email protected]>
Co-authored-by: Viraj <[email protected]>
Co-authored-by: Devanshu Kumar Singh <[email protected]>
Co-authored-by: ellipsis-dev[bot] <65095814+ellipsis-dev[bot]@users.noreply.github.com>

Author: 4 people

python/composio/client/collections.py

@@ -30,7 +30,7 @@
     Trigger,
     TriggerType,
 )
-from composio.constants import PUSHER_CLUSTER, PUSHER_KEY
+from composio.constants import DEFAULT_ENTITY_ID, PUSHER_CLUSTER, PUSHER_KEY
 from composio.exceptions import (
     ErrorFetchingResource,
     InvalidParams,
@@ -92,7 +92,7 @@ class ConnectedAccountModel(BaseModel):
     connectionParams: AuthConnectionParamsModel
 
     clientUniqueUserId: t.Optional[str] = None
-    entityId: t.Optional[str] = None
+    entityId: str = DEFAULT_ENTITY_ID
 
     # Override arbitrary model config.
     model_config: ConfigDict = ConfigDict(  # type: ignore

python/composio/tools/toolset.py

@@ -1668,7 +1668,9 @@ def _validate_connection_ids(
             return valid
         raise InvalidConnectedAccount(f"Invalid connected accounts found: {invalid}")
 
-    def check_connected_account(self, action: ActionType) -> None:
+    def check_connected_account(
+        self, action: ActionType, entity_id: t.Optional[str] = None
+    ) -> None:
         """Check if connected account is required and if required it exists or not."""
         action = Action(action)
         if action.no_auth or action.is_runtime:
@@ -1686,6 +1688,7 @@ def check_connected_account(self, action: ActionType) -> None:
         if action.app not in [
             connection.appUniqueId.upper()  # Normalize app names/ids coming from API
             for connection in self._connected_accounts
+            if entity_id is None or connection.clientUniqueUserId == entity_id
         ]:
             raise ConnectedAccountNotFoundError(
                 f"No connected account found for app `{action.app}`; "
@@ -1799,7 +1802,7 @@ def _execute_remote(
             action=action
         )
         if auth is None:
-            self.check_connected_account(action=action)
+            self.check_connected_account(action=action, entity_id=entity_id)
 
         output = self.client.get_entity(  # pylint: disable=protected-access
             id=entity_id

python/tests/test_tools/test_toolset.py

@@ -12,7 +12,11 @@
 from pydantic import BaseModel, Field
 
 from composio import Action, App, Trigger
-from composio.exceptions import ApiKeyNotProvidedError, ComposioSDKError
+from composio.exceptions import (
+    ApiKeyNotProvidedError,
+    ComposioSDKError,
+    ConnectedAccountNotFoundError,
+)
 from composio.tools.base.abs import action_registry, tool_registry
 from composio.tools.base.runtime import action as custom_action
 from composio.tools.local.filetool.tool import Filetool, FindFile
@@ -289,6 +293,22 @@ def postprocess(response: dict) -> dict:
         assert postprocessor_called
 
 
+def test_entity_id_validation_in_check_connected_accounts() -> None:
+    """Test whether check_connected_account raises error with invalid entity_id"""
+    toolset = ComposioToolSet()
+    with pytest.raises(
+        ConnectedAccountNotFoundError,
+        match=(
+            "No connected account found for app `GMAIL`; "
+            "Run `composio add gmail` to fix this"
+        ),
+    ):
+        toolset.check_connected_account(
+            action=Action.GMAIL_FETCH_EMAILS,
+            entity_id="some_very_random_obviously_wrong_entity_id",
+        )
+
+
 def test_check_connected_accounts_flag() -> None:
     """Test the `check_connected_accounts` flag on `get_tools()`."""