feat(cli): Add suport for multi-workspace authentication#1151
feat(cli): Add suport for multi-workspace authentication#1151
Conversation
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
📦 Publishable packages changed
Add |
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughThis pull request introduces multi-workspace support to the Cua CLI. It adds workspace-scoped credential management, enhances OAuth authentication to return workspace metadata, extends the auth command with workspace listing/management capabilities, and introduces a new workspace command for switching between cached workspaces with browser-based fallback authentication. Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant CLI as Workspace CLI
participant Store as CredentialStore
participant Browser as Browser Auth
participant API as CUA API
User->>CLI: workspace set my-workspace
CLI->>Store: get_workspace_api_key('my-workspace')
alt Cached key exists
Store-->>CLI: api_key
CLI->>API: GET /v1/me (with api_key)
API-->>CLI: 200 + workspace metadata
CLI->>Store: set_active_workspace('my-workspace')
CLI-->>User: ✓ Workspace activated
else No cached key
CLI->>Browser: authenticate_via_browser(workspace_slug='my-workspace')
Browser->>API: OAuth flow
API-->>Browser: AuthResult(token, slug, name, org)
Browser-->>CLI: AuthResult
CLI->>Store: save_workspace(slug, token, name, org)
CLI->>Store: set_active_workspace(slug)
CLI-->>User: ✓ Workspace activated
end
sequenceDiagram
participant User
participant CLI as Auth CLI
participant API as CUA API
participant Store as CredentialStore
User->>CLI: cua auth login
CLI->>CLI: authenticate_via_browser()
CLI->>API: OAuth callback with auth_token + workspace metadata
API-->>CLI: AuthResult
CLI->>API: GET /v1/me (with auth_token)
API-->>CLI: workspace_slug, name, org
CLI->>Store: save_workspace(slug, token, name, org)
CLI->>Store: set_active_workspace(slug)
CLI-->>User: ✓ Logged in to workspace
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
📝 Coding Plan
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment Tip CodeRabbit can generate a title for your PR based on the changes with custom instructions.Set the |
| description="Cua CLI - Unified command-line interface for Computer-Use Agents", | ||
| formatter_class=argparse.RawDescriptionHelpFormatter, | ||
| epilog=""" | ||
| Examples: |
There was a problem hiding this comment.
I removed these examples because they were annoying me when trying to view the commands in the cli.
Can add them back if we think they're hepful.
📦 Publishable packages changed
Add |
📦 Publishable packages changed
Add |
📦 Publishable packages changed
Add |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 7
🧹 Nitpick comments (1)
libs/python/cua-cli/cua_cli/auth/store.py (1)
205-213: Return workspaces in a deterministic order.This list now drives both the interactive picker and the
remaining[0]fallback inlibs/python/cua-cli/cua_cli/commands/auth.py. Without sorting here, the next active workspace depends on SQLite row order and can change across runs.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@libs/python/cua-cli/cua_cli/auth/store.py` around lines 205 - 213, The returned workspace list is nondeterministic; after building result from rows/slugs/active, sort it deterministically before returning (e.g., result.sort(key=lambda ws: ws["slug"]) so order is stable across runs and the interactive picker / remaining[0] fallback behave predictably). Ensure you perform the sort on the result list (referencing variables result, rows, slugs, active) and then return the sorted list.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@libs/python/cua-cli/cua_cli/auth/store.py`:
- Around line 269-273: The current clear_credentials() implementation deletes
all workspaces via clear_all_workspaces() and the legacy API key (API_KEY_NAME);
split this behavior by adding a new function (e.g., clear_legacy_credentials()
or clear_legacy_session()) that only deletes the legacy key
(store.delete(API_KEY_NAME)) and update callers in commands/auth.py (the
no-active-workspace branches of cmd_logout() and cmd_status()) to call the new
narrow function, while keeping clear_credentials() as the existing "clear
everything" helper used for the --all path.
In `@libs/python/cua-cli/cua_cli/commands/auth.py`:
- Around line 10-20: Re-run isort to normalize the import ordering in this
module: reorder the top-level imports so standard/library imports come first,
then third-party, then local/package imports, and ensure the multi-import from
cua_cli.auth.store is alphabetized (clear_credentials, delete_workspace,
get_active_workspace, get_api_key, list_workspaces, save_api_key,
save_workspace, set_active_workspace) and that authenticate_via_browser is
positioned according to isort rules; then save the file so the CI Python lint
job passes.
- Around line 295-315: The 401 handler is mutating cached workspaces even when
the failing token may have come from the environment (CUA_API_KEY); update the
logic around the _fetch_me call and subsequent 401 branch to first determine the
credential source (e.g., via get_api_key() or by checking
os.environ["CUA_API_KEY"]/a returned source flag), then only delete or switch
the cached active workspace (using get_active_workspace, delete_workspace,
list_workspaces, set_active_workspace, or
_get_store().delete(ACTIVE_WORKSPACE_KEY)) when the failing credential is the
stored workspace token—not when the token came from the environment; apply the
same source-aware check to the similar cleanup at lines ~328-330.
- Around line 171-178: If the workspace slug is missing, avoid creating a
synthetic "workspace:default:*" entry; instead treat this as a legacy-login
fallback: change the branch that currently calls save_workspace(slug, api_key,
name, org_name) and set_active_workspace(slug) to detect missing ws.get("slug")
and call the legacy-save path (the same logic used in the 202-205 legacy block)
and clear any active workspace marker (call the function that clears active
workspace) rather than creating a synthetic workspace; apply the same adjustment
to the other similar block around lines 190-205 so both paths consistently
fallback to legacy behavior and clear the active workspace when slug is absent.
In `@libs/python/cua-cli/cua_cli/commands/workspace.py`:
- Around line 65-88: The validator currently collapses transport failures and
expired tokens into a single False result; change _validate_workspace_key to
distinguish outcomes: return (True, data) when status == 200, return (False,
data) when status == 401 (explicitly expired/invalid), and return (None, {}) for
transport errors (catch aiohttp.ClientError, asyncio.TimeoutError and other
exceptions) or non-401/non-200 statuses so callers know the API was unreachable
or returned an unexpected status; update cmd_set (the caller that currently
treats any False as "expired") to treat None as "transient/unreachable" and
avoid deleting the cached workspace token in that case.
- Around line 135-159: The interactive branch currently picks a workspace from
list_workspaces and calls set_active_workspace without verifying the cached
session; replicate the explicit-slug validation flow by validating the selected
workspace's session (e.g., attempt the same /v1/me check or call the same
validation helper used for the slug path) before calling set_active_workspace;
if validation fails, remove the stale cached keys (same cleanup used in the
explicit path), print an error and return non-zero, otherwise proceed to
set_active_workspace and print_success for the selected workspace.
- Around line 118-129: The handler currently accepts any non-empty
result.workspace_slug and switches to it; change it to validate the returned
workspace against the originally requested slug (compare result.workspace_slug
to the input/requested slug variable used to start the browser flow) and abort
if they differ: do not call save_workspace or set_active_workspace, instead
print an error (or print_failure) explaining the mismatch and return/non-zero
exit; keep the existing print_success path only for exact matches. Use the
existing symbols result.workspace_slug, save_workspace, set_active_workspace,
and print_success when locating the code to modify.
---
Nitpick comments:
In `@libs/python/cua-cli/cua_cli/auth/store.py`:
- Around line 205-213: The returned workspace list is nondeterministic; after
building result from rows/slugs/active, sort it deterministically before
returning (e.g., result.sort(key=lambda ws: ws["slug"]) so order is stable
across runs and the interactive picker / remaining[0] fallback behave
predictably). Ensure you perform the sort on the result list (referencing
variables result, rows, slugs, active) and then return the sorted list.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 412cf2f7-a8d7-495b-8eaf-6437278a4056
📒 Files selected for processing (5)
libs/python/cua-cli/cua_cli/auth/browser.pylibs/python/cua-cli/cua_cli/auth/store.pylibs/python/cua-cli/cua_cli/commands/auth.pylibs/python/cua-cli/cua_cli/commands/workspace.pylibs/python/cua-cli/cua_cli/main.py
| def clear_credentials() -> None: | ||
| """Clear all stored credentials.""" | ||
| _get_store().clear() | ||
| """Clear all stored credentials including workspaces and legacy key.""" | ||
| store = _get_store() | ||
| clear_all_workspaces() | ||
| store.delete(API_KEY_NAME) |
There was a problem hiding this comment.
Split all-workspace clearing from legacy-session cleanup.
libs/python/cua-cli/cua_cli/commands/auth.py still uses this helper in the "no active workspace" branches of cmd_logout() and cmd_status(). With the new implementation, those paths now wipe every cached workspace instead of just clearing the current legacy token. Keep this helper for --all and add a narrower legacy-only clear for those callers.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@libs/python/cua-cli/cua_cli/auth/store.py` around lines 269 - 273, The
current clear_credentials() implementation deletes all workspaces via
clear_all_workspaces() and the legacy API key (API_KEY_NAME); split this
behavior by adding a new function (e.g., clear_legacy_credentials() or
clear_legacy_session()) that only deletes the legacy key
(store.delete(API_KEY_NAME)) and update callers in commands/auth.py (the
no-active-workspace branches of cmd_logout() and cmd_status()) to call the new
narrow function, while keeping clear_credentials() as the existing "clear
everything" helper used for the --all path.
| from cua_cli.auth.browser import authenticate_via_browser | ||
| from cua_cli.auth.store import clear_credentials, get_api_key, save_api_key | ||
| from cua_cli.auth.store import ( | ||
| clear_credentials, | ||
| delete_workspace, | ||
| get_active_workspace, | ||
| get_api_key, | ||
| list_workspaces, | ||
| save_api_key, | ||
| save_workspace, | ||
| set_active_workspace, | ||
| ) |
There was a problem hiding this comment.
Re-run isort on this import block.
CI is already failing the Python lint job on this file, so this needs to be normalized before the PR can merge.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@libs/python/cua-cli/cua_cli/commands/auth.py` around lines 10 - 20, Re-run
isort to normalize the import ordering in this module: reorder the top-level
imports so standard/library imports come first, then third-party, then
local/package imports, and ensure the multi-import from cua_cli.auth.store is
alphabetized (clear_credentials, delete_workspace, get_active_workspace,
get_api_key, list_workspaces, save_api_key, save_workspace,
set_active_workspace) and that authenticate_via_browser is positioned according
to isort rules; then save the file so the CI Python lint job passes.
| ws = data.get("workspace", {}) | ||
| org = data.get("organization", {}) | ||
| slug = ws.get("slug", "default") | ||
| name = ws.get("name", slug) | ||
| org_name = org.get("name", "") | ||
| save_workspace(slug, api_key, name, org_name) | ||
| set_active_workspace(slug) | ||
| print_success(f"Authenticated and switched to workspace: {name} ({slug})") |
There was a problem hiding this comment.
Treat missing workspace metadata as a real legacy-login fallback.
Lines 171-178 currently invent workspace:default:*, while Lines 202-205 save a legacy key without clearing any existing active workspace. Both paths can make a "successful" login keep using the wrong credentials. If the slug is missing, stay on the legacy path and clear the active-workspace marker instead of creating a synthetic workspace entry.
Also applies to: 190-205
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@libs/python/cua-cli/cua_cli/commands/auth.py` around lines 171 - 178, If the
workspace slug is missing, avoid creating a synthetic "workspace:default:*"
entry; instead treat this as a legacy-login fallback: change the branch that
currently calls save_workspace(slug, api_key, name, org_name) and
set_active_workspace(slug) to detect missing ws.get("slug") and call the
legacy-save path (the same logic used in the 202-205 legacy block) and clear any
active workspace marker (call the function that clears active workspace) rather
than creating a synthetic workspace; apply the same adjustment to the other
similar block around lines 190-205 so both paths consistently fallback to legacy
behavior and clear the active workspace when slug is absent.
| status_code, data = _fetch_me(api_key) | ||
| except Exception as e: | ||
| print_error(f"Failed to reach API: {e}") | ||
| return 1 | ||
|
|
||
| if status_code == 401: | ||
| clear_credentials() | ||
| active = get_active_workspace() | ||
| if active: | ||
| delete_workspace(active) | ||
| print_info(f"Removed expired credentials for workspace: {active}") | ||
| remaining = list_workspaces() | ||
| if remaining: | ||
| set_active_workspace(remaining[0]["slug"]) | ||
| print_info(f"Switched to workspace: {remaining[0]['name']} ({remaining[0]['slug']})") | ||
| else: | ||
| from cua_cli.auth.store import _get_store, ACTIVE_WORKSPACE_KEY | ||
|
|
||
| _get_store().delete(ACTIVE_WORKSPACE_KEY) | ||
| else: | ||
| clear_credentials() | ||
| print_info("Removed expired credentials.") |
There was a problem hiding this comment.
Don't mutate cached workspaces when the failing token came from CUA_API_KEY.
get_api_key() prefers the environment, but this 401 handler and the new [active: ...] label both assume the active-workspace credential was used. With CUA_API_KEY set, a failed status call can delete unrelated cached workspaces and still report the wrong active workspace. Track the credential source before cleaning up or labeling the session.
Also applies to: 328-330
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@libs/python/cua-cli/cua_cli/commands/auth.py` around lines 295 - 315, The 401
handler is mutating cached workspaces even when the failing token may have come
from the environment (CUA_API_KEY); update the logic around the _fetch_me call
and subsequent 401 branch to first determine the credential source (e.g., via
get_api_key() or by checking os.environ["CUA_API_KEY"]/a returned source flag),
then only delete or switch the cached active workspace (using
get_active_workspace, delete_workspace, list_workspaces, set_active_workspace,
or _get_store().delete(ACTIVE_WORKSPACE_KEY)) when the failing credential is the
stored workspace token—not when the token came from the environment; apply the
same source-aware check to the similar cleanup at lines ~328-330.
| def _validate_workspace_key(api_key: str) -> tuple[bool, dict]: | ||
| """Validate an API key by calling /v1/me. Returns (valid, data).""" | ||
|
|
||
| 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()) | ||
| return status_code == 200, data | ||
| except Exception: | ||
| return False, {} |
There was a problem hiding this comment.
Differentiate expired keys from reachability failures.
_validate_workspace_key() returns False for timeouts, DNS errors, and other transport problems, and cmd_set() treats every False as "expired" and deletes the cached workspace. A transient API outage should not discard a valid token.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@libs/python/cua-cli/cua_cli/commands/workspace.py` around lines 65 - 88, The
validator currently collapses transport failures and expired tokens into a
single False result; change _validate_workspace_key to distinguish outcomes:
return (True, data) when status == 200, return (False, data) when status == 401
(explicitly expired/invalid), and return (None, {}) for transport errors (catch
aiohttp.ClientError, asyncio.TimeoutError and other exceptions) or
non-401/non-200 statuses so callers know the API was unreachable or returned an
unexpected status; update cmd_set (the caller that currently treats any False as
"expired") to treat None as "transient/unreachable" and avoid deleting the
cached workspace token in that case.
| if result.workspace_slug: | ||
| 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})" | ||
| ) |
There was a problem hiding this comment.
Reject browser auth for the wrong workspace.
cua workspace set <slug> currently accepts any non-empty result.workspace_slug and switches to it. If the browser flow comes back with a different workspace than the one requested, the CLI silently lands in the wrong tenant.
Suggested guard
- if result.workspace_slug:
+ if result.workspace_slug == slug:
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})"
)
+ elif result.workspace_slug:
+ print_error(
+ f"Authentication returned workspace '{result.workspace_slug}', "
+ f"expected '{slug}'."
+ )
+ return 1
else:
print_error("Authentication did not return workspace metadata.")
return 1📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if result.workspace_slug: | |
| 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})" | |
| ) | |
| if result.workspace_slug == slug: | |
| 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})" | |
| ) | |
| elif result.workspace_slug: | |
| print_error( | |
| f"Authentication returned workspace '{result.workspace_slug}', " | |
| f"expected '{slug}'." | |
| ) | |
| return 1 | |
| else: | |
| print_error("Authentication did not return workspace metadata.") | |
| return 1 |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@libs/python/cua-cli/cua_cli/commands/workspace.py` around lines 118 - 129,
The handler currently accepts any non-empty result.workspace_slug and switches
to it; change it to validate the returned workspace against the originally
requested slug (compare result.workspace_slug to the input/requested slug
variable used to start the browser flow) and abort if they differ: do not call
save_workspace or set_active_workspace, instead print an error (or
print_failure) explaining the mismatch and return/non-zero exit; keep the
existing print_success path only for exact matches. Use the existing symbols
result.workspace_slug, save_workspace, set_active_workspace, and print_success
when locating the code to modify.
| # No slug — interactive selection from cached workspaces | ||
| workspaces = list_workspaces() | ||
| if not workspaces: | ||
| print_info("No workspaces cached. 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] | ||
| set_active_workspace(selected["slug"]) | ||
| print_success(f"Switched to workspace: {selected['name']} ({selected['slug']})") |
There was a problem hiding this comment.
Validate the picked workspace before marking it active.
The explicit-slug path checks /v1/me and cleans up stale keys, but the interactive branch just flips active_workspace. That lets users "successfully" switch into an expired cached entry and only fail on the next command.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@libs/python/cua-cli/cua_cli/commands/workspace.py` around lines 135 - 159,
The interactive branch currently picks a workspace from list_workspaces and
calls set_active_workspace without verifying the cached session; replicate the
explicit-slug validation flow by validating the selected workspace's session
(e.g., attempt the same /v1/me check or call the same validation helper used for
the slug path) before calling set_active_workspace; if validation fails, remove
the stale cached keys (same cleanup used in the explicit path), print an error
and return non-zero, otherwise proceed to set_active_workspace and print_success
for the selected workspace.
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
📦 Publishable packages changed
Add |
📦 Publishable packages changed
Add |
📦 Publishable packages changed
Add |
📦 Publishable packages changed
Add |
2 participants
Summary
Add gcloud-style multi-workspace switching. Users can
cua auth loginmultiple times to cache credentials per workspace, then switch between them withcua workspace set <slug>without re-authenticating.What changed
auth/store.py) — Namespaced KV keys per workspace (workspace:<slug>:api_key/name/org), with legacyapi_keyfallback for backward compat. No schema migration needed.auth/browser.py) — NewAuthResultdataclass. Parsesworkspace_slug,workspace_name,org_namefrom callback. Accepts optionalworkspace_slughint for pre-selection.commands/auth.py) —loginsaves workspace metadata for both browser and--api-keyflows.logoutsupports--workspace <slug>and--allflags. Newlistsubcommand shows cached workspaces grouped by org.statusshows active workspace label.commands/workspace.py) — Newcua workspace set [slug](aliascua ws set). Validates cached keys via/v1/me, falls back to browser auth, supports interactive picker.main.py) — Registers workspace command.Backward compat
Existing users with a single
api_keyin the DB are unaffected —get_api_key()falls back to it. Nextcua auth loginmigrates to the new format naturally. No manual migration needed.Summary by CodeRabbit
workspace(aliasws) command for switching between saved workspaces via interactive selection.auth listcommand to display cached workspaces grouped by organization, with active status indicators.--workspace) and bulk (--all) removal options.