From 688a615fb231f95728076ba2f01edc0f7eec2fc1 Mon Sep 17 00:00:00 2001
From: Kerem Yilmaz <kerem@skyvern.com>
Date: Sun, 26 May 2024 14:51:53 -0700
Subject: [PATCH] introduce permission checker to task create api (#368)

---
 poetry.lock                                     | 17 ++++++++++++++++-
 pyproject.toml                                  |  1 +
 skyvern/forge/agent_functions.py                |  3 ++-
 skyvern/forge/sdk/core/permissions/__init__.py  |  0
 .../permissions/permission_checker_factory.py   | 13 +++++++++++++
 .../sdk/core/permissions/permission_checkers.py | 14 ++++++++++++++
 skyvern/forge/sdk/routes/agent_protocol.py      |  2 ++
 7 files changed, 48 insertions(+), 2 deletions(-)
 create mode 100644 skyvern/forge/sdk/core/permissions/__init__.py
 create mode 100644 skyvern/forge/sdk/core/permissions/permission_checker_factory.py
 create mode 100644 skyvern/forge/sdk/core/permissions/permission_checkers.py

diff --git a/poetry.lock b/poetry.lock
index 95bd6b4f..67ac35f7 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -6061,6 +6061,21 @@ docs = ["myst-parser[linkify]", "sphinx", "sphinx-rtd-theme"]
 release = ["twine"]
 test = ["pylint", "pytest", "pytest-black", "pytest-cov", "pytest-pylint"]
 
+[[package]]
+name = "stripe"
+version = "9.7.0"
+description = "Python bindings for the Stripe API"
+optional = false
+python-versions = ">=3.6"
+files = [
+    {file = "stripe-9.7.0-py2.py3-none-any.whl", hash = "sha256:266af2f4ff23ca3cdc73c332cf2776fd2e1084b1e9379b03fb981b1040ed69f0"},
+    {file = "stripe-9.7.0.tar.gz", hash = "sha256:af694723bf7968cea18a956641dcef786ec1f72de8022718fa3ef5f7868bd430"},
+]
+
+[package.dependencies]
+requests = {version = ">=2.20", markers = "python_version >= \"3.0\""}
+typing-extensions = {version = ">=4.5.0", markers = "python_version >= \"3.7\""}
+
 [[package]]
 name = "structlog"
 version = "23.3.0"
@@ -7190,4 +7205,4 @@ testing = ["big-O", "jaraco.functools", "jaraco.itertools", "more-itertools", "p
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.11,<3.12"
-content-hash = "0dc27842c8de420ea8e4d4573e2e062f4a99d27b0e8e4c398abea4b7137b7e2d"
+content-hash = "525799973a084af745be22b7c8ead47729df7ed19e0ec1cfc94f25b20b893693"
diff --git a/pyproject.toml b/pyproject.toml
index ab733d0d..a9a32a4f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -55,6 +55,7 @@ supabase = "^2.4.3"
 aioredlock = "^0.7.3"
 fpdf = "^1.7.2"
 pypdf = "^4.2.0"
+stripe = "^9.7.0"
 
 [tool.poetry.group.dev.dependencies]
 isort = "^5.13.2"
diff --git a/skyvern/forge/agent_functions.py b/skyvern/forge/agent_functions.py
index bbdafe22..cc51a561 100644
--- a/skyvern/forge/agent_functions.py
+++ b/skyvern/forge/agent_functions.py
@@ -1,5 +1,6 @@
 from playwright.async_api import Page
 
+from skyvern.exceptions import StepTerminationError
 from skyvern.forge import app
 from skyvern.forge.async_operations import AsyncOperation
 from skyvern.forge.sdk.models import Organization, Step, StepStatus
@@ -33,7 +34,7 @@ class AgentFunction:
 
         can_execute = has_valid_task_status and has_valid_step_status and has_no_running_steps
         if not can_execute:
-            raise Exception(f"Cannot execute step. Reasons: {reasons}, Step: {step}")
+            raise StepTerminationError(step_id=step.step_id, reason="Cannot execute step. Reasons: {reasons}")
 
     def generate_async_operations(
         self,
diff --git a/skyvern/forge/sdk/core/permissions/__init__.py b/skyvern/forge/sdk/core/permissions/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/skyvern/forge/sdk/core/permissions/permission_checker_factory.py b/skyvern/forge/sdk/core/permissions/permission_checker_factory.py
new file mode 100644
index 00000000..fbd1058c
--- /dev/null
+++ b/skyvern/forge/sdk/core/permissions/permission_checker_factory.py
@@ -0,0 +1,13 @@
+from skyvern.forge.sdk.core.permissions.permission_checkers import NoopPermissionChecker, PermissionChecker
+
+
+class PermissionCheckerFactory:
+    __instance: PermissionChecker = NoopPermissionChecker()
+
+    @staticmethod
+    def get_instance() -> PermissionChecker:
+        return PermissionCheckerFactory.__instance
+
+    @staticmethod
+    def set_instance(permission_checker: PermissionChecker) -> None:
+        PermissionCheckerFactory.__instance = permission_checker
diff --git a/skyvern/forge/sdk/core/permissions/permission_checkers.py b/skyvern/forge/sdk/core/permissions/permission_checkers.py
new file mode 100644
index 00000000..2049ee3a
--- /dev/null
+++ b/skyvern/forge/sdk/core/permissions/permission_checkers.py
@@ -0,0 +1,14 @@
+import abc
+
+from skyvern.forge.sdk.models import Organization
+
+
+class PermissionChecker(abc.ABC):
+    @abc.abstractmethod
+    async def check(self, organization: Organization) -> None:
+        pass
+
+
+class NoopPermissionChecker(PermissionChecker):
+    async def check(self, organization: Organization) -> None:
+        return
diff --git a/skyvern/forge/sdk/routes/agent_protocol.py b/skyvern/forge/sdk/routes/agent_protocol.py
index c4d761b3..6f32ea9a 100644
--- a/skyvern/forge/sdk/routes/agent_protocol.py
+++ b/skyvern/forge/sdk/routes/agent_protocol.py
@@ -11,6 +11,7 @@ from skyvern.exceptions import StepNotFound
 from skyvern.forge import app
 from skyvern.forge.sdk.artifact.models import Artifact, ArtifactType
 from skyvern.forge.sdk.core import skyvern_context
+from skyvern.forge.sdk.core.permissions.permission_checker_factory import PermissionCheckerFactory
 from skyvern.forge.sdk.core.security import generate_skyvern_signature
 from skyvern.forge.sdk.executor.factory import AsyncExecutorFactory
 from skyvern.forge.sdk.models import Organization, Step
@@ -99,6 +100,7 @@ async def create_agent_task(
     x_max_steps_override: Annotated[int | None, Header()] = None,
 ) -> CreateTaskResponse:
     analytics.capture("skyvern-oss-agent-task-create", data={"url": task.url})
+    await PermissionCheckerFactory.get_instance().check(current_org)
 
     if current_org and current_org.organization_name == "CoverageCat":
         task.proxy_location = ProxyLocation.RESIDENTIAL