diff --git a/tests/test_session.py b/tests/test_session.py index 6409dc7f..14dcdd24 100644 --- a/tests/test_session.py +++ b/tests/test_session.py @@ -51,6 +51,7 @@ def session_constants(self): "roles": ["admin"], "permissions": ["read"], "entitlements": ["feature_1"], + "feature_flags": ["flag1", "flag2"], "exp": int(current_datetime.timestamp()) + 3600, "iat": int(current_datetime.timestamp()), } @@ -244,6 +245,7 @@ def test_authenticate_success(self, session_constants, mock_user_management): "roles": ["admin"], "permissions": ["read"], "entitlements": ["feature_1"], + "feature_flags": ["flag1", "flag2"], } with patch.object(Session, "unseal_data", return_value=mock_session), patch( @@ -263,6 +265,7 @@ def test_authenticate_success(self, session_constants, mock_user_management): assert response.roles == ["admin"] assert response.permissions == ["read"] assert response.entitlements == ["feature_1"] + assert response.feature_flags == ["flag1", "flag2"] assert response.user.id == session_constants["USER_ID"] assert response.impersonator is None @@ -312,6 +315,7 @@ def test_authenticate_success_with_roles( "roles": ["admin", "member"], "permissions": ["read", "write"], "entitlements": ["feature_1"], + "feature_flags": ["flag1", "flag2"], } with patch.object(Session, "unseal_data", return_value=mock_session), patch( @@ -331,6 +335,7 @@ def test_authenticate_success_with_roles( assert response.roles == ["admin", "member"] assert response.permissions == ["read", "write"] assert response.entitlements == ["feature_1"] + assert response.feature_flags == ["flag1", "flag2"] assert response.user.id == session_constants["USER_ID"] assert response.impersonator is None @@ -410,6 +415,7 @@ def test_refresh_success(self, session_constants, mock_user_management): "roles": ["admin"], "permissions": ["read"], "entitlements": ["feature_1"], + "feature_flags": ["flag1", "flag2"], }, ): response = session.refresh() @@ -511,6 +517,7 @@ async def test_refresh_success(self, session_constants, mock_user_management): "roles": ["admin"], "permissions": ["read"], "entitlements": ["feature_1"], + "feature_flags": ["flag1", "flag2"], }, ): response = await session.refresh() diff --git a/workos/session.py b/workos/session.py index 2052534e..f0ad6c45 100644 --- a/workos/session.py +++ b/workos/session.py @@ -4,6 +4,7 @@ from functools import lru_cache import json from typing import Any, Dict, Optional, Union, cast + import jwt from jwt import PyJWKClient from cryptography.fernet import Fernet @@ -107,6 +108,7 @@ def authenticate( entitlements=decoded.get("entitlements", None), user=session["user"], impersonator=session.get("impersonator", None), + feature_flags=decoded.get("feature_flags", None), ) def refresh( @@ -235,6 +237,7 @@ def refresh( entitlements=decoded.get("entitlements", None), user=auth_response.user, impersonator=auth_response.impersonator, + feature_flags=decoded.get("feature_flags", None), ) except Exception as e: return RefreshWithSessionCookieErrorResponse( @@ -326,6 +329,7 @@ async def refresh( entitlements=decoded.get("entitlements", None), user=auth_response.user, impersonator=auth_response.impersonator, + feature_flags=decoded.get("feature_flags", None), ) except Exception as e: return RefreshWithSessionCookieErrorResponse( diff --git a/workos/types/user_management/session.py b/workos/types/user_management/session.py index 1be8025f..06869a97 100644 --- a/workos/types/user_management/session.py +++ b/workos/types/user_management/session.py @@ -24,6 +24,7 @@ class AuthenticateWithSessionCookieSuccessResponse(WorkOSModel): user: User impersonator: Optional[Impersonator] = None entitlements: Optional[Sequence[str]] = None + feature_flags: Optional[Sequence[str]] = None class AuthenticateWithSessionCookieErrorResponse(WorkOSModel):