chore: import upstream snapshot with attribution

libs/python/cua-cli/cua_cli/commands/workspace.py

@@ -0,0 +1,205 @@
+"""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 <command>")
+        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