From 2de21b3449d3bc37574e74b5b8447090357aea87 Mon Sep 17 00:00:00 2001
From: RMidhunSuresh <hi@midhun.dev>
Date: Wed, 27 Apr 2022 15:34:02 +0530
Subject: [PATCH] Also check if token registration stage is present

---
 src/viewmodels/AccountSetupViewModel.ts | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/viewmodels/AccountSetupViewModel.ts b/src/viewmodels/AccountSetupViewModel.ts
index daccd6f..9e46c51 100644
--- a/src/viewmodels/AccountSetupViewModel.ts
+++ b/src/viewmodels/AccountSetupViewModel.ts
@@ -33,9 +33,11 @@ export class AccountSetupViewModel extends ViewModel {
                         "m.login.dummy"
                     ];
                     for (const flow of flows) {
-                        // Find the first flow that does not contain any unsupported stages
+                        // Find the first flow that does not contain any unsupported stages but contains Token registration stage.
                         const containsUnsupportedStage = flow.stages.some(stage => !allowedStages.includes(stage));
-                        if (!containsUnsupportedStage) {
+                        const containsTokenStage = flow.stages.includes("m.login.registration_token") ||
+                            flow.stages.includes("org.matrix.msc3231.login.registration_token");
+                        if (!containsUnsupportedStage && containsTokenStage) {
                             return flow;
                         }
                     }