From 27da621f44065f477e2d680aefc5aff7dcf888dd Mon Sep 17 00:00:00 2001
From: Kerem Yilmaz <kerem@skyvern.com>
Date: Sat, 29 Jun 2024 23:55:05 -0700
Subject: [PATCH] add api key expired message to the 403 when an api key is
 expired/invalid (#532)

---
 skyvern/forge/sdk/db/client.py                | 19 ++++++++++---------
 .../forge/sdk/services/org_auth_service.py    |  7 +++++++
 2 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/skyvern/forge/sdk/db/client.py b/skyvern/forge/sdk/db/client.py
index b18f01195..60d8e0ece 100644
--- a/skyvern/forge/sdk/db/client.py
+++ b/skyvern/forge/sdk/db/client.py
@@ -566,18 +566,19 @@ async def validate_org_auth_token(
         organization_id: str,
         token_type: OrganizationAuthTokenType,
         token: str,
+        valid: bool | None = True,
     ) -> OrganizationAuthToken | None:
         try:
             async with self.Session() as session:
-                if token_obj := (
-                    await session.scalars(
-                        select(OrganizationAuthTokenModel)
-                        .filter_by(organization_id=organization_id)
-                        .filter_by(token_type=token_type)
-                        .filter_by(token=token)
-                        .filter_by(valid=True)
-                    )
-                ).first():
+                query = (
+                    select(OrganizationAuthTokenModel)
+                    .filter_by(organization_id=organization_id)
+                    .filter_by(token_type=token_type)
+                    .filter_by(token=token)
+                )
+                if valid is not None:
+                    query = query.filter_by(valid=valid)
+                if token_obj := (await session.scalars(query)).first():
                     return convert_to_organization_auth_token(token_obj)
                 else:
                     return None
diff --git a/skyvern/forge/sdk/services/org_auth_service.py b/skyvern/forge/sdk/services/org_auth_service.py
index 6b48acf57..8dd3ad213 100644
--- a/skyvern/forge/sdk/services/org_auth_service.py
+++ b/skyvern/forge/sdk/services/org_auth_service.py
@@ -109,6 +109,7 @@ async def _get_current_org_cached(x_api_key: str, db: AgentDB) -> Organization:
         organization_id=organization.organization_id,
         token_type=OrganizationAuthTokenType.api,
         token=x_api_key,
+        valid=None,
     )
     if not api_key_db_obj:
         raise HTTPException(
@@ -116,6 +117,12 @@ async def _get_current_org_cached(x_api_key: str, db: AgentDB) -> Organization:
             detail="Invalid credentials",
         )
 
+    if api_key_db_obj.valid is False:
+        raise HTTPException(
+            status_code=status.HTTP_403_FORBIDDEN,
+            detail="Your API key has expired. Please retrieve the latest one from https://app.skyvern.com/settings",
+        )
+
     # set organization_id in skyvern context and log context
     context = skyvern_context.current()
     if context: