diff --git a/src/viewmodels/AccountSetupViewModel.ts b/src/viewmodels/AccountSetupViewModel.ts
index a774c1f63953f1bfb923e9521ca1f02ac7b55a43..6c874e931a994c903ab19a75ccfb49e3a65adf19 100644
--- a/src/viewmodels/AccountSetupViewModel.ts
+++ b/src/viewmodels/AccountSetupViewModel.ts
@@ -10,6 +10,7 @@ export class AccountSetupViewModel extends ViewModel {
     private _termsStage?: any;
     private _password: string;
     private _username: string;
+    private _registration: any;
 
     constructor(options) {
         super(options);
@@ -24,7 +25,8 @@ export class AccountSetupViewModel extends ViewModel {
         for (let i = 0; i < maxAttempts; ++i) {
             try {
                 this._username = `${this._config.username_prefix}-${generateUsername(10)}`;
-                let stage = await this._client.startRegistration(this._homeserver, this._username, this._password, "Chatterbox");
+                this._registration = await this._client.startRegistration(this._homeserver, this._username, this._password, "Chatterbox");
+                const stage = await this._registration.start();
                 if (stage.type === "m.login.terms") {
                     this._termsStage = stage;
                     this.emitChange("termsStage");
@@ -41,8 +43,8 @@ export class AccountSetupViewModel extends ViewModel {
 
     async completeRegistration() {
         let stage = this._termsStage;
-        while (typeof stage !== "string") {
-            stage = await stage.complete();
+        while (stage) {
+            stage = await this._registration.submitStage(stage);
         }
         const loginPromise = this.login(this._username, this._password);
         this.navigation.push("timeline", loginPromise);