adding jwt_leeway param to allow configuring the leeway param on PyJWT jwt.decode method

Author: Th3R3p0

workos/_base_client.py

@@ -25,6 +25,7 @@ class BaseClient(ClientConfiguration):
     _base_url: str
     _client_id: str
     _request_timeout: int
+    _jwt_leeway: float
 
     def __init__(
         self,
@@ -33,6 +34,7 @@ def __init__(
         client_id: Optional[str],
         base_url: Optional[str] = None,
         request_timeout: Optional[int] = None,
+        jwt_leeway: float = 0,
     ) -> None:
         api_key = api_key or os.getenv("WORKOS_API_KEY")
         if api_key is None:
@@ -63,6 +65,8 @@ def __init__(
             if request_timeout
             else int(os.getenv("WORKOS_REQUEST_TIMEOUT", DEFAULT_REQUEST_TIMEOUT))
         )
+        
+        self._jwt_leeway = jwt_leeway
 
     @property
     @abstractmethod
@@ -122,3 +126,7 @@ def client_id(self) -> str:
     @property
     def request_timeout(self) -> int:
         return self._request_timeout
+
+    @property
+    def jwt_leeway(self) -> float:
+        return self._jwt_leeway

workos/async_client.py

@@ -28,12 +28,14 @@ def __init__(
         client_id: Optional[str] = None,
         base_url: Optional[str] = None,
         request_timeout: Optional[int] = None,
+        jwt_leeway: float = 0,
     ):
         super().__init__(
             api_key=api_key,
             client_id=client_id,
             base_url=base_url,
             request_timeout=request_timeout,
+            jwt_leeway=jwt_leeway,
         )
         self._http_client = AsyncHTTPClient(
             api_key=self._api_key,

workos/client.py

@@ -28,12 +28,14 @@ def __init__(
         client_id: Optional[str] = None,
         base_url: Optional[str] = None,
         request_timeout: Optional[int] = None,
+        jwt_leeway: float = 0,
     ):
         super().__init__(
             api_key=api_key,
             client_id=client_id,
             base_url=base_url,
             request_timeout=request_timeout,
+            jwt_leeway=jwt_leeway,
         )
         self._http_client = SyncHTTPClient(
             api_key=self._api_key,

workos/session.py

@@ -28,6 +28,7 @@ class SessionModule(Protocol):
     cookie_password: str
     jwks: PyJWKClient
     jwk_algorithms: List[str]
+    jwt_leeway: float
 
     def __init__(
         self,
@@ -36,6 +37,7 @@ def __init__(
         client_id: str,
         session_data: str,
         cookie_password: str,
+        jwt_leeway: float = 0,
     ) -> None:
         # If the cookie password is not provided, throw an error
         if cookie_password is None or cookie_password == "":
@@ -45,6 +47,7 @@ def __init__(
         self.client_id = client_id
         self.session_data = session_data
         self.cookie_password = cookie_password
+        self.jwt_leeway = jwt_leeway
 
         self.jwks = PyJWKClient(self.user_management.get_jwks_url())
 
@@ -89,6 +92,7 @@ def authenticate(
             signing_key.key,
             algorithms=self.jwk_algorithms,
             options={"verify_aud": False},
+            leeway=self.jwt_leeway,
         )
 
         return AuthenticateWithSessionCookieSuccessResponse(
@@ -136,6 +140,7 @@ def _is_valid_jwt(self, token: str) -> bool:
                 signing_key.key,
                 algorithms=self.jwk_algorithms,
                 options={"verify_aud": False},
+                leeway=self.jwt_leeway,
             )
             return True
         except jwt.exceptions.InvalidTokenError:
@@ -167,6 +172,7 @@ def __init__(
         client_id: str,
         session_data: str,
         cookie_password: str,
+        jwt_leeway: float = 0,
     ) -> None:
         # If the cookie password is not provided, throw an error
         if cookie_password is None or cookie_password == "":
@@ -176,6 +182,7 @@ def __init__(
         self.client_id = client_id
         self.session_data = session_data
         self.cookie_password = cookie_password
+        self.jwt_leeway = jwt_leeway
 
         self.jwks = PyJWKClient(self.user_management.get_jwks_url())
 
@@ -228,6 +235,7 @@ def refresh(
                 signing_key.key,
                 algorithms=self.jwk_algorithms,
                 options={"verify_aud": False},
+                leeway=self.jwt_leeway,
             )
 
             return RefreshWithSessionCookieSuccessResponse(
@@ -257,6 +265,7 @@ def __init__(
         client_id: str,
         session_data: str,
         cookie_password: str,
+        jwt_leeway: float = 0,
     ) -> None:
         # If the cookie password is not provided, throw an error
         if cookie_password is None or cookie_password == "":
@@ -266,6 +275,7 @@ def __init__(
         self.client_id = client_id
         self.session_data = session_data
         self.cookie_password = cookie_password
+        self.jwt_leeway = jwt_leeway
 
         self.jwks = PyJWKClient(self.user_management.get_jwks_url())
 
@@ -318,6 +328,7 @@ async def refresh(
                 signing_key.key,
                 algorithms=self.jwk_algorithms,
                 options={"verify_aud": False},
+                leeway=self.jwt_leeway,
             )
 
             return RefreshWithSessionCookieSuccessResponse(

workos/user_management.py

@@ -866,6 +866,7 @@ def load_sealed_session(
             client_id=self._http_client.client_id,
             session_data=sealed_session,
             cookie_password=cookie_password,
+            jwt_leeway=self._client_configuration.jwt_leeway,
         )
 
     def get_user(self, user_id: str) -> User:
@@ -1491,6 +1492,7 @@ async def load_sealed_session(
             client_id=self._http_client.client_id,
             session_data=sealed_session,
             cookie_password=cookie_password,
+            jwt_leeway=self._client_configuration.jwt_leeway,
         )
 
     async def get_user(self, user_id: str) -> User: