diff --git a/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/AuthenticationService.java b/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/AuthenticationService.java
index e631c9a04f2..248dbbf4f5a 100644
--- a/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/AuthenticationService.java
+++ b/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/AuthenticationService.java
@@ -92,7 +92,7 @@ private AuthServiceResponse processCommonAuthResponse(AuthServiceRequestWrapper
         AuthServiceResponse authServiceResponse = new AuthServiceResponse();
 
         if (isAuthFlowSuccessful(request)) {
-            handleSuccessAuthResponse(request, authServiceResponse);
+            handleSuccessAuthResponse(request, response, authServiceResponse);
         } else if (isAuthFlowFailed(request, response)) {
             handleFailedAuthResponse(request, response, authServiceResponse);
         } else if (isAuthFlowIncomplete(request)) {
@@ -124,9 +124,10 @@ private void handleIntermediateAuthResponse(AuthServiceRequestWrapper request, A
         authServiceResponse.setData(responseData);
     }
 
-    private void handleSuccessAuthResponse(AuthServiceRequestWrapper request, AuthServiceResponse authServiceResponse) {
+    private void handleSuccessAuthResponse(AuthServiceRequestWrapper request, AuthServiceResponseWrapper response,
+                                           AuthServiceResponse authServiceResponse) throws AuthServiceException {
 
-        authServiceResponse.setSessionDataKey(request.getSessionDataKey());
+        authServiceResponse.setSessionDataKey(getFlowCompletionSessionDataKey(request, response));
         authServiceResponse.setFlowStatus(AuthServiceConstants.FlowStatus.SUCCESS_COMPLETED);
     }
 
@@ -272,8 +273,19 @@ private boolean isSentToRetryPageOnMissingContext(AuthServiceRequestWrapper requ
                 Boolean.TRUE.equals(request.getAttribute(FrameworkConstants.IS_SENT_TO_RETRY))) {
             Map<String, String> queryParams = AuthServiceUtils.extractQueryParams(response.getRedirectURL());
             return StringUtils.equals(queryParams.get(FrameworkConstants.STATUS_PARAM),
-                                      FrameworkConstants.ERROR_STATUS_AUTH_CONTEXT_NULL);
+                    FrameworkConstants.ERROR_STATUS_AUTH_CONTEXT_NULL);
         }
         return false;
     }
+
+    private String getFlowCompletionSessionDataKey(AuthServiceRequestWrapper request,
+                                                   AuthServiceResponseWrapper response) throws AuthServiceException {
+
+        String completionSessionDataKey = (String) request.getAttribute(FrameworkConstants.SESSION_DATA_KEY);
+        if (StringUtils.isBlank(completionSessionDataKey)) {
+            completionSessionDataKey = response.getSessionDataKey();
+        }
+
+        return completionSessionDataKey;
+    }
 }
diff --git a/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/model/auth/service/AuthServiceResponseWrapper.java b/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/model/auth/service/AuthServiceResponseWrapper.java
index fde174136a7..fafe596c6aa 100644
--- a/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/model/auth/service/AuthServiceResponseWrapper.java
+++ b/components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/model/auth/service/AuthServiceResponseWrapper.java
@@ -20,6 +20,7 @@
 
 import org.wso2.carbon.identity.application.authentication.framework.exception.auth.service.AuthServiceException;
 import org.wso2.carbon.identity.application.authentication.framework.model.CommonAuthResponseWrapper;
+import org.wso2.carbon.identity.application.authentication.framework.util.FrameworkConstants;
 import org.wso2.carbon.identity.application.authentication.framework.util.auth.service.AuthServiceConstants;
 import org.wso2.carbon.identity.application.authentication.framework.util.auth.service.AuthServiceUtils;
 
@@ -65,6 +66,18 @@ public boolean isErrorResponse() throws AuthServiceException {
         return Boolean.parseBoolean(queryParams.get(AuthServiceConstants.AUTH_FAILURE_PARAM));
     }
 
+    /**
+     * Get the sessionDataKey related to the authentication flow.
+     *
+     * @return String of sessionDataKey.
+     * @throws AuthServiceException
+     */
+    public String getSessionDataKey() throws AuthServiceException {
+
+        Map<String, String> queryParams = AuthServiceUtils.extractQueryParams(getRedirectURL());
+        return queryParams.get(FrameworkConstants.SESSION_DATA_KEY);
+    }
+
     @Override
     public String getRedirectURL() {