feat(backend): Ensure validity of OAuth credentials during graph execution

autogpt_platform/autogpt_libs/autogpt_libs/supabase_integration_credentials_store/__init__.py

@@ -1,8 +1,9 @@
 from .store import SupabaseIntegrationCredentialsStore
-from .types import APIKeyCredentials, OAuth2Credentials
+from .types import Credentials, APIKeyCredentials, OAuth2Credentials
 
 __all__ = [
     "SupabaseIntegrationCredentialsStore",
+    "Credentials",
     "APIKeyCredentials",
     "OAuth2Credentials",
 ]

autogpt_platform/autogpt_libs/autogpt_libs/supabase_integration_credentials_store/store.py

@@ -1,8 +1,12 @@
 import secrets
 from datetime import datetime, timedelta, timezone
-from typing import cast
+from typing import TYPE_CHECKING, cast
 
-from supabase import Client
+if TYPE_CHECKING:
+    from redis import Redis
+    from supabase import Client
+
+from autogpt_libs.utils.synchronize import RedisKeyedMutex
 
 from .types import (
     Credentials,
@@ -14,26 +18,28 @@
 
 
 class SupabaseIntegrationCredentialsStore:
-    def __init__(self, supabase: Client):
+    def __init__(self, supabase: "Client", redis: "Redis"):
         self.supabase = supabase
+        self.locks = RedisKeyedMutex(redis)
 
     def add_creds(self, user_id: str, credentials: Credentials) -> None:
-        if self.get_creds_by_id(user_id, credentials.id):
-            raise ValueError(
-                f"Can not re-create existing credentials with ID {credentials.id} "
-                f"for user with ID {user_id}"
+        with self.locked_user_metadata(user_id):
+            if self.get_creds_by_id(user_id, credentials.id):
+                raise ValueError(
+                    f"Can not re-create existing credentials #{credentials.id} "
+                    f"for user #{user_id}"
+                )
+            self._set_user_integration_creds(
+                user_id, [*self.get_all_creds(user_id), credentials]
             )
-        self._set_user_integration_creds(
-            user_id, [*self.get_all_creds(user_id), credentials]
-        )
 
     def get_all_creds(self, user_id: str) -> list[Credentials]:
         user_metadata = self._get_user_metadata(user_id)
         return UserMetadata.model_validate(user_metadata).integration_credentials
 
     def get_creds_by_id(self, user_id: str, credentials_id: str) -> Credentials | None:
-        credentials = self.get_all_creds(user_id)
-        return next((c for c in credentials if c.id == credentials_id), None)
+        all_credentials = self.get_all_creds(user_id)
+        return next((c for c in all_credentials if c.id == credentials_id), None)
 
     def get_creds_by_provider(self, user_id: str, provider: str) -> list[Credentials]:
         credentials = self.get_all_creds(user_id)
@@ -44,42 +50,45 @@ def get_authorized_providers(self, user_id: str) -> list[str]:
         return list(set(c.provider for c in credentials))
 
     def update_creds(self, user_id: str, updated: Credentials) -> None:
-        current = self.get_creds_by_id(user_id, updated.id)
-        if not current:
-            raise ValueError(
-                f"Credentials with ID {updated.id} "
-                f"for user with ID {user_id} not found"
-            )
-        if type(current) is not type(updated):
-            raise TypeError(
-                f"Can not update credentials with ID {updated.id} "
-                f"from type {type(current)} "
-                f"to type {type(updated)}"
-            )
-
-        # Ensure no scopes are removed when updating credentials
-        if (
-            isinstance(updated, OAuth2Credentials)
-            and isinstance(current, OAuth2Credentials)
-            and not set(updated.scopes).issuperset(current.scopes)
-        ):
-            raise ValueError(
-                f"Can not update credentials with ID {updated.id} "
-                f"and scopes {current.scopes} "
-                f"to more restrictive set of scopes {updated.scopes}"
-            )
-
-        # Update the credentials
-        updated_credentials_list = [
-            updated if c.id == updated.id else c for c in self.get_all_creds(user_id)
-        ]
-        self._set_user_integration_creds(user_id, updated_credentials_list)
+        with self.locked_user_metadata(user_id):
+            current = self.get_creds_by_id(user_id, updated.id)
+            if not current:
+                raise ValueError(
+                    f"Credentials with ID {updated.id} "
+                    f"for user with ID {user_id} not found"
+                )
+            if type(current) is not type(updated):
+                raise TypeError(
+                    f"Can not update credentials with ID {updated.id} "
+                    f"from type {type(current)} "
+                    f"to type {type(updated)}"
+                )
+
+            # Ensure no scopes are removed when updating credentials
+            if (
+                isinstance(updated, OAuth2Credentials)
+                and isinstance(current, OAuth2Credentials)
+                and not set(updated.scopes).issuperset(current.scopes)
+            ):
+                raise ValueError(
+                    f"Can not update credentials with ID {updated.id} "
+                    f"and scopes {current.scopes} "
+                    f"to more restrictive set of scopes {updated.scopes}"
+                )
+
+            # Update the credentials
+            updated_credentials_list = [
+                updated if c.id == updated.id else c
+                for c in self.get_all_creds(user_id)
+            ]
+            self._set_user_integration_creds(user_id, updated_credentials_list)
 
     def delete_creds_by_id(self, user_id: str, credentials_id: str) -> None:
-        filtered_credentials = [
-            c for c in self.get_all_creds(user_id) if c.id != credentials_id
-        ]
-        self._set_user_integration_creds(user_id, filtered_credentials)
+        with self.locked_user_metadata(user_id):
+            filtered_credentials = [
+                c for c in self.get_all_creds(user_id) if c.id != credentials_id
+            ]
+            self._set_user_integration_creds(user_id, filtered_credentials)
 
     async def store_state_token(
         self, user_id: str, provider: str, scopes: list[str]
@@ -94,14 +103,15 @@ async def store_state_token(
             scopes=scopes,
         )
 
-        user_metadata = self._get_user_metadata(user_id)
-        oauth_states = user_metadata.get("integration_oauth_states", [])
-        oauth_states.append(state.model_dump())
-        user_metadata["integration_oauth_states"] = oauth_states
+        with self.locked_user_metadata(user_id):
+            user_metadata = self._get_user_metadata(user_id)
+            oauth_states = user_metadata.get("integration_oauth_states", [])
+            oauth_states.append(state.model_dump())
+            user_metadata["integration_oauth_states"] = oauth_states
 
-        self.supabase.auth.admin.update_user_by_id(
-            user_id, {"user_metadata": user_metadata}
-        )
+            self.supabase.auth.admin.update_user_by_id(
+                user_id, {"user_metadata": user_metadata}
+            )
 
         return token
 
@@ -136,29 +146,30 @@ async def get_any_valid_scopes_from_state_token(
         return []
 
     async def verify_state_token(self, user_id: str, token: str, provider: str) -> bool:
-        user_metadata = self._get_user_metadata(user_id)
-        oauth_states = user_metadata.get("integration_oauth_states", [])
-
-        now = datetime.now(timezone.utc)
-        valid_state = next(
-            (
-                state
-                for state in oauth_states
-                if state["token"] == token
-                and state["provider"] == provider
-                and state["expires_at"] > now.timestamp()
-            ),
-            None,
-        )
-
-        if valid_state:
-            # Remove the used state
-            oauth_states.remove(valid_state)
-            user_metadata["integration_oauth_states"] = oauth_states
-            self.supabase.auth.admin.update_user_by_id(
-                user_id, {"user_metadata": user_metadata}
+        with self.locked_user_metadata(user_id):
+            user_metadata = self._get_user_metadata(user_id)
+            oauth_states = user_metadata.get("integration_oauth_states", [])
+
+            now = datetime.now(timezone.utc)
+            valid_state = next(
+                (
+                    state
+                    for state in oauth_states
+                    if state["token"] == token
+                    and state["provider"] == provider
+                    and state["expires_at"] > now.timestamp()
+                ),
+                None,
             )
-            return True
+
+            if valid_state:
+                # Remove the used state
+                oauth_states.remove(valid_state)
+                user_metadata["integration_oauth_states"] = oauth_states
+                self.supabase.auth.admin.update_user_by_id(
+                    user_id, {"user_metadata": user_metadata}
+                )
+                return True
 
         return False
 
@@ -178,3 +189,7 @@ def _get_user_metadata(self, user_id: str) -> UserMetadataRaw:
         if not response.user:
             raise ValueError(f"User with ID {user_id} not found")
         return cast(UserMetadataRaw, response.user.user_metadata)
+
+    def locked_user_metadata(self, user_id: str):
+        key = (self.supabase.supabase_url, f"user:{user_id}", "metadata")
+        return self.locks.locked(key)

autogpt_platform/autogpt_libs/autogpt_libs/utils/synchronize.py

@@ -0,0 +1,56 @@
+from contextlib import contextmanager
+from threading import Lock
+from typing import TYPE_CHECKING, Any
+
+from expiringdict import ExpiringDict
+
+if TYPE_CHECKING:
+    from redis import Redis
+    from redis.lock import Lock as RedisLock
+
+
+class RedisKeyedMutex:
+    """
+    This class provides a mutex that can be locked and unlocked by a specific key,
+    using Redis as a distributed locking provider.
+    It uses an ExpiringDict to automatically clear the mutex after a specified timeout,
+    in case the key is not unlocked for a specified duration, to prevent memory leaks.
+    """
+
+    def __init__(self, redis: "Redis", timeout: int | None = 60):
+        self.redis = redis
+        self.timeout = timeout
+        self.locks: dict[Any, "RedisLock"] = ExpiringDict(
+            max_len=6000, max_age_seconds=self.timeout
+        )
+        self.locks_lock = Lock()
+
+    @contextmanager
+    def locked(self, key: Any):
+        lock = self.acquire(key)
+        try:
+            yield
+        finally:
+            lock.release()
+
+    def acquire(self, key: Any) -> "RedisLock":
+        """Acquires and returns a lock with the given key"""
+        with self.locks_lock:
+            if key not in self.locks:
+                self.locks[key] = self.redis.lock(
+                    str(key), self.timeout, thread_local=False
+                )
+            lock = self.locks[key]
+        lock.acquire()
+        return lock
+
+    def release(self, key: Any):
+        if lock := self.locks.get(key):
+            lock.release()
+
+    def release_all_locks(self):
+        """Call this on process termination to ensure all locks are released"""
+        self.locks_lock.acquire(blocking=False)
+        for lock in self.locks.values():
+            if lock.locked() and lock.owned():
+                lock.release()

autogpt_platform/autogpt_libs/pyproject.toml

@@ -8,6 +8,7 @@ packages = [{ include = "autogpt_libs" }]
 
 [tool.poetry.dependencies]
 colorama = "^0.4.6"
+expiringdict = "^1.2.2"
 google-cloud-logging = "^3.8.0"
 pydantic = "^2.8.2"
 pydantic-settings = "^2.5.2"
@@ -16,6 +17,9 @@ python = ">=3.10,<4.0"
 python-dotenv = "^1.0.1"
 supabase = "^2.7.2"
 
+[tool.poetry.group.dev.dependencies]
+redis = "^5.0.8"
+
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"

autogpt_platform/backend/backend/app.py

@@ -25,7 +25,8 @@ def main(**kwargs):
     """
 
     from backend.executor import ExecutionManager, ExecutionScheduler
-    from backend.server import AgentServer, WebsocketServer
+    from backend.server.rest_api import AgentServer
+    from backend.server.ws_api import WebsocketServer
 
     run_processes(
         ExecutionManager(),

autogpt_platform/backend/backend/data/execution.py

@@ -3,7 +3,6 @@
 from multiprocessing import Manager
 from typing import Any, Generic, TypeVar
 
-from autogpt_libs.supabase_integration_credentials_store.types import Credentials
 from prisma.enums import AgentExecutionStatus
 from prisma.models import (
     AgentGraphExecution,
@@ -26,7 +25,6 @@ class GraphExecution(BaseModel):
     graph_exec_id: str
     graph_id: str
     start_node_execs: list["NodeExecution"]
-    node_input_credentials: dict[str, Credentials]  # dict[node_id, Credentials]
 
 
 class NodeExecution(BaseModel):