"""Workspace commands for Cua CLI.""" import argparse import asyncio import os import aiohttp from cua_cli.auth.browser import authenticate_via_browser from cua_cli.auth.store import ( _get_store, delete_workspace, get_workspace_api_key, list_workspaces, save_workspace, set_active_workspace, ) from cua_cli.utils.async_utils import run_async from cua_cli.utils.output import print_error, print_info, print_success from cua_core.http import cua_version_headers DEFAULT_API_BASE = "https://api.cua.ai" def _get_api_base() -> str: return os.environ.get("CUA_API_BASE", DEFAULT_API_BASE).rstrip("/") def register_parser(subparsers: argparse._SubParsersAction) -> None: """Register the workspace command and subcommands.""" ws_parser = subparsers.add_parser( "workspace", help="Workspace commands", description="Manage workspaces", aliases=["ws"], ) ws_subparsers = ws_parser.add_subparsers( dest="workspace_command", help="Workspace command", ) set_parser = ws_subparsers.add_parser( "set", help="Switch active workspace", description="Switch to a different workspace", ) set_parser.add_argument( "slug", nargs="?", help="Workspace slug to switch to", ) def execute(args: argparse.Namespace) -> int: """Execute workspace command based on subcommand.""" cmd = getattr(args, "workspace_command", None) if cmd == "set": return cmd_set(args) else: print_error("Usage: cua workspace ") print_info("Commands: set") return 1 def _validate_workspace_key(api_key: str) -> tuple[bool | None, dict]: """Validate an API key by calling /v1/me. Returns: (True, data) — valid key (HTTP 200) (False, data) — explicitly invalid/expired (HTTP 401) (None, {}) — transport error or unexpected HTTP status """ async def _do(): url = f"{_get_api_base()}/v1/me" headers = { "Authorization": f"Bearer {api_key}", "Accept": "application/json", **cua_version_headers(), } async with aiohttp.ClientSession() as session: timeout = aiohttp.ClientTimeout(total=10) async with session.get(url, headers=headers, timeout=timeout) as resp: try: data = await resp.json(content_type=None) except Exception: return resp.status, {} return resp.status, data try: status_code, data = run_async(_do()) if status_code == 200: return True, data elif status_code == 401: return False, data else: return None, {} except (aiohttp.ClientError, asyncio.TimeoutError, OSError): return None, {} except Exception: return None, {} def cmd_set(args: argparse.Namespace) -> int: """Handle the workspace set command.""" slug = getattr(args, "slug", None) if slug: # Check if we have a stored key for this workspace stored_key = get_workspace_api_key(slug) if stored_key: # Validate the stored key valid, _data = _validate_workspace_key(stored_key) if valid is True: set_active_workspace(slug) ws_name = _get_store().get(f"workspace:{slug}:name") or slug print_success(f"Switched to workspace: {ws_name} ({slug})") return 0 elif valid is None: # Transport error — don't delete cached credentials print_error( f"Could not verify credentials for '{slug}' (API unreachable). Try again later." ) return 1 else: # Explicitly expired (401) — remove and fall through to browser auth print_info(f"Credentials for '{slug}' are expired. Re-authenticating...") delete_workspace(slug) # Not authenticated or stale — trigger browser auth with hint try: result = run_async(authenticate_via_browser(workspace_slug=slug)) except (TimeoutError, RuntimeError) as e: print_error(str(e)) return 1 if result.workspace_slug: if result.workspace_slug != slug: print_error( f"Authenticated workspace '{result.workspace_slug}' does not match " f"requested workspace '{slug}'." ) return 1 save_workspace( result.workspace_slug, result.token, result.workspace_name, result.org_name, ) set_active_workspace(result.workspace_slug) print_success( f"Authenticated and switched to workspace: " f"{result.workspace_name} ({result.workspace_slug})" ) else: print_error("Authentication did not return workspace metadata.") return 1 return 0 # No slug — interactive selection from authenticated workspaces workspaces = list_workspaces() if not workspaces: print_info("No authenticated workspaces. Run 'cua auth login' to add one.") return 0 print_info("Select a workspace:") for i, ws in enumerate(workspaces, 1): marker = " (active)" if ws["is_active"] else "" name_part = f" - {ws['name']}" if ws["name"] else "" print(f" {i}. {ws['slug']}{name_part}{marker}") try: choice = input("\nEnter number: ").strip() idx = int(choice) - 1 if idx < 0 or idx >= len(workspaces): print_error("Invalid selection.") return 1 except (ValueError, EOFError, KeyboardInterrupt): print_error("Invalid selection.") return 1 selected = workspaces[idx] # Validate the selected workspace's session before switching stored_key = get_workspace_api_key(selected["slug"]) if stored_key: valid, _data = _validate_workspace_key(stored_key) if valid is None: print_error( f"Could not verify credentials for '{selected['slug']}' (API unreachable). " f"Try again later." ) return 1 elif valid is False: delete_workspace(selected["slug"]) print_error( f"Credentials for '{selected['slug']}' are expired. " f"Run 'cua auth login' to re-authenticate." ) return 1 set_active_workspace(selected["slug"]) print_success(f"Switched to workspace: {selected['name']} ({selected['slug']})") return 0