feat: enhance MFA retry logic with configurable options (#1441)

apps/cli/src/commands/auth.command.ts

@@ -603,7 +603,7 @@ Examples:
 						error.name === 'ExitPromptError' ||
 						error.message?.includes('force closed')
 					) {
-						ui.displayWarning('\nMFA verification cancelled by user');
+						ui.displayWarning(' MFA verification cancelled by user');
 						throw new AuthenticationError(
 							'MFA verification cancelled',
 							'MFA_VERIFICATION_FAILED'
@@ -612,7 +612,15 @@ Examples:
 					throw error;
 				}
 			},
-			3 // Max attempts
+			{
+				maxAttempts: 3,
+				onInvalidCode: (_attempt: number, remaining: number) => {
+					// Callback invoked when invalid code is entered
+					if (remaining > 0) {
+						ui.displayError(`Invalid MFA code. Please try again.`);
+					}
+				}
+			}
 		);
 
 		// Handle result from core

packages/tm-core/src/modules/auth/managers/auth-manager.spec.ts

@@ -195,7 +195,7 @@ describe('AuthManager - MFA Retry Logic', () => {
 			const result = await authManager.verifyMFAWithRetry(
 				'factor-123',
 				codeProvider,
-				3
+				{ maxAttempts: 3 }
 			);
 
 			expect(result.success).toBe(true);
@@ -232,7 +232,7 @@ describe('AuthManager - MFA Retry Logic', () => {
 			const result = await authManager.verifyMFAWithRetry(
 				'factor-123',
 				codeProvider,
-				3
+				{ maxAttempts: 3 }
 			);
 
 			expect(result.success).toBe(true);
@@ -254,7 +254,7 @@ describe('AuthManager - MFA Retry Logic', () => {
 			const result = await authManager.verifyMFAWithRetry(
 				'factor-123',
 				codeProvider,
-				3
+				{ maxAttempts: 3 }
 			);
 
 			expect(result.success).toBe(false);
@@ -276,7 +276,9 @@ describe('AuthManager - MFA Retry Logic', () => {
 			vi.spyOn(authManager, 'verifyMFA').mockRejectedValue(networkError);
 
 			await expect(
-				authManager.verifyMFAWithRetry('factor-123', codeProvider, 3)
+				authManager.verifyMFAWithRetry('factor-123', codeProvider, {
+					maxAttempts: 3
+				})
 			).rejects.toThrow('Network error');
 
 			// Should not retry on non-INVALID_MFA_CODE errors
@@ -295,7 +297,7 @@ describe('AuthManager - MFA Retry Logic', () => {
 			const result = await authManager.verifyMFAWithRetry(
 				'factor-123',
 				codeProvider,
-				5 // Custom max attempts
+				{ maxAttempts: 5 } // Custom max attempts
 			);
 
 			expect(result.success).toBe(false);
@@ -311,7 +313,7 @@ describe('AuthManager - MFA Retry Logic', () => {
 				new AuthenticationError('Invalid MFA code', 'INVALID_MFA_CODE')
 			);
 
-			// Don't pass maxAttempts - should default to 3
+			// Don't pass options - should default to 3
 			const result = await authManager.verifyMFAWithRetry(
 				'factor-123',
 				codeProvider
@@ -328,24 +330,151 @@ describe('AuthManager - MFA Retry Logic', () => {
 
 			// Test with 0
 			await expect(
-				authManager.verifyMFAWithRetry('factor-123', codeProvider, 0)
+				authManager.verifyMFAWithRetry('factor-123', codeProvider, {
+					maxAttempts: 0
+				})
 			).rejects.toThrow(TypeError);
 
 			await expect(
-				authManager.verifyMFAWithRetry('factor-123', codeProvider, 0)
+				authManager.verifyMFAWithRetry('factor-123', codeProvider, {
+					maxAttempts: 0
+				})
 			).rejects.toThrow('Invalid maxAttempts value: 0. Must be at least 1.');
 
 			// Test with negative
 			await expect(
-				authManager.verifyMFAWithRetry('factor-123', codeProvider, -1)
+				authManager.verifyMFAWithRetry('factor-123', codeProvider, {
+					maxAttempts: -1
+				})
 			).rejects.toThrow(TypeError);
 
 			await expect(
-				authManager.verifyMFAWithRetry('factor-123', codeProvider, -1)
+				authManager.verifyMFAWithRetry('factor-123', codeProvider, {
+					maxAttempts: -1
+				})
 			).rejects.toThrow('Invalid maxAttempts value: -1. Must be at least 1.');
 
 			// Verify code provider was never called
 			expect(codeProvider).not.toHaveBeenCalled();
 		});
+
+		it('should invoke onInvalidCode callback when invalid code is entered', async () => {
+			const authManager = AuthManager.getInstance();
+			const codeProvider = vi.fn(async () => '000000');
+			const onInvalidCode = vi.fn();
+
+			// Mock verification to always fail
+			vi.spyOn(authManager, 'verifyMFA').mockRejectedValue(
+				new AuthenticationError('Invalid MFA code', 'INVALID_MFA_CODE')
+			);
+
+			await authManager.verifyMFAWithRetry('factor-123', codeProvider, {
+				maxAttempts: 3,
+				onInvalidCode
+			});
+
+			// Should be called 3 times (after each failed attempt)
+			expect(onInvalidCode).toHaveBeenCalledTimes(3);
+
+			// Verify callback arguments: (attempt, remaining)
+			expect(onInvalidCode).toHaveBeenNthCalledWith(1, 1, 2); // 1st attempt, 2 remaining
+			expect(onInvalidCode).toHaveBeenNthCalledWith(2, 2, 1); // 2nd attempt, 1 remaining
+			expect(onInvalidCode).toHaveBeenNthCalledWith(3, 3, 0); // 3rd attempt, 0 remaining
+		});
+
+		it('should not invoke onInvalidCode callback on successful verification', async () => {
+			const authManager = AuthManager.getInstance();
+			const codeProvider = vi.fn(async () => '123456');
+			const onInvalidCode = vi.fn();
+
+			// Mock successful verification
+			vi.spyOn(authManager, 'verifyMFA').mockResolvedValue({
+				token: 'test-token',
+				userId: 'test-user',
+				email: 'test@example.com',
+				tokenType: 'standard',
+				savedAt: new Date().toISOString()
+			});
+
+			const result = await authManager.verifyMFAWithRetry(
+				'factor-123',
+				codeProvider,
+				{
+					maxAttempts: 3,
+					onInvalidCode
+				}
+			);
+
+			expect(result.success).toBe(true);
+			expect(onInvalidCode).not.toHaveBeenCalled();
+		});
+
+		it('should work without onInvalidCode callback (backward compatibility)', async () => {
+			const authManager = AuthManager.getInstance();
+			const codeProvider = vi.fn(async () => '123456');
+
+			// Mock successful verification
+			vi.spyOn(authManager, 'verifyMFA').mockResolvedValue({
+				token: 'test-token',
+				userId: 'test-user',
+				email: 'test@example.com',
+				tokenType: 'standard',
+				savedAt: new Date().toISOString()
+			});
+
+			// Call without onInvalidCode - should not throw
+			const result = await authManager.verifyMFAWithRetry(
+				'factor-123',
+				codeProvider,
+				{
+					maxAttempts: 3
+				}
+			);
+
+			expect(result.success).toBe(true);
+		});
+
+		it('should invoke onInvalidCode with correct remaining attempts', async () => {
+			const authManager = AuthManager.getInstance();
+			let attemptCount = 0;
+			const codeProvider = vi.fn(async () => {
+				attemptCount++;
+				return `code-${attemptCount}`;
+			});
+			const onInvalidCode = vi.fn();
+
+			// Fail twice, then succeed
+			vi.spyOn(authManager, 'verifyMFA')
+				.mockRejectedValueOnce(
+					new AuthenticationError('Invalid MFA code', 'INVALID_MFA_CODE')
+				)
+				.mockRejectedValueOnce(
+					new AuthenticationError('Invalid MFA code', 'INVALID_MFA_CODE')
+				)
+				.mockResolvedValueOnce({
+					token: 'test-token',
+					userId: 'test-user',
+					email: 'test@example.com',
+					tokenType: 'standard',
+					savedAt: new Date().toISOString()
+				});
+
+			const result = await authManager.verifyMFAWithRetry(
+				'factor-123',
+				codeProvider,
+				{
+					maxAttempts: 3,
+					onInvalidCode
+				}
+			);
+
+			expect(result.success).toBe(true);
+			expect(result.attemptsUsed).toBe(3);
+
+			// Verify callback was called for the first two failed attempts
+			expect(onInvalidCode).toHaveBeenCalledTimes(2);
+			expect(onInvalidCode).toHaveBeenNthCalledWith(1, 1, 2); // 1st attempt, 2 remaining
+			expect(onInvalidCode).toHaveBeenNthCalledWith(2, 2, 1); // 2nd attempt, 1 remaining
+		});
 	});
 });

packages/tm-core/src/modules/auth/managers/auth-manager.ts

@@ -130,7 +130,7 @@ export class AuthManager {
 	 *
 	 * @param factorId - MFA factor ID from the MFA_REQUIRED error
 	 * @param codeProvider - Function that prompts for and returns the MFA code
-	 * @param maxAttempts - Maximum number of verification attempts (default: 3)
+	 * @param options - Optional configuration for retry behavior
 	 * @returns Result object with success status, attempts used, and credentials if successful
 	 *
 	 * @example
@@ -138,7 +138,10 @@ export class AuthManager {
 	 * const result = await authManager.verifyMFAWithRetry(
 	 *   factorId,
 	 *   async () => await promptUserForMFACode(),
-	 *   3
+	 *   {
+	 *     maxAttempts: 3,
+	 *     onInvalidCode: (attempt, remaining) => console.log(`Invalid code. ${remaining} attempts remaining.`)
+	 *   }
 	 * );
 	 *
 	 * if (result.success) {
@@ -151,8 +154,14 @@ export class AuthManager {
 	async verifyMFAWithRetry(
 		factorId: string,
 		codeProvider: () => Promise<string>,
-		maxAttempts = 3
+		options?: {
+			maxAttempts?: number;
+			onInvalidCode?: (attempt: number, remaining: number) => void;
+		}
 	): Promise<MFAVerificationResult> {
+		const maxAttempts = options?.maxAttempts ?? 3;
+		const onInvalidCode = options?.onInvalidCode;
+
 		// Guard against invalid maxAttempts values
 		if (maxAttempts < 1) {
 			throw new TypeError(
@@ -175,6 +184,14 @@ export class AuthManager {
 					error instanceof AuthenticationError &&
 					error.code === 'INVALID_MFA_CODE'
 				) {
+					// Calculate remaining attempts
+					const remaining = maxAttempts - attempt;
+
+					// Notify callback of invalid code
+					if (onInvalidCode) {
+						onInvalidCode(attempt, remaining);
+					}
+
 					// If we've exhausted attempts, return failure
 					if (attempt >= maxAttempts) {
 						return {