feat: Add YNAB budget selector to resolve multi-budget users

README.md

@@ -58,7 +58,8 @@ moneyflow --since 2024-06-01  # Fetch from specific date
 
 **First-time Monarch Money setup:** You'll need your 2FA secret key. See the [Monarch Money setup guide](https://moneyflow.dev/guide/monarch).
 
-**First-time YNAB setup:** You'll need a Personal Access Token from your YNAB account settings. See the [YNAB setup guide](https://moneyflow.dev/guide/ynab).
+**First-time YNAB setup:** You'll need a Personal Access Token from your YNAB account settings. If you have multiple
+budgets, you'll be prompted to select one. See the [YNAB setup guide](https://moneyflow.dev/guide/ynab).
 
 ---
 

docs/guide/ynab.md

@@ -69,10 +69,18 @@ Create a NEW password to encrypt your stored credentials:
 !!! info
     Token encrypted with AES-128/PBKDF2 (100k iterations) at `~/.moneyflow/credentials.enc`
 
-### 4. Initial Data Load
+### 4. Budget Selection (if multiple budgets)
 
-**Note:** If you have multiple YNAB budgets, moneyflow will automatically use the first one. Multi-budget selection
-UI is not yet implemented.
+If you have multiple YNAB budgets, moneyflow will show a budget selector screen:
+
+- Use **↑/↓** arrow keys to navigate between budgets
+- Press **Enter** to select the highlighted budget
+- Press **Esc** to cancel
+
+The selected budget will be saved for future sessions, so you won't need to select it again unless you switch
+accounts or delete the account profile.
+
+### 5. Initial Data Load
 
 moneyflow will:
 

moneyflow/account_manager.py

@@ -38,16 +38,20 @@ class Account:
     backend_type: BackendType  # Backend type (monarch, ynab, amazon, demo)
     created_at: str  # ISO timestamp when account was created
     last_used: Optional[str] = None  # ISO timestamp when last accessed
+    budget_id: Optional[str] = None  # For YNAB: the specific budget ID to use
 
     def to_dict(self) -> Dict:
         """Convert to dictionary for JSON serialization."""
-        return {
+        result = {
             "id": self.id,
             "name": self.name,
             "backend_type": self.backend_type,
             "created_at": self.created_at,
             "last_used": self.last_used,
         }
+        if self.budget_id is not None:
+            result["budget_id"] = self.budget_id
+        return result
 
     @staticmethod
     def from_dict(data: Dict) -> "Account":
@@ -58,6 +62,7 @@ def from_dict(data: Dict) -> "Account":
             backend_type=data["backend_type"],
             created_at=data["created_at"],
             last_used=data.get("last_used"),
+            budget_id=data.get("budget_id"),
         )
 
 
@@ -188,7 +193,11 @@ def generate_account_id(self, backend_type: str, account_name: str) -> str:
         return f"{account_id}-{counter}"
 
     def create_account(
-        self, name: str, backend_type: BackendType, account_id: Optional[str] = None
+        self,
+        name: str,
+        backend_type: BackendType,
+        account_id: Optional[str] = None,
+        budget_id: Optional[str] = None,
     ) -> Account:
         """
         Create a new account profile.
@@ -197,6 +206,7 @@ def create_account(
             name: User-friendly display name
             backend_type: Backend type (monarch, ynab, amazon, demo)
             account_id: Optional custom ID (generated if not provided)
+            budget_id: Optional budget ID for YNAB accounts
 
         Returns:
             Created Account object
@@ -226,6 +236,7 @@ def create_account(
             backend_type=backend_type,
             created_at=datetime.now().isoformat(),
             last_used=None,
+            budget_id=budget_id,
         )
 
         # Add to registry and save

moneyflow/app.py

@@ -34,8 +34,9 @@
 from textual.reactive import reactive
 from textual.widgets import DataTable, Footer, Header, LoadingIndicator, Static
 
-from .account_manager import AccountManager
+from .account_manager import Account, AccountManager
 from .app_controller import AppController
+from .backend_config import get_backend_config
 from .backends import DemoBackend, get_backend
 from .cache_manager import CacheManager
 from .credentials import CredentialManager
@@ -53,6 +54,7 @@
 # Screen imports
 from .screens.account_name_input_screen import AccountNameInputScreen
 from .screens.account_selector_screen import AccountSelectorScreen
+from .screens.budget_selector_screen import BudgetSelectorScreen
 from .screens.credential_screens import (
     BackendSelectionScreen,
     CredentialSetupScreen,
@@ -561,6 +563,93 @@ async def _handle_account_selection(self):
             # Success - return account info
             return account.id, profile_dir, creds
 
+    def _get_ynab_budget_id(self, account_id: str) -> Optional[str]:
+        """
+        Look up the budget_id for a YNAB account.
+
+        Args:
+            account_id: The account ID to look up
+
+        Returns:
+            The budget_id if found, None otherwise
+        """
+        config_path = Path(self.config_dir) if self.config_dir else None
+        account_manager = AccountManager(config_dir=config_path)
+        account = account_manager.get_account(account_id)
+
+        if account and account.budget_id:
+            return account.budget_id
+        return None
+
+    async def _handle_ynab_budget_selection(
+        self,
+        creds: dict,
+        account: Account,
+        account_manager: AccountManager,
+    ) -> Optional[str]:
+        """
+        Handle YNAB-specific budget selection after credentials are set up.
+
+        This method encapsulates the YNAB budget selection flow:
+        1. Create temporary backend and login
+        2. Fetch available budgets
+        3. Show budget selector if multiple budgets exist
+        4. Update account with selected budget_id
+
+        Args:
+            creds: Credentials dictionary (must contain "password" with YNAB token)
+            account: The newly created Account object
+            account_manager: AccountManager for updating the account
+
+        Returns:
+            budget_id if successful, None if user cancelled or error occurred.
+            On error/cancel, this method handles cleanup (clearing creds, deleting account).
+        """
+        logger = get_logger(__name__)
+        temp_backend = get_backend("ynab")
+
+        try:
+            await temp_backend.login(password=creds["password"])
+            budgets = await temp_backend.get_budgets()  # type: ignore
+
+            budget_id = None
+            if len(budgets) > 1:
+                # Show budget selector
+                budget_id = await self.push_screen(
+                    BudgetSelectorScreen(budgets), wait_for_dismiss=True
+                )
+
+                if budget_id is None:
+                    # User cancelled budget selection - clean up
+                    creds.clear()
+                    account_manager.delete_account(account.id)
+                    return None
+            elif len(budgets) == 1:
+                budget_id = budgets[0]["id"]
+
+            # Update the account with the selected budget_id
+            if budget_id:
+                account.budget_id = budget_id
+                registry = account_manager.load_registry()
+                for i, acc in enumerate(registry.accounts):
+                    if acc.id == account.id:
+                        registry.accounts[i] = account
+                        break
+                account_manager.save_registry(registry)
+
+            return budget_id
+
+        except Exception:
+            logger.error("Failed to fetch YNAB budgets during account setup")
+            creds.clear()
+            account_manager.delete_account(account.id)
+            return None
+        finally:
+            # Always clear temporary backend auth to minimize credential exposure
+            # Note: creds dict is intentionally NOT cleared here on success because
+            # it's returned and used for the actual backend login later.
+            temp_backend.clear_auth()
+
     async def _handle_add_new_account(self, account_manager: AccountManager):
         """
         Handle adding a new account.
@@ -588,7 +677,7 @@ async def _handle_add_new_account(self, account_manager: AccountManager):
         if not backend_type:
             return None  # User cancelled
 
-        # Step 3: Create account profile
+        # Step 3: Create account profile first
         try:
             account = account_manager.create_account(name=account_name, backend_type=backend_type)
         except ValueError as e:
@@ -597,21 +686,16 @@ async def _handle_add_new_account(self, account_manager: AccountManager):
             logger.error(f"Failed to create account: {e}")
             return None
 
-        # Step 4: Get credentials (if backend requires auth)
-        from .backend_config import BackendConfig
+        # Step 4: Get/store credentials and handle backend-specific setup
+        from .backend_config import get_backend_config
 
-        backend_config = {
-            "monarch": BackendConfig.for_monarch(),
-            "ynab": BackendConfig.for_ynab(),
-            "amazon": BackendConfig.for_amazon(),
-            "demo": BackendConfig.for_demo(),
-        }.get(backend_type, BackendConfig.for_monarch())
+        backend_config = get_backend_config(backend_type)
 
         # Get profile directory
         profile_dir = account_manager.get_profile_dir(account.id)
 
         if backend_config.requires_auth:
-            # Show credential setup
+            # Show credential setup with profile_dir so it saves to the right place
             creds = await self.push_screen(
                 CredentialSetupScreen(backend_type=backend_type, profile_dir=profile_dir),
                 wait_for_dismiss=True,
@@ -621,18 +705,28 @@ async def _handle_add_new_account(self, account_manager: AccountManager):
                 # User cancelled - delete the account we just created
                 account_manager.delete_account(account.id)
                 return None
+
+            # Handle backend-specific post-credential setup (e.g., YNAB budget selection)
+            if backend_type == "ynab":
+                budget_id = await self._handle_ynab_budget_selection(
+                    creds, account, account_manager
+                )
+                if budget_id is None:
+                    # User cancelled or error occurred - cleanup already handled
+                    return None
         else:
             # Backend doesn't need credentials (Amazon, Demo)
             creds = {"backend_type": backend_type}
 
         return account.id, profile_dir, creds
 
-    async def _login_with_retry(self, creds, loading_status):
+    async def _login_with_retry(self, creds, loading_status, budget_id=None):
         """Login with retry logic for robustness.
 
         Args:
             creds: Credentials dict
             loading_status: Loading status widget
+            budget_id: Optional budget ID for YNAB accounts
 
         Returns:
             bool: True on success, False on failure
@@ -657,13 +751,22 @@ async def login_operation():
             """Login with automatic retry on session expiration."""
             try:
                 logger.debug("Attempting login with saved session...")
-                await self.backend.login(
-                    email=creds["email"],
-                    password=creds["password"],
-                    use_saved_session=True,  # Try saved session first
-                    save_session=True,
-                    mfa_secret_key=creds["mfa_secret"],
-                )
+
+                # Simple login - budget selection happens during account creation
+                login_kwargs = {
+                    "email": creds["email"],
+                    "password": creds["password"],
+                    "use_saved_session": True,
+                    "save_session": True,
+                    "mfa_secret_key": creds["mfa_secret"],
+                }
+
+                # For YNAB, include budget_id if available
+                if backend_type == "ynab" and budget_id:
+                    login_kwargs["budget_id"] = budget_id
+
+                await self.backend.login(**login_kwargs)
+
                 logger.debug("Login succeeded!")
                 return True
             except Exception as e:
@@ -988,22 +1091,19 @@ async def initialize_data(self) -> None:
                     loading_status.update("🎮 DEMO MODE - Loading sample data...")
                 else:
                     # Load account info to get backend_type
-                    from moneyflow.account_manager import AccountManager
-
                     config_path = Path(self.config_dir) if self.config_dir else None
                     account_manager = AccountManager(config_dir=config_path)
                     account = account_manager.get_account(account_id)
 
                     if account and account.backend_type == "amazon" and profile_dir:
                         # Initialize Amazon backend with profile-scoped database
-                        from moneyflow.backend_config import BackendConfig
                         from moneyflow.backends.amazon import AmazonBackend
 
                         db_path = str(profile_dir / "amazon.db")
                         self.backend = AmazonBackend(
                             db_path=db_path, config_dir=self.config_dir, profile_dir=profile_dir
                         )
-                        self.backend_config = BackendConfig.for_amazon()
+                        self.backend_config = get_backend_config("amazon")
                         loading_status.update("📦 Loading Amazon data...")
 
             # Step 2: Initialize backend (if not already set)
@@ -1017,26 +1117,22 @@ async def initialize_data(self) -> None:
                     backend_kwargs["profile_dir"] = str(self._preconfigured_profile_dir)
 
                 self.backend = get_backend(backend_type, **backend_kwargs)
-
-                from moneyflow.backend_config import BackendConfig
-
-                if backend_type == "ynab":
-                    self.backend_config = BackendConfig.for_ynab()
-                elif backend_type == "monarch":
-                    self.backend_config = BackendConfig.for_monarch()
-                elif backend_type == "amazon":
-                    self.backend_config = BackendConfig.for_amazon()
-                elif backend_type == "demo":
-                    self.backend_config = BackendConfig.for_demo()
+                self.backend_config = get_backend_config(backend_type)
 
                 # Step 3: Login with retry logic
-                login_success = await self._login_with_retry(creds, loading_status)
+                # For YNAB, get budget_id from account if available
+                budget_id = None
+                if backend_type == "ynab" and account_id:
+                    budget_id = self._get_ynab_budget_id(account_id)
+
+                login_success = await self._login_with_retry(creds, loading_status, budget_id)
                 if not login_success:
                     has_error = True
                     return
             elif self.backend and not self.demo_mode:
                 # Backend exists but might need login
                 if self.backend_config.requires_auth and creds:
+                    # For pre-configured backends, we don't have account_id to look up budget_id
                     login_success = await self._login_with_retry(creds, loading_status)
                     if not login_success:
                         has_error = True
@@ -1054,14 +1150,8 @@ async def initialize_data(self) -> None:
                 if self.demo_mode:
                     determined_backend_type = "demo"
                 elif self.backend:
-                    # Get backend type from backend instance
-                    backend_class = self.backend.__class__.__name__
-                    if "Amazon" in backend_class:
-                        determined_backend_type = "amazon"
-                    elif "Monarch" in backend_class:
-                        determined_backend_type = "monarch"
-                    elif "YNAB" in backend_class:
-                        determined_backend_type = "ynab"
+                    # Get backend type from backend instance (Open/Closed Principle)
+                    determined_backend_type = self.backend.get_backend_type()
                 elif creds:
                     determined_backend_type = creds.get("backend_type")
 
@@ -2262,7 +2352,7 @@ def launch_amazon_mode(
     Uses the AmazonBackend with data stored in SQLite.
     Data must be imported first using: moneyflow amazon import <csv>
     """
-    from moneyflow.backend_config import BackendConfig
+    from moneyflow.backend_config import get_backend_config
     from moneyflow.backends.amazon import AmazonBackend
 
     # Initialize logging
@@ -2276,7 +2366,7 @@ def launch_amazon_mode(
     try:
         # Create Amazon backend and config
         backend = AmazonBackend(db_path=db_path, config_dir=config_dir, profile_dir=profile_dir)
-        config = BackendConfig.for_amazon()
+        config = get_backend_config("amazon")
 
         # Create MoneyflowApp in Amazon mode
         app = MoneyflowApp(