feat(mfa): Implement session persistence for MFA tokens in the client

Author: subhankarmaiti

examples/MFA.md

@@ -30,6 +30,9 @@ The Auth0 MFA API allows you to manage multi-factor authentication for users in
     - [Verify with OOB (SMS or Email)](#verify-with-oob-sms-or-email)
     - [Verify with OTP](#verify-with-otp)
     - [Verify with Recovery Code](#verify-with-recovery-code)
+  - [Session Persistence](#session-persistence)
+    - [Automatic Session Update](#automatic-session-update)
+    - [Manual Session Update](#manual-session-update)
   - [Complete MFA Flow Examples](#complete-mfa-flow-examples)
     - [Enrollment Flow](#enrollment-flow)
     - [Challenge Flow](#challenge-flow)
@@ -311,34 +314,43 @@ try:
     verify_response = await server_client.mfa.verify({
         "mfa_token": mfa_token,
         "oob_code": challenge.oob_code,
-        "binding_code": "123456"  # Code user received via SMS/Email
+        "binding_code": "123456",  # Code user received via SMS/Email
+        "persist": True,  # Persist tokens to session store
+        "audience": "https://api.example.com",  # Required when persist=True
+        "scope": "openid profile email"  # Optional scope
     })
 
     access_token = verify_response.access_token
     id_token = verify_response.id_token
-    token_type = verify_response.token_type
 
     print(f"MFA verification successful!")
     print(f"Access Token: {access_token}")
     print(f"ID Token: {id_token}")
+    print("Tokens have been persisted to session store")
 
 except Exception as error:
     print(f"Verification failed: {error}")
 ```
 
+> [!NOTE]
+> Setting `persist=True` automatically updates the session store with the new tokens, similar to nextjs-auth0 and auth0-spa-js SDKs. This eliminates the need for manual token management after MFA verification.
+
 ### Verify with OTP
 
 ```python
 try:
     verify_response = await server_client.mfa.verify({
         "mfa_token": mfa_token,
-        "otp": "123456"  # 6-digit code from authenticator app
+        "otp": "123456",  # 6-digit code from authenticator app
+        "persist": True,  # Persist tokens to session store
+        "audience": "https://api.example.com",  # Required when persist=True
+        "scope": "openid profile email"
     })
 
     access_token = verify_response.access_token
-    id_token = verify_response.id_token
 
     print("MFA verification successful!")
+    print("Tokens have been persisted to session store")
 
 except Exception as error:
     print(f"Invalid OTP code: {error}")
@@ -352,17 +364,64 @@ Recovery codes can be used to complete MFA verification without initiating a cha
 try:
     verify_response = await server_client.mfa.verify({
         "mfa_token": mfa_token,
-        "recovery_code": "XXXX-XXXX-XXXX"  # One of the recovery codes
+        "recovery_code": "XXXX-XXXX-XXXX",  # One of the recovery codes
+        "persist": True,  # Persist tokens to session store
+        "audience": "https://api.example.com"  # Required when persist=True
     })
 
     access_token = verify_response.access_token
 
     print("MFA verification successful using recovery code!")
+    print("Tokens have been persisted to session store")
 
 except Exception as error:
     print(f"Verification failed: {error}")
 ```
 
+## Session Persistence
+
+By default, `verify()` returns tokens without persisting them to the session store. However, you can automatically persist tokens by setting `persist=True`, similar to how nextjs-auth0 and auth0-spa-js handle MFA.
+
+### Automatic Session Update
+
+When you set `persist=True`, the SDK will:
+1. Update the session's `access_token` for the specified audience
+2. Update the session's `id_token` if present
+3. Add the token to the `token_sets` array with expiration information
+
+```python
+verify_response = await server_client.mfa.verify({
+    "mfa_token": mfa_token,
+    "otp": "123456",
+    "persist": True,  # Enable automatic persistence
+    "audience": "https://api.example.com",  # Required when persist=True
+    "scope": "openid profile email"  # Optional
+})
+
+# Tokens are now available in the session store
+# User can call server_client.get_user() to access updated session
+user = await server_client.get_user()
+```
+
+### Manual Session Update
+
+If you prefer to manage session updates yourself:
+
+```python
+verify_response = await server_client.mfa.verify({
+    "mfa_token": mfa_token,
+    "otp": "123456"
+    # persist=False (default)
+})
+
+# Handle token storage manually if needed
+access_token = verify_response.access_token
+id_token = verify_response.id_token
+
+# Store tokens in your application's session management
+await my_session_store.update_tokens(access_token, id_token)
+```
+
 ## Complete MFA Flow Examples
 
 ### Enrollment Flow
@@ -396,10 +455,14 @@ async def handle_mfa_enrollment_flow(server_client, mfa_token):
         # Verify enrollment
         verify_response = await server_client.mfa.verify({
             "mfa_token": mfa_token,
-            "otp": user_code
+            "otp": user_code,
+            "persist": True,
+            "audience": "https://api.example.com",
+            "scope": "openid profile email"
         })
 
         print("MFA enrollment successful!")
+        print("Tokens have been persisted to session store")
         return verify_response.access_token
 
     except Exception as error:
@@ -441,17 +504,24 @@ async def handle_mfa_challenge_flow(server_client, mfa_token):
             user_code = input("Enter 6-digit code from authenticator: ")
             verify_response = await server_client.mfa.verify({
                 "mfa_token": mfa_token,
-                "otp": user_code
+                "otp": user_code,
+                "persist": True,
+                "audience": "https://api.example.com",
+                "scope": "openid profile email"
             })
         else:
             user_code = input(f"Enter code from {selected_auth.authenticator_type}: ")
             verify_response = await server_client.mfa.verify({
                 "mfa_token": mfa_token,
                 "oob_code": challenge.oob_code,
-                "binding_code": user_code
+                "binding_code": user_code,
+                "persist": True,
+                "audience": "https://api.example.com",
+                "scope": "openid profile email"
             })
 
         print("MFA verification successful!")
+        print("Tokens have been persisted to session store")
         return verify_response.access_token
 
     except Exception as error:

src/auth0_server_python/auth_server/mfa_client.py

@@ -50,13 +50,17 @@ def __init__(
         domain: str,
         client_id: str,
         client_secret: str,
-        secret: str
+        secret: str,
+        state_store=None,
+        state_identifier: str = "_a0_session"
     ):
         self._domain = domain
         self._base_url = f"https://{domain}"
         self._client_id = client_id
         self._client_secret = client_secret
         self._secret = secret
+        self._state_store = state_store
+        self._state_identifier = state_identifier
 
     # ============================================================================
     # MFA TOKEN ENCRYPTION / DECRYPTION
@@ -283,6 +287,10 @@ async def verify(
                 - 'otp': OTP code
                 - 'oob_code' + 'binding_code': OOB verification
                 - 'recovery_code': Recovery code
+                - 'persist': bool (optional, default=False) - Persist tokens to state store
+                - 'audience': str (optional, required if persist=True) - Audience for token_set
+                - 'scope': str (optional) - Scope for token_set
+                - 'store_options': dict (optional) - Store-specific options
 
         Returns:
             MfaVerifyResponse with access_token, token_type, etc.
@@ -348,11 +356,99 @@ async def verify(
                     )
 
                 token_response = response.json()
-                return MfaVerifyResponse(**token_response)
+                verify_response = MfaVerifyResponse(**token_response)
+
+                # Persist tokens to state store if requested
+                if options.get("persist") and self._state_store:
+                    await self._persist_mfa_tokens(
+                        verify_response=verify_response,
+                        options=options
+                    )
+
+                return verify_response
 
         except (MfaVerifyError, MfaRequiredError):
             raise
         except Exception as e:
             raise MfaVerifyError(
                 f"Unexpected error during MFA verification: {str(e)}"
             )
+
+    async def _persist_mfa_tokens(
+        self,
+        verify_response: MfaVerifyResponse,
+        options: dict[str, Any]
+    ) -> None:
+        """
+        Persist MFA verification tokens to the state store.
+
+        Updates the session with the new access_token and id_token from MFA verification.
+
+        Args:
+            verify_response: The response from verify() containing tokens
+            options: Dict containing:
+                - 'audience': str - Audience for token_set
+                - 'scope': str (optional) - Scope for token_set
+                - 'store_options': dict (optional) - Store-specific options
+        """
+        from auth0_server_python.auth_types import TokenSet, StateData
+        import time
+
+        audience = options.get("audience")
+        scope = options.get("scope")
+        store_options = options.get("store_options")
+
+        if not audience:
+            raise MfaVerifyError(
+                "audience is required when persist=True"
+            )
+
+        try:
+            # Get existing state
+            state_data = await self._state_store.get(
+                self._state_identifier,
+                store_options
+            )
+
+            if not state_data:
+                raise MfaVerifyError(
+                    "No existing session found to update with MFA tokens"
+                )
+
+            # Parse state data
+            existing_state = StateData(**state_data) if isinstance(state_data, dict) else state_data
+
+            # Update id_token if present
+            if verify_response.id_token:
+                existing_state.id_token = verify_response.id_token
+
+            # Create token_set for the access_token
+            expires_in = verify_response.get("expires_in", 86400)  # Default 24 hours
+            expires_at = int(time.time()) + expires_in
+
+            new_token_set = TokenSet(
+                audience=audience,
+                access_token=verify_response.access_token,
+                scope=scope,
+                expires_at=expires_at
+            )
+
+            # Add to token_sets, replacing any existing token_set for this audience
+            existing_state.token_sets = [
+                ts for ts in existing_state.token_sets if ts.audience != audience
+            ]
+            existing_state.token_sets.append(new_token_set)
+
+            # Persist updated state
+            await self._state_store.set(
+                self._state_identifier,
+                existing_state.model_dump() if hasattr(existing_state, 'model_dump') else existing_state,
+                options=store_options
+            )
+
+        except MfaVerifyError:
+            raise
+        except Exception as e:
+            raise MfaVerifyError(
+                f"Failed to persist MFA tokens to state store: {str(e)}"
+            )

src/auth0_server_python/auth_server/server_client.py

@@ -131,7 +131,9 @@ def __init__(
             domain=self._domain,
             client_id=self._client_id,
             client_secret=self._client_secret,
-            secret=self._secret
+            secret=self._secret,
+            state_store=self._state_store,
+            state_identifier=self._state_identifier
         )
 
     @property