Extract feature flags (#480)

Co-authored-by: Kendall Strautman Swarthout <36613477+kendallstrautman@users.noreply.github.com>

Author: gobbledygook88

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()

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(

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):