feat(auth): implement view data persistence in 2FA flows

Add persistence to two-factor authentication in the extension login flow. Implements caching of form state to improve user experience when navigating between authentication steps. Includes feature flag for controlled rollout.

libs/auth/src/angular/two-factor-auth/child-components/two-factor-auth-authenticator.component.html

@@ -1,6 +1,13 @@
 <ng-container>
   <bit-form-field>
     <bit-label>{{ "verificationCode" | i18n }}</bit-label>
-    <input bitInput type="text" appAutofocus appInputVerbatim [formControl]="tokenFormControl" />
+    <input
+      bitInput
+      type="text"
+      appAutofocus
+      appInputVerbatim
+      [formControl]="tokenFormControl"
+      (keyup)="onTokenChange($event)"
+    />
   </bit-form-field>
 </ng-container>

libs/auth/src/angular/two-factor-auth/child-components/two-factor-auth-authenticator.component.ts

@@ -1,5 +1,5 @@
 import { CommonModule } from "@angular/common";
-import { Component, Input } from "@angular/core";
+import { Component, Input, Output, EventEmitter } from "@angular/core";
 import { ReactiveFormsModule, FormsModule, FormControl } from "@angular/forms";
 
 import { JslibModule } from "@bitwarden/angular/jslib.module";
@@ -32,4 +32,10 @@ import {
 })
 export class TwoFactorAuthAuthenticatorComponent {
   @Input({ required: true }) tokenFormControl: FormControl | undefined = undefined;
+  @Output() tokenChange = new EventEmitter<{ token: string }>();
+
+  onTokenChange(event: Event) {
+    const tokenValue = (event.target as HTMLInputElement).value || "";
+    this.tokenChange.emit({ token: tokenValue });
+  }
 }

libs/auth/src/angular/two-factor-auth/child-components/two-factor-auth-email/two-factor-auth-component-email-cache.service.spec.ts

@@ -0,0 +1,165 @@
+import { TestBed } from "@angular/core/testing";
+import { mock, MockProxy } from "jest-mock-extended";
+import { BehaviorSubject } from "rxjs";
+
+import { ViewCacheService } from "@bitwarden/angular/platform/abstractions/view-cache.service";
+import { FeatureFlag } from "@bitwarden/common/enums/feature-flag.enum";
+import { ConfigService } from "@bitwarden/common/platform/abstractions/config/config.service";
+
+import {
+  TwoFactorAuthEmailComponentCache,
+  TwoFactorAuthEmailComponentCacheService,
+} from "./two-factor-auth-email-component-cache.service";
+
+describe("TwoFactorAuthEmailCache", () => {
+  describe("fromJSON", () => {
+    it("returns null when input is null", () => {
+      const result = TwoFactorAuthEmailComponentCache.fromJSON(null as any);
+      expect(result).toBeNull();
+    });
+
+    it("creates a TwoFactorAuthEmailCache instance from valid JSON", () => {
+      const jsonData = { emailSent: true };
+      const result = TwoFactorAuthEmailComponentCache.fromJSON(jsonData);
+
+      expect(result).not.toBeNull();
+      expect(result).toBeInstanceOf(TwoFactorAuthEmailComponentCache);
+      expect(result?.emailSent).toBe(true);
+    });
+  });
+});
+
+describe("TwoFactorAuthEmailComponentCacheService", () => {
+  let service: TwoFactorAuthEmailComponentCacheService;
+  let mockViewCacheService: MockProxy<ViewCacheService>;
+  let mockConfigService: MockProxy<ConfigService>;
+  let cacheData: BehaviorSubject<TwoFactorAuthEmailComponentCache | null>;
+  let mockSignal: any;
+
+  beforeEach(() => {
+    mockViewCacheService = mock<ViewCacheService>();
+    mockConfigService = mock<ConfigService>();
+    cacheData = new BehaviorSubject<TwoFactorAuthEmailComponentCache | null>(null);
+    mockSignal = jest.fn(() => cacheData.getValue());
+    mockSignal.set = jest.fn((value: TwoFactorAuthEmailComponentCache | null) =>
+      cacheData.next(value),
+    );
+    mockViewCacheService.signal.mockReturnValue(mockSignal);
+
+    TestBed.configureTestingModule({
+      providers: [
+        TwoFactorAuthEmailComponentCacheService,
+        { provide: ViewCacheService, useValue: mockViewCacheService },
+        { provide: ConfigService, useValue: mockConfigService },
+      ],
+    });
+
+    service = TestBed.inject(TwoFactorAuthEmailComponentCacheService);
+  });
+
+  it("creates the service", () => {
+    expect(service).toBeTruthy();
+  });
+
+  describe("init", () => {
+    it("sets featureEnabled to true when flag is enabled", async () => {
+      mockConfigService.getFeatureFlag.mockResolvedValue(true);
+
+      await service.init();
+
+      expect(mockConfigService.getFeatureFlag).toHaveBeenCalledWith(
+        FeatureFlag.PM9115_TwoFactorExtensionDataPersistence,
+      );
+
+      service.cacheData({ emailSent: true });
+      expect(mockSignal.set).toHaveBeenCalled();
+    });
+
+    it("sets featureEnabled to false when flag is disabled", async () => {
+      mockConfigService.getFeatureFlag.mockResolvedValue(false);
+
+      await service.init();
+
+      expect(mockConfigService.getFeatureFlag).toHaveBeenCalledWith(
+        FeatureFlag.PM9115_TwoFactorExtensionDataPersistence,
+      );
+
+      service.cacheData({ emailSent: true });
+      expect(mockSignal.set).not.toHaveBeenCalled();
+    });
+  });
+
+  describe("cacheData", () => {
+    beforeEach(async () => {
+      mockConfigService.getFeatureFlag.mockResolvedValue(true);
+      await service.init();
+    });
+
+    it("caches email sent state when feature is enabled", () => {
+      service.cacheData({ emailSent: true });
+
+      expect(mockSignal.set).toHaveBeenCalledWith({
+        emailSent: true,
+      });
+    });
+
+    it("does not cache data when feature is disabled", async () => {
+      mockConfigService.getFeatureFlag.mockResolvedValue(false);
+      await service.init();
+
+      service.cacheData({ emailSent: true });
+
+      expect(mockSignal.set).not.toHaveBeenCalled();
+    });
+  });
+
+  describe("clearCachedData", () => {
+    beforeEach(async () => {
+      mockConfigService.getFeatureFlag.mockResolvedValue(true);
+      await service.init();
+    });
+
+    it("clears cached data when feature is enabled", () => {
+      service.clearCachedData();
+
+      expect(mockSignal.set).toHaveBeenCalledWith(null);
+    });
+
+    it("does not clear cached data when feature is disabled", async () => {
+      mockConfigService.getFeatureFlag.mockResolvedValue(false);
+      await service.init();
+
+      service.clearCachedData();
+
+      expect(mockSignal.set).not.toHaveBeenCalled();
+    });
+  });
+
+  describe("getCachedData", () => {
+    beforeEach(async () => {
+      mockConfigService.getFeatureFlag.mockResolvedValue(true);
+      await service.init();
+    });
+
+    it("returns cached data when feature is enabled", () => {
+      const testData = new TwoFactorAuthEmailComponentCache();
+      testData.emailSent = true;
+      cacheData.next(testData);
+
+      const result = service.getCachedData();
+
+      expect(result).toEqual(testData);
+      expect(mockSignal).toHaveBeenCalled();
+    });
+
+    it("returns null when feature is disabled", async () => {
+      mockConfigService.getFeatureFlag.mockResolvedValue(false);
+      await service.init();
+
+      const result = service.getCachedData();
+
+      expect(result).toBeNull();
+      expect(mockSignal).not.toHaveBeenCalled();
+    });
+  });
+});

libs/auth/src/angular/two-factor-auth/child-components/two-factor-auth-email/two-factor-auth-email-component-cache.service.spec.ts

@@ -0,0 +1,165 @@
+import { TestBed } from "@angular/core/testing";
+import { mock, MockProxy } from "jest-mock-extended";
+import { BehaviorSubject } from "rxjs";
+
+import { ViewCacheService } from "@bitwarden/angular/platform/abstractions/view-cache.service";
+import { FeatureFlag } from "@bitwarden/common/enums/feature-flag.enum";
+import { ConfigService } from "@bitwarden/common/platform/abstractions/config/config.service";
+
+import {
+  TwoFactorAuthEmailComponentCache,
+  TwoFactorAuthEmailComponentCacheService,
+} from "./two-factor-auth-email-component-cache.service";
+
+describe("TwoFactorAuthEmailComponentCache", () => {
+  describe("fromJSON", () => {
+    it("returns null when input is null", () => {
+      const result = TwoFactorAuthEmailComponentCache.fromJSON(null as any);
+      expect(result).toBeNull();
+    });
+
+    it("creates a TwoFactorAuthEmailCache instance from valid JSON", () => {
+      const jsonData = { emailSent: true };
+      const result = TwoFactorAuthEmailComponentCache.fromJSON(jsonData);
+
+      expect(result).not.toBeNull();
+      expect(result).toBeInstanceOf(TwoFactorAuthEmailComponentCache);
+      expect(result?.emailSent).toBe(true);
+    });
+  });
+});
+
+describe("TwoFactorAuthEmailComponentCacheService", () => {
+  let service: TwoFactorAuthEmailComponentCacheService;
+  let mockViewCacheService: MockProxy<ViewCacheService>;
+  let mockConfigService: MockProxy<ConfigService>;
+  let cacheData: BehaviorSubject<TwoFactorAuthEmailComponentCache | null>;
+  let mockSignal: any;
+
+  beforeEach(() => {
+    mockViewCacheService = mock<ViewCacheService>();
+    mockConfigService = mock<ConfigService>();
+    cacheData = new BehaviorSubject<TwoFactorAuthEmailComponentCache | null>(null);
+    mockSignal = jest.fn(() => cacheData.getValue());
+    mockSignal.set = jest.fn((value: TwoFactorAuthEmailComponentCache | null) =>
+      cacheData.next(value),
+    );
+    mockViewCacheService.signal.mockReturnValue(mockSignal);
+
+    TestBed.configureTestingModule({
+      providers: [
+        TwoFactorAuthEmailComponentCacheService,
+        { provide: ViewCacheService, useValue: mockViewCacheService },
+        { provide: ConfigService, useValue: mockConfigService },
+      ],
+    });
+
+    service = TestBed.inject(TwoFactorAuthEmailComponentCacheService);
+  });
+
+  it("creates the service", () => {
+    expect(service).toBeTruthy();
+  });
+
+  describe("init", () => {
+    it("sets featureEnabled to true when flag is enabled", async () => {
+      mockConfigService.getFeatureFlag.mockResolvedValue(true);
+
+      await service.init();
+
+      expect(mockConfigService.getFeatureFlag).toHaveBeenCalledWith(
+        FeatureFlag.PM9115_TwoFactorExtensionDataPersistence,
+      );
+
+      service.cacheData({ emailSent: true });
+      expect(mockSignal.set).toHaveBeenCalled();
+    });
+
+    it("sets featureEnabled to false when flag is disabled", async () => {
+      mockConfigService.getFeatureFlag.mockResolvedValue(false);
+
+      await service.init();
+
+      expect(mockConfigService.getFeatureFlag).toHaveBeenCalledWith(
+        FeatureFlag.PM9115_TwoFactorExtensionDataPersistence,
+      );
+
+      service.cacheData({ emailSent: true });
+      expect(mockSignal.set).not.toHaveBeenCalled();
+    });
+  });
+
+  describe("cacheData", () => {
+    beforeEach(async () => {
+      mockConfigService.getFeatureFlag.mockResolvedValue(true);
+      await service.init();
+    });
+
+    it("caches email sent state when feature is enabled", () => {
+      service.cacheData({ emailSent: true });
+
+      expect(mockSignal.set).toHaveBeenCalledWith({
+        emailSent: true,
+      });
+    });
+
+    it("does not cache data when feature is disabled", async () => {
+      mockConfigService.getFeatureFlag.mockResolvedValue(false);
+      await service.init();
+
+      service.cacheData({ emailSent: true });
+
+      expect(mockSignal.set).not.toHaveBeenCalled();
+    });
+  });
+
+  describe("clearCachedData", () => {
+    beforeEach(async () => {
+      mockConfigService.getFeatureFlag.mockResolvedValue(true);
+      await service.init();
+    });
+
+    it("clears cached data when feature is enabled", () => {
+      service.clearCachedData();
+
+      expect(mockSignal.set).toHaveBeenCalledWith(null);
+    });
+
+    it("does not clear cached data when feature is disabled", async () => {
+      mockConfigService.getFeatureFlag.mockResolvedValue(false);
+      await service.init();
+
+      service.clearCachedData();
+
+      expect(mockSignal.set).not.toHaveBeenCalled();
+    });
+  });
+
+  describe("getCachedData", () => {
+    beforeEach(async () => {
+      mockConfigService.getFeatureFlag.mockResolvedValue(true);
+      await service.init();
+    });
+
+    it("returns cached data when feature is enabled", () => {
+      const testData = new TwoFactorAuthEmailComponentCache();
+      testData.emailSent = true;
+      cacheData.next(testData);
+
+      const result = service.getCachedData();
+
+      expect(result).toEqual(testData);
+      expect(mockSignal).toHaveBeenCalled();
+    });
+
+    it("returns null when feature is disabled", async () => {
+      mockConfigService.getFeatureFlag.mockResolvedValue(false);
+      await service.init();
+
+      const result = service.getCachedData();
+
+      expect(result).toBeNull();
+      expect(mockSignal).not.toHaveBeenCalled();
+    });
+  });
+});

libs/auth/src/angular/two-factor-auth/child-components/two-factor-auth-email/two-factor-auth-email-component-cache.service.ts

@@ -0,0 +1,95 @@
+import { inject, Injectable, WritableSignal } from "@angular/core";
+import { Jsonify } from "type-fest";
+
+import { ViewCacheService } from "@bitwarden/angular/platform/abstractions/view-cache.service";
+import { FeatureFlag } from "@bitwarden/common/enums/feature-flag.enum";
+import { ConfigService } from "@bitwarden/common/platform/abstractions/config/config.service";
+
+/**
+ * The key for the email two factor auth component cache.
+ */
+export const TWO_FACTOR_AUTH_EMAIL_COMPONENT_CACHE_KEY = "two-factor-auth-email-component-cache";
+
+/**
+ * Cache model for the email two factor auth component.
+ */
+export class TwoFactorAuthEmailComponentCache {
+  emailSent: boolean = false;
+
+  static fromJSON(
+    obj: Partial<Jsonify<TwoFactorAuthEmailComponentCache>>,
+  ): TwoFactorAuthEmailComponentCache | null {
+    // Return null if the cache is empty
+    if (obj == null) {
+      return null;
+    }
+
+    return Object.assign(new TwoFactorAuthEmailComponentCache(), obj);
+  }
+}
+
+/**
+ * Cache service for the two factor auth email component.
+ */
+@Injectable()
+export class TwoFactorAuthEmailComponentCacheService {
+  private viewCacheService: ViewCacheService = inject(ViewCacheService);
+  private configService: ConfigService = inject(ConfigService);
+
+  /** True when the feature flag is enabled */
+  private featureEnabled: boolean = false;
+
+  /**
+   * Signal for the cached email state.
+   */
+  private emailCache: WritableSignal<TwoFactorAuthEmailComponentCache | null> =
+    this.viewCacheService.signal<TwoFactorAuthEmailComponentCache | null>({
+      key: TWO_FACTOR_AUTH_EMAIL_COMPONENT_CACHE_KEY,
+      initialValue: null,
+      deserializer: TwoFactorAuthEmailComponentCache.fromJSON,
+    });
+
+  /**
+   * Must be called once before interacting with the cached data.
+   */
+  async init() {
+    this.featureEnabled = await this.configService.getFeatureFlag(
+      FeatureFlag.PM9115_TwoFactorExtensionDataPersistence,
+    );
+  }
+
+  /**
+   * Cache the email sent state.
+   */
+  cacheData(data: { emailSent: boolean }): void {
+    if (!this.featureEnabled) {
+      return;
+    }
+
+    this.emailCache.set({
+      emailSent: data.emailSent,
+    } as TwoFactorAuthEmailComponentCache);
+  }
+
+  /**
+   * Clear the cached email data.
+   */
+  clearCachedData(): void {
+    if (!this.featureEnabled) {
+      return;
+    }
+
+    this.emailCache.set(null);
+  }
+
+  /**
+   * Get whether the email has been sent.
+   */
+  getCachedData(): TwoFactorAuthEmailComponentCache | null {
+    if (!this.featureEnabled) {
+      return null;
+    }
+
+    return this.emailCache();
+  }
+}

libs/auth/src/angular/two-factor-auth/child-components/two-factor-auth-email/two-factor-auth-email.component.html

@@ -1,6 +1,13 @@
 <bit-form-field class="!tw-mb-0">
   <bit-label>{{ "verificationCode" | i18n }}</bit-label>
-  <input bitInput type="text" appAutofocus appInputVerbatim [formControl]="tokenFormControl" />
+  <input
+    bitInput
+    type="text"
+    appAutofocus
+    appInputVerbatim
+    [formControl]="tokenFormControl"
+    (keyup)="onTokenChange($event)"
+  />
 </bit-form-field>
 
 <div class="tw-mb-4">

libs/auth/src/angular/two-factor-auth/child-components/two-factor-auth-email/two-factor-auth-email.component.ts

@@ -1,5 +1,5 @@
 import { CommonModule } from "@angular/common";
-import { Component, Input, OnInit } from "@angular/core";
+import { Component, Input, OnInit, Output, EventEmitter } from "@angular/core";
 import { ReactiveFormsModule, FormsModule, FormControl } from "@angular/forms";
 
 import { JslibModule } from "@bitwarden/angular/jslib.module";
@@ -22,6 +22,7 @@ import {
   ToastService,
 } from "@bitwarden/components";
 
+import { TwoFactorAuthEmailComponentCacheService } from "./two-factor-auth-email-component-cache.service";
 import { TwoFactorAuthEmailComponentService } from "./two-factor-auth-email-component.service";
 
 @Component({
@@ -40,14 +41,20 @@ import { TwoFactorAuthEmailComponentService } from "./two-factor-auth-email-comp
     AsyncActionsModule,
     FormsModule,
   ],
-  providers: [],
+  providers: [
+    {
+      provide: TwoFactorAuthEmailComponentCacheService,
+      useClass: TwoFactorAuthEmailComponentCacheService,
+    },
+  ],
 })
 export class TwoFactorAuthEmailComponent implements OnInit {
   @Input({ required: true }) tokenFormControl: FormControl | undefined = undefined;
+  @Output() tokenChange = new EventEmitter<{ token: string }>();
 
   twoFactorEmail: string | undefined = undefined;
-  emailPromise: Promise<any> | undefined = undefined;
-  tokenValue: string = "";
+  emailPromise: Promise<any> | undefined;
+  emailSent = false;
 
   constructor(
     protected i18nService: I18nService,
@@ -59,14 +66,22 @@ export class TwoFactorAuthEmailComponent implements OnInit {
     protected appIdService: AppIdService,
     private toastService: ToastService,
     private twoFactorAuthEmailComponentService: TwoFactorAuthEmailComponentService,
+    private cacheService: TwoFactorAuthEmailComponentCacheService,
   ) {}
 
   async ngOnInit(): Promise<void> {
     await this.twoFactorAuthEmailComponentService.openPopoutIfApprovedForEmail2fa?.();
+    await this.cacheService.init();
+
+    // Check if email was already sent
+    const cachedData = this.cacheService.getCachedData();
+    if (cachedData?.emailSent) {
+      this.emailSent = true;
+    }
 
     const providers = await this.twoFactorService.getProviders();
 
-    if (!providers) {
+    if (!providers || providers.size === 0) {
       throw new Error("User has no 2FA Providers");
     }
 
@@ -78,11 +93,20 @@ export class TwoFactorAuthEmailComponent implements OnInit {
 
     this.twoFactorEmail = email2faProviderData.Email;
 
-    if (providers.size > 1) {
+    if (!this.emailSent) {
       await this.sendEmail(false);
     }
   }
 
+  /**
+   * Emits the token value to the parent component
+   * @param event - The event object from the input field
+   */
+  onTokenChange(event: Event) {
+    const tokenValue = (event.target as HTMLInputElement).value || "";
+    this.tokenChange.emit({ token: tokenValue });
+  }
+
   async sendEmail(doToast: boolean) {
     if (this.emailPromise !== undefined) {
       return;
@@ -113,6 +137,10 @@ export class TwoFactorAuthEmailComponent implements OnInit {
       request.authRequestId = (await this.loginStrategyService.getAuthRequestId()) ?? "";
       this.emailPromise = this.apiService.postTwoFactorEmail(request);
       await this.emailPromise;
+
+      this.emailSent = true;
+      this.cacheService.cacheData({ emailSent: this.emailSent });
+
       if (doToast) {
         this.toastService.showToast({
           variant: "success",

libs/auth/src/angular/two-factor-auth/two-factor-auth-component-cache.service.spec.ts

@@ -0,0 +1,191 @@
+import { TestBed } from "@angular/core/testing";
+import { mock, MockProxy } from "jest-mock-extended";
+import { BehaviorSubject } from "rxjs";
+
+import { ViewCacheService } from "@bitwarden/angular/platform/abstractions/view-cache.service";
+import { TwoFactorProviderType } from "@bitwarden/common/auth/enums/two-factor-provider-type";
+import { FeatureFlag } from "@bitwarden/common/enums/feature-flag.enum";
+import { ConfigService } from "@bitwarden/common/platform/abstractions/config/config.service";
+
+import {
+  TwoFactorAuthComponentCache,
+  TwoFactorAuthComponentCacheService,
+  TwoFactorAuthComponentData,
+} from "./two-factor-auth-component-cache.service";
+
+describe("TwoFactorAuthCache", () => {
+  describe("fromJSON", () => {
+    it("returns null when input is null", () => {
+      const result = TwoFactorAuthComponentCache.fromJSON(null as any);
+      expect(result).toBeNull();
+    });
+
+    it("creates a TwoFactorAuthCache instance from valid JSON", () => {
+      const jsonData = {
+        token: "123456",
+        remember: true,
+        selectedProviderType: TwoFactorProviderType.Email,
+      };
+      const result = TwoFactorAuthComponentCache.fromJSON(jsonData as any);
+
+      expect(result).not.toBeNull();
+      expect(result).toBeInstanceOf(TwoFactorAuthComponentCache);
+      expect(result?.token).toBe("123456");
+      expect(result?.remember).toBe(true);
+      expect(result?.selectedProviderType).toBe(TwoFactorProviderType.Email);
+    });
+  });
+});
+
+describe("TwoFactorAuthComponentCacheService", () => {
+  let service: TwoFactorAuthComponentCacheService;
+  let mockViewCacheService: MockProxy<ViewCacheService>;
+  let mockConfigService: MockProxy<ConfigService>;
+  let cacheData: BehaviorSubject<TwoFactorAuthComponentCache | null>;
+  let mockSignal: any;
+
+  beforeEach(() => {
+    mockViewCacheService = mock<ViewCacheService>();
+    mockConfigService = mock<ConfigService>();
+    cacheData = new BehaviorSubject<TwoFactorAuthComponentCache | null>(null);
+    mockSignal = jest.fn(() => cacheData.getValue());
+    mockSignal.set = jest.fn((value: TwoFactorAuthComponentCache | null) => cacheData.next(value));
+    mockViewCacheService.signal.mockReturnValue(mockSignal);
+
+    TestBed.configureTestingModule({
+      providers: [
+        TwoFactorAuthComponentCacheService,
+        { provide: ViewCacheService, useValue: mockViewCacheService },
+        { provide: ConfigService, useValue: mockConfigService },
+      ],
+    });
+
+    service = TestBed.inject(TwoFactorAuthComponentCacheService);
+  });
+
+  it("creates the service", () => {
+    expect(service).toBeTruthy();
+  });
+
+  describe("init", () => {
+    it("sets featureEnabled to true when flag is enabled", async () => {
+      mockConfigService.getFeatureFlag.mockResolvedValue(true);
+
+      await service.init();
+
+      expect(mockConfigService.getFeatureFlag).toHaveBeenCalledWith(
+        FeatureFlag.PM9115_TwoFactorExtensionDataPersistence,
+      );
+
+      service.cacheData({ token: "123456" });
+      expect(mockSignal.set).toHaveBeenCalled();
+    });
+
+    it("sets featureEnabled to false when flag is disabled", async () => {
+      mockConfigService.getFeatureFlag.mockResolvedValue(false);
+
+      await service.init();
+
+      expect(mockConfigService.getFeatureFlag).toHaveBeenCalledWith(
+        FeatureFlag.PM9115_TwoFactorExtensionDataPersistence,
+      );
+
+      service.cacheData({ token: "123456" });
+      expect(mockSignal.set).not.toHaveBeenCalled();
+    });
+  });
+
+  describe("cacheData", () => {
+    beforeEach(async () => {
+      mockConfigService.getFeatureFlag.mockResolvedValue(true);
+      await service.init();
+    });
+
+    it("caches complete data when feature is enabled", () => {
+      const testData: TwoFactorAuthComponentData = {
+        token: "123456",
+        remember: true,
+        selectedProviderType: TwoFactorProviderType.Email,
+      };
+
+      service.cacheData(testData);
+
+      expect(mockSignal.set).toHaveBeenCalledWith({
+        token: "123456",
+        remember: true,
+        selectedProviderType: TwoFactorProviderType.Email,
+      });
+    });
+
+    it("caches partial data when feature is enabled", () => {
+      service.cacheData({ token: "123456" });
+
+      expect(mockSignal.set).toHaveBeenCalledWith({
+        token: "123456",
+        remember: undefined,
+        selectedProviderType: undefined,
+      });
+    });
+
+    it("does not cache data when feature is disabled", async () => {
+      mockConfigService.getFeatureFlag.mockResolvedValue(false);
+      await service.init();
+
+      service.cacheData({ token: "123456" });
+
+      expect(mockSignal.set).not.toHaveBeenCalled();
+    });
+  });
+
+  describe("clearCachedData", () => {
+    beforeEach(async () => {
+      mockConfigService.getFeatureFlag.mockResolvedValue(true);
+      await service.init();
+    });
+
+    it("clears cached data when feature is enabled", () => {
+      service.clearCachedData();
+
+      expect(mockSignal.set).toHaveBeenCalledWith(null);
+    });
+
+    it("does not clear cached data when feature is disabled", async () => {
+      mockConfigService.getFeatureFlag.mockResolvedValue(false);
+      await service.init();
+
+      service.clearCachedData();
+
+      expect(mockSignal.set).not.toHaveBeenCalled();
+    });
+  });
+
+  describe("getCachedData", () => {
+    beforeEach(async () => {
+      mockConfigService.getFeatureFlag.mockResolvedValue(true);
+      await service.init();
+    });
+
+    it("returns cached data when feature is enabled", () => {
+      const testData = new TwoFactorAuthComponentCache();
+      testData.token = "123456";
+      testData.remember = true;
+      testData.selectedProviderType = TwoFactorProviderType.Email;
+      cacheData.next(testData);
+
+      const result = service.getCachedData();
+
+      expect(result).toEqual(testData);
+      expect(mockSignal).toHaveBeenCalled();
+    });
+
+    it("returns null when feature is disabled", async () => {
+      mockConfigService.getFeatureFlag.mockResolvedValue(false);
+      await service.init();
+
+      const result = service.getCachedData();
+
+      expect(result).toBeNull();
+      expect(mockSignal).not.toHaveBeenCalled();
+    });
+  });
+});

libs/auth/src/angular/two-factor-auth/two-factor-auth-component-cache.service.ts

@@ -0,0 +1,105 @@
+import { inject, Injectable, WritableSignal } from "@angular/core";
+import { Jsonify } from "type-fest";
+
+import { ViewCacheService } from "@bitwarden/angular/platform/abstractions/view-cache.service";
+import { TwoFactorProviderType } from "@bitwarden/common/auth/enums/two-factor-provider-type";
+import { FeatureFlag } from "@bitwarden/common/enums/feature-flag.enum";
+import { ConfigService } from "@bitwarden/common/platform/abstractions/config/config.service";
+
+const TWO_FACTOR_AUTH_COMPONENT_CACHE_KEY = "two-factor-auth-component-cache";
+
+/**
+ * Cache model for the two factor authentication data.
+ */
+export class TwoFactorAuthComponentCache {
+  token: string | undefined = undefined;
+  remember: boolean | undefined = undefined;
+  selectedProviderType: TwoFactorProviderType | undefined = undefined;
+
+  static fromJSON(
+    obj: Partial<Jsonify<TwoFactorAuthComponentCache>>,
+  ): TwoFactorAuthComponentCache | null {
+    // Return null if the cache is empty
+    if (obj == null) {
+      return null;
+    }
+
+    return Object.assign(new TwoFactorAuthComponentCache(), obj);
+  }
+}
+
+export interface TwoFactorAuthComponentData {
+  token?: string;
+  remember?: boolean;
+  selectedProviderType?: TwoFactorProviderType;
+}
+
+/**
+ * Cache service used for the two factor auth component.
+ */
+@Injectable()
+export class TwoFactorAuthComponentCacheService {
+  private viewCacheService: ViewCacheService = inject(ViewCacheService);
+  private configService: ConfigService = inject(ConfigService);
+
+  /** True when the `PM9115_TwoFactorExtensionDataPersistence` flag is enabled */
+  private featureEnabled: boolean = false;
+
+  /**
+   * Signal for the cached TwoFactorAuthData.
+   */
+  private twoFactorAuthComponentCache: WritableSignal<TwoFactorAuthComponentCache | null> =
+    this.viewCacheService.signal<TwoFactorAuthComponentCache | null>({
+      key: TWO_FACTOR_AUTH_COMPONENT_CACHE_KEY,
+      initialValue: null,
+      deserializer: TwoFactorAuthComponentCache.fromJSON,
+    });
+
+  constructor() {}
+
+  /**
+   * Must be called once before interacting with the cached data.
+   */
+  async init() {
+    this.featureEnabled = await this.configService.getFeatureFlag(
+      FeatureFlag.PM9115_TwoFactorExtensionDataPersistence,
+    );
+  }
+
+  /**
+   * Update the cache with the new TwoFactorAuthData.
+   */
+  cacheData(data: TwoFactorAuthComponentData): void {
+    if (!this.featureEnabled) {
+      return;
+    }
+
+    this.twoFactorAuthComponentCache.set({
+      token: data.token,
+      remember: data.remember,
+      selectedProviderType: data.selectedProviderType,
+    } as TwoFactorAuthComponentCache);
+  }
+
+  /**
+   * Clears the cached TwoFactorAuthData.
+   */
+  clearCachedData(): void {
+    if (!this.featureEnabled) {
+      return;
+    }
+
+    this.twoFactorAuthComponentCache.set(null);
+  }
+
+  /**
+   * Returns the cached TwoFactorAuthData (when available).
+   */
+  getCachedData(): TwoFactorAuthComponentCache | null {
+    if (!this.featureEnabled) {
+      return null;
+    }
+
+    return this.twoFactorAuthComponentCache();
+  }
+}

libs/auth/src/angular/two-factor-auth/two-factor-auth.component.html

@@ -13,11 +13,13 @@
   >
     <app-two-factor-auth-email
       [tokenFormControl]="tokenFormControl"
+      (tokenChange)="saveFormDataWithPartialData($event)"
       *ngIf="selectedProviderType === providerType.Email"
     />
 
     <app-two-factor-auth-authenticator
       [tokenFormControl]="tokenFormControl"
+      (tokenChange)="saveFormDataWithPartialData($event)"
       *ngIf="selectedProviderType === providerType.Authenticator"
     />
     <app-two-factor-auth-yubikey
@@ -36,7 +38,7 @@
     />
     <bit-form-control *ngIf="!hideRememberMe()">
       <bit-label>{{ "dontAskAgainOnThisDeviceFor30Days" | i18n }}</bit-label>
-      <input type="checkbox" bitCheckbox formControlName="remember" />
+      <input type="checkbox" bitCheckbox formControlName="remember" (change)="onRememberChange()" />
     </bit-form-control>
 
     <app-two-factor-auth-webauthn

libs/auth/src/angular/two-factor-auth/two-factor-auth.component.spec.ts

@@ -1,8 +1,6 @@
-// FIXME: Update this file to be type safe and remove this and next line
-// @ts-strict-ignore
 import { Component } from "@angular/core";
 import { ComponentFixture, TestBed } from "@angular/core/testing";
-import { ActivatedRoute, convertToParamMap, Router } from "@angular/router";
+import { ActivatedRoute, Router, convertToParamMap } from "@angular/router";
 import { mock, MockProxy } from "jest-mock-extended";
 import { BehaviorSubject } from "rxjs";
 
@@ -38,6 +36,7 @@ import { DialogService, ToastService } from "@bitwarden/components";
 
 import { AnonLayoutWrapperDataService } from "../anon-layout/anon-layout-wrapper-data.service";
 
+import { TwoFactorAuthComponentCacheService } from "./two-factor-auth-component-cache.service";
 import { TwoFactorAuthComponentService } from "./two-factor-auth-component.service";
 import { TwoFactorAuthComponent } from "./two-factor-auth.component";
 
@@ -72,6 +71,7 @@ describe("TwoFactorAuthComponent", () => {
   let anonLayoutWrapperDataService: MockProxy<AnonLayoutWrapperDataService>;
   let mockEnvService: MockProxy<EnvironmentService>;
   let mockLoginSuccessHandlerService: MockProxy<LoginSuccessHandlerService>;
+  let mockTwoFactorAuthCompCacheService: MockProxy<TwoFactorAuthComponentCacheService>;
 
   let mockUserDecryptionOpts: {
     noMasterPassword: UserDecryptionOptions;
@@ -112,6 +112,10 @@ describe("TwoFactorAuthComponent", () => {
 
     anonLayoutWrapperDataService = mock<AnonLayoutWrapperDataService>();
 
+    mockTwoFactorAuthCompCacheService = mock<TwoFactorAuthComponentCacheService>();
+    mockTwoFactorAuthCompCacheService.getCachedData.mockReturnValue(null);
+    mockTwoFactorAuthCompCacheService.init.mockResolvedValue();
+
     mockUserDecryptionOpts = {
       noMasterPassword: new UserDecryptionOptions({
         hasMasterPassword: false,
@@ -155,7 +159,9 @@ describe("TwoFactorAuthComponent", () => {
       }),
     };
 
-    selectedUserDecryptionOptions = new BehaviorSubject<UserDecryptionOptions>(undefined);
+    selectedUserDecryptionOptions = new BehaviorSubject<UserDecryptionOptions>(
+      mockUserDecryptionOpts.withMasterPassword,
+    );
     mockUserDecryptionOptionsService.userDecryptionOptions$ = selectedUserDecryptionOptions;
 
     TestBed.configureTestingModule({
@@ -194,6 +200,10 @@ describe("TwoFactorAuthComponent", () => {
         { provide: EnvironmentService, useValue: mockEnvService },
         { provide: AnonLayoutWrapperDataService, useValue: anonLayoutWrapperDataService },
         { provide: LoginSuccessHandlerService, useValue: mockLoginSuccessHandlerService },
+        {
+          provide: TwoFactorAuthComponentCacheService,
+          useValue: mockTwoFactorAuthCompCacheService,
+        },
       ],
     });
 

libs/auth/src/angular/two-factor-auth/two-factor-auth.component.ts

@@ -60,6 +60,10 @@ import { TwoFactorAuthDuoComponent } from "./child-components/two-factor-auth-du
 import { TwoFactorAuthEmailComponent } from "./child-components/two-factor-auth-email/two-factor-auth-email.component";
 import { TwoFactorAuthWebAuthnComponent } from "./child-components/two-factor-auth-webauthn/two-factor-auth-webauthn.component";
 import { TwoFactorAuthYubikeyComponent } from "./child-components/two-factor-auth-yubikey.component";
+import {
+  TwoFactorAuthComponentCacheService,
+  TwoFactorAuthComponentData,
+} from "./two-factor-auth-component-cache.service";
 import {
   DuoLaunchAction,
   LegacyKeyMigrationAction,
@@ -90,7 +94,11 @@ import {
     TwoFactorAuthYubikeyComponent,
     TwoFactorAuthWebAuthnComponent,
   ],
-  providers: [],
+  providers: [
+    {
+      provide: TwoFactorAuthComponentCacheService,
+    },
+  ],
 })
 export class TwoFactorAuthComponent implements OnInit, OnDestroy {
   @ViewChild("continueButton", { read: ElementRef, static: false }) continueButton:
@@ -160,6 +168,7 @@ export class TwoFactorAuthComponent implements OnInit, OnDestroy {
     private anonLayoutWrapperDataService: AnonLayoutWrapperDataService,
     private environmentService: EnvironmentService,
     private loginSuccessHandlerService: LoginSuccessHandlerService,
+    private twoFactorAuthComponentCacheService: TwoFactorAuthComponentCacheService,
   ) {}
 
   async ngOnInit() {
@@ -168,7 +177,33 @@ export class TwoFactorAuthComponent implements OnInit, OnDestroy {
 
     this.listenForAuthnSessionTimeout();
 
-    await this.setSelected2faProviderType();
+    // Initialize the cache
+    await this.twoFactorAuthComponentCacheService.init();
+
+    // Load cached form data if available
+    let loadedCachedProviderType = false;
+    const cachedData = this.twoFactorAuthComponentCacheService.getCachedData();
+    if (cachedData) {
+      if (cachedData.token) {
+        this.form.patchValue({ token: cachedData.token });
+      }
+      if (cachedData.remember !== undefined) {
+        this.form.patchValue({ remember: cachedData.remember });
+      }
+      if (cachedData.selectedProviderType !== undefined) {
+        this.selectedProviderType = cachedData.selectedProviderType;
+        loadedCachedProviderType = true;
+      }
+    }
+
+    // If we don't have a cached provider type, set it to the default and cache it
+    if (!loadedCachedProviderType) {
+      this.selectedProviderType = await this.initializeSelected2faProviderType();
+      this.twoFactorAuthComponentCacheService.cacheData({
+        selectedProviderType: this.selectedProviderType,
+      });
+    }
+
     await this.set2faProvidersAndData();
     await this.setAnonLayoutDataByTwoFactorProviderType();
 
@@ -181,7 +216,29 @@ export class TwoFactorAuthComponent implements OnInit, OnDestroy {
     this.loading = false;
   }
 
-  private async setSelected2faProviderType() {
+  /**
+   * Save specific form data fields to the cache
+   */
+  async saveFormDataWithPartialData(data: Partial<TwoFactorAuthComponentData>) {
+    // Get current cached data
+    const currentData = this.twoFactorAuthComponentCacheService.getCachedData();
+
+    this.twoFactorAuthComponentCacheService.cacheData({
+      token: data?.token ?? currentData?.token ?? "",
+      remember: data?.remember ?? currentData?.remember ?? false,
+      selectedProviderType: data?.selectedProviderType ?? currentData?.selectedProviderType,
+    });
+  }
+
+  /**
+   * Save the remember value to the cache when the checkbox is checked or unchecked
+   */
+  async onRememberChange() {
+    const rememberValue = !!this.rememberFormControl.value;
+    await this.saveFormDataWithPartialData({ remember: rememberValue });
+  }
+
+  private async initializeSelected2faProviderType(): Promise<TwoFactorProviderType> {
     const webAuthnSupported = this.platformUtilsService.supportsWebAuthn(this.win);
 
     if (
@@ -190,18 +247,19 @@ export class TwoFactorAuthComponent implements OnInit, OnDestroy {
     ) {
       const webAuthn2faResponse = this.activatedRoute.snapshot.paramMap.get("webAuthnResponse");
       if (webAuthn2faResponse) {
-        this.selectedProviderType = TwoFactorProviderType.WebAuthn;
-        return;
+        return TwoFactorProviderType.WebAuthn;
       }
     }
 
-    this.selectedProviderType = await this.twoFactorService.getDefaultProvider(webAuthnSupported);
+    return await this.twoFactorService.getDefaultProvider(webAuthnSupported);
   }
 
   private async set2faProvidersAndData() {
     this.twoFactorProviders = await this.twoFactorService.getProviders();
-    const providerData = this.twoFactorProviders?.get(this.selectedProviderType);
-    this.selectedProviderData = providerData;
+    if (this.selectedProviderType !== undefined) {
+      const providerData = this.twoFactorProviders?.get(this.selectedProviderType);
+      this.selectedProviderData = providerData;
+    }
   }
 
   private listenForAuthnSessionTimeout() {
@@ -267,6 +325,13 @@ export class TwoFactorAuthComponent implements OnInit, OnDestroy {
     // In all flows but WebAuthn, the remember value is taken from the form.
     const rememberValue = remember ?? this.rememberFormControl.value ?? false;
 
+    // Cache form data before submitting
+    this.twoFactorAuthComponentCacheService.cacheData({
+      token: tokenValue,
+      remember: rememberValue,
+      selectedProviderType: this.selectedProviderType,
+    });
+
     try {
       this.formPromise = this.loginStrategyService.logInTwoFactor(
         new TokenTwoFactorRequest(this.selectedProviderType, tokenValue, rememberValue),
@@ -274,6 +339,7 @@ export class TwoFactorAuthComponent implements OnInit, OnDestroy {
       );
       const authResult: AuthResult = await this.formPromise;
       this.logService.info("Successfully submitted two factor token");
+
       await this.handleAuthResult(authResult);
     } catch {
       this.logService.error("Error submitting two factor token");
@@ -299,6 +365,13 @@ export class TwoFactorAuthComponent implements OnInit, OnDestroy {
       this.selectedProviderType = response.type;
       await this.setAnonLayoutDataByTwoFactorProviderType();
 
+      // Update the cached provider type when a new one is chosen
+      this.twoFactorAuthComponentCacheService.cacheData({
+        token: "",
+        remember: false,
+        selectedProviderType: response.type,
+      });
+
       this.form.reset();
       this.form.updateValueAndValidity();
     }
@@ -376,6 +449,9 @@ export class TwoFactorAuthComponent implements OnInit, OnDestroy {
   }
 
   private async handleAuthResult(authResult: AuthResult) {
+    // Clear form cache
+    this.twoFactorAuthComponentCacheService.clearCachedData();
+
     if (await this.handleMigrateEncryptionKey(authResult)) {
       return; // stop login process
     }

libs/common/src/enums/feature-flag.enum.ts

@@ -17,6 +17,7 @@ export enum FeatureFlag {
 
   /* Auth */
   PM9112_DeviceApprovalPersistence = "pm-9112-device-approval-persistence",
+  PM9115_TwoFactorExtensionDataPersistence = "pm-9115-two-factor-extension-data-persistence",
 
   /* Autofill */
   BlockBrowserInjectionsByDomain = "block-browser-injections-by-domain",
@@ -116,6 +117,7 @@ export const DefaultFeatureFlagValue = {
 
   /* Auth */
   [FeatureFlag.PM9112_DeviceApprovalPersistence]: FALSE,
+  [FeatureFlag.PM9115_TwoFactorExtensionDataPersistence]: FALSE,
 
   /* Billing */
   [FeatureFlag.TrialPaymentOptional]: FALSE,