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

Pwuts · 2024-09-26T15:25:30Z

Resolves Ensure OAuth credentials are fresh before using them in graph execution #8179

Changes 🏗️

`backend`

Added Pydantic model (de)serialization support to @expose decorator

`backend.executor`

Change credential injection mechanism to acquire credentials just before execution
- Also locks the credentials for the duration of the execution

`backend.server`

Add IntegrationCredentialsManager to handle and synchronize credentials-related operations
Make backend.server.integrations module with router, creds_manager, and utils in it

`autogpt_libs`

Add mutexes to SupabaseIntegrationCredentialsStore to ensure thread-safety
Move KeyedMutex to autogpt_libs.utils.synchronize

…graph execution

netlify · 2024-09-26T15:26:10Z

✅ Deploy Preview for auto-gpt-docs ready!

Name	Link
🔨 Latest commit	`fab700b`
🔍 Latest deploy log	https://app.netlify.com/sites/auto-gpt-docs/deploys/66fd4f3502b98900096b4776
😎 Deploy Preview	https://deploy-preview-8191--auto-gpt-docs.netlify.app
📱 Preview on mobile	Toggle QR Code... Use your smartphone camera to open QR code link.

To edit notification comments on pull requests, go to your Netlify site configuration.

- Move `KeyedMutex` to `autogpt_libs.utils.synchronize` - Add locks to transactions in `SupabaseIntegrationCredentialsStore`

github-actions · 2024-09-27T20:32:17Z

This pull request has conflicts with the base branch, please resolve those so we can evaluate the pull request.

…entials JIT Also: - Add exposed `acquire_credentials`, `release_credentials` methods to `AgentServer` - Add Pydantic model auto-deserialization support to `@expose` decorator - Make `ExecutionManager.agent_server_client` cached

…-are-fresh-before-using-them-in

github-actions · 2024-09-27T22:33:24Z

Conflicts have been resolved! 🎉 A maintainer will review the pull request shortly.

github-actions · 2024-09-29T21:26:57Z

This pull request has conflicts with the base branch, please resolve those so we can evaluate the pull request.

…-are-fresh-before-using-them-in

CLAassistant · 2024-09-30T09:57:34Z

All committers have signed the CLA.

github-actions · 2024-09-30T09:57:43Z

Conflicts have been resolved! 🎉 A maintainer will review the pull request shortly.

majdyz

Code looks good, though we are using KeyedMutex in many places, where this is planned to be killed:
#8197

We need to sync on this or make a helper lies in a single place for an easy replacement

majdyz · 2024-09-30T18:16:18Z

autogpt_platform/autogpt_libs/autogpt_libs/supabase_integration_credentials_store/store.py

@@ -14,26 +17,29 @@


 class SupabaseIntegrationCredentialsStore:
+    locks = KeyedMutex()


FYI: This will be gone very soon: https://github.com/Significant-Gravitas/AutoGPT/pull/8197/files#diff-39d77f5b792b124e499a265155ab201f0a211a16c57129cec51d32ba4f0c695d

Let me know if there will be any issue migrating this to Redis.

It shouldn't be an issue, other than extra time spent. I just need the locking service to be available within autogpt_libs.

majdyz · 2024-09-30T18:19:53Z

autogpt_platform/autogpt_libs/autogpt_libs/supabase_integration_credentials_store/store.py

+            key = (key,)
+        key = (self.supabase.supabase_url, *key)
+
+        self.locks.lock(key)


I think it's better to add a key prefix here to avoid clashing with another possible usage

Isn't that the point of locking resources? It should clash if it's simultaneously in use elsewhere.

The atomic unit of data here is the entire user metadata object (because it needs to be read/written in its entirety), so that dictates the lock scope: with self.locked(user_id)

locks is shared across all SupabaseIntegrationCredentialsStore instances in a process

I'm using supabase_url as a key prefix to account for the edge case of multiple connections to different Supabase instances

majdyz · 2024-09-30T18:26:05Z

autogpt_platform/backend/backend/executor/manager.py

+        # Release lock on credentials ASAP
+        if credentials:
+            api_client.release_credentials(user_id, credentials.id)
+            credentials = None


Why do we need this if the finally block already does the same

Like the comment says:

Release lock on credentials ASAP

because between this and the finally block being called, some other stuff happens and I want to avoid the extra latency

majdyz · 2024-09-30T18:28:40Z

autogpt_platform/backend/backend/executor/manager.py

@@ -731,7 +724,7 @@ def add_execution(
            raise Exception(f"Graph #{graph_id} not found.")

        graph.validate_graph(for_run=True)
-        node_input_credentials = self._get_node_input_credentials(graph, user_id)
+        self._validate_node_input_credentials(graph, user_id)


Don't we need to do this on each node execution?

Why? This checks if the credentials exist / are available. If this is all good at the start of graph execution, we can start executing. If halfway through execution the user somehow deletes the in-use credentials, there's no use in re-checking if they exist before trying to use them because it will result in execution failure either way.

majdyz · 2024-09-30T18:31:29Z

autogpt_platform/backend/backend/executor/manager.py

@@ -720,6 +711,8 @@ def cleanup(self):

    @property
    def agent_server_client(self) -> "AgentServer":
+        # Since every single usage of this property happens from a different thread,
+        # there is no value in caching it.
        return get_agent_server_client()


Sharing client between the different thread is forbidden.
The comment is contradicting with the reality though, what's the objective? should we remove the property decorator then?

I put that comment there to prevent future confusion and time wasted on this. I tried to make the property cached because it seemed like a waste to re-create the client for every single use. Then I found out that that's illegal and will cause errors because a client can only be owned by one thread. I reverted the technical change and left a note for the next person who might think it's a good idea to re-use the client.

autogpt_platform/backend/backend/server/integrations/creds_manager.py

majdyz · 2024-09-30T18:36:13Z

autogpt_platform/backend/backend/server/integrations/creds_manager.py

+        key = (self.store.supabase.supabase_url, *key)
+        self._locks.unlock(key)
+
+    @contextmanager


At this point we might as well bring this feature to synchronize.py

majdyz · 2024-09-30T18:37:00Z

autogpt_platform/backend/backend/server/integrations/creds_manager.py

+
+    def delete(self, user_id: str, credentials_id: str) -> None:
+        with self._locked(user_id, credentials_id):
+            self.store.delete_creds_by_id(user_id, credentials_id)


Do we actually need lock for it?

Without this lock, we would allow deleting a set of credentials while it's in use.

majdyz · 2024-09-30T18:41:45Z

autogpt_platform/backend/backend/util/service.py

@@ -40,9 +52,52 @@ def wrapper(*args, **kwargs):
            logger.exception(msg)
            raise Exception(msg, e)

+    # Register custom serializers and deserializers for annotated Pydantic models
+    for name, annotation in func.__annotations__.items():


Nice!
Though can we extract this out as a function, we definitely need a test for the peeling logic, something like test_types.py testing things like Union[A,B], A|B, list[dict[str, int]]

IMO properly documenting the intended use and/or limitations is more efficient and effective here than building exhaustive unit tests. That's what the "Gotcha" in the docstring is for. Not because unit tests are bad, but predicting use cases is hard and accounting for all of the potential use cases in unit tests would take up way too much time and make this unnecessarily complex.

This implementation is based on how we currently use it. It supports Annotated and Union, and IMO that's as complex as it should get with parameter or return annotations containing Pydantic models.

The complexity of the peeling logic also increases with the complexity of the types that you want it to support. Regarding your examples:

list[dict[str, int]] isn't relevant here, because we only care about Pydantic models

Union[A, B] == A | B, which is supported

Extracted the type peeling logic into a helper function: 94fba3c

….get(..)`

…-are-fresh-before-using-them-in

github-actions · 2024-10-03T16:38:47Z

This pull request has conflicts with the base branch, please resolve those so we can evaluate the pull request.

feat(platform/executor): Refresh OAuth credentials before starting a …

ffcf200

…graph execution

github-actions bot added platform/backend AutoGPT Platform - Back end size/m labels Sep 26, 2024

feat(libs): Make SupabaseIntegrationCredentialsStore thread safe

5850b13

- Move `KeyedMutex` to `autogpt_libs.utils.synchronize` - Add locks to transactions in `SupabaseIntegrationCredentialsStore`

github-actions bot added size/l conflicts Automatically applied to PRs with merge conflicts and removed size/m labels Sep 27, 2024

Pwuts added 2 commits September 28, 2024 00:16

feat(backend): Add IntegrationCredentialsManager

1f4d0ad

github-actions bot added size/xl and removed size/l labels Sep 27, 2024

Merge branch 'master' into reinier/open-1891-ensure-oauth-credentials…

478b59c

…-are-fresh-before-using-them-in

github-actions bot removed the conflicts Automatically applied to PRs with merge conflicts label Sep 27, 2024

Pwuts added the platform/backend/executor label Sep 27, 2024

Pwuts marked this pull request as ready for review September 27, 2024 22:33

Pwuts requested a review from a team as a code owner September 27, 2024 22:33

Pwuts requested review from aarushik93, majdyz and ntindle and removed request for a team September 27, 2024 22:33

Pwuts force-pushed the reinier/open-1891-ensure-oauth-credentials-are-fresh-before-using-them-in branch from d48fbfe to fda2cf6 Compare September 27, 2024 22:48

Pwuts added 2 commits September 28, 2024 00:59

Remove unused _get_provider_oauth_handler from executor.manager

8267679

fix process forking issue in AgentServer

5f3ef42

Pwuts force-pushed the reinier/open-1891-ensure-oauth-credentials-are-fresh-before-using-them-in branch from fda2cf6 to 5f3ef42 Compare September 27, 2024 23:00

Pwuts added 2 commits September 28, 2024 01:59

remove broken agent_server_client caching

da29cce

lint

b464e76

github-actions bot added the conflicts Automatically applied to PRs with merge conflicts label Sep 29, 2024

Merge branch 'master' into reinier/open-1891-ensure-oauth-credentials…

dadd262

…-are-fresh-before-using-them-in

github-actions bot removed the conflicts Automatically applied to PRs with merge conflicts label Sep 30, 2024

Pwuts changed the title ~~feat(platform/executor): Ensure validity of OAuth credentials during graph execution~~ feat(backend): Ensure validity of OAuth credentials during graph execution Sep 30, 2024

Pwuts enabled auto-merge (squash) September 30, 2024 16:41

majdyz previously approved these changes Sep 30, 2024

View reviewed changes

fix refresh check concurrency issue in `IntegrationCredentialsManager…

2095326

….get(..)`

Pwuts dismissed majdyz’s stale review via 2095326 October 1, 2024 13:02

Extract annotation type peeling logic from @expose

94fba3c

Pwuts mentioned this pull request Oct 1, 2024

feat(platform): List and revoke credentials in user profile #8207

Open

Merge branch 'master' into reinier/open-1891-ensure-oauth-credentials…

fab700b

…-are-fresh-before-using-them-in

github-actions bot added the conflicts Automatically applied to PRs with merge conflicts label Oct 3, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

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

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

Pwuts commented Sep 26, 2024 •

edited

Loading

netlify bot commented Sep 26, 2024 •

edited

Loading

github-actions bot commented Sep 27, 2024

github-actions bot commented Sep 27, 2024

github-actions bot commented Sep 29, 2024

CLAassistant commented Sep 30, 2024 •

edited

Loading

github-actions bot commented Sep 30, 2024

majdyz left a comment

majdyz Sep 30, 2024

Pwuts Oct 1, 2024 •

edited

Loading

majdyz Sep 30, 2024

Pwuts Oct 1, 2024

majdyz Sep 30, 2024

Pwuts Oct 1, 2024 •

edited

Loading

majdyz Sep 30, 2024

Pwuts Oct 1, 2024 •

edited

Loading

majdyz Sep 30, 2024

Pwuts Oct 1, 2024

majdyz Sep 30, 2024

majdyz Sep 30, 2024

Pwuts Oct 1, 2024

majdyz Sep 30, 2024

Pwuts Oct 1, 2024

Pwuts Oct 1, 2024

github-actions bot commented Oct 3, 2024

		@@ -14,26 +17,29 @@


		class SupabaseIntegrationCredentialsStore:
		locks = KeyedMutex()

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

Are you sure you want to change the base?

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

Conversation

Pwuts commented Sep 26, 2024 • edited Loading

Changes 🏗️

backend

backend.executor

backend.server

autogpt_libs

netlify bot commented Sep 26, 2024 • edited Loading

✅ Deploy Preview for auto-gpt-docs ready!

github-actions bot commented Sep 27, 2024

github-actions bot commented Sep 27, 2024

github-actions bot commented Sep 29, 2024

CLAassistant commented Sep 30, 2024 • edited Loading

github-actions bot commented Sep 30, 2024

majdyz left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Pwuts Oct 1, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Pwuts Oct 1, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Pwuts Oct 1, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

github-actions bot commented Oct 3, 2024

Pwuts commented Sep 26, 2024 •

edited

Loading

`backend`

`backend.executor`

`backend.server`

`autogpt_libs`

netlify bot commented Sep 26, 2024 •

edited

Loading

CLAassistant commented Sep 30, 2024 •

edited

Loading

Pwuts Oct 1, 2024 •

edited

Loading

Pwuts Oct 1, 2024 •

edited

Loading

Pwuts Oct 1, 2024 •

edited

Loading