llamactl 0.3.0a18__py3-none-any.whl → 0.3.0a20__py3-none-any.whl

llama_deploy/cli/commands/auth.py

@@ -1,18 +1,36 @@
 import asyncio
+import json
+import os
+import platform
+import subprocess
+import sys
+import webbrowser
 
 import click
+import httpx
 import questionary
-from llama_deploy.cli.client import get_control_plane_client
+from llama_deploy.cli.auth.client import (
+    DeviceAuthorizationRequest,
+    OIDCClient,
+    PlatformAuthClient,
+    PlatformAuthDiscoveryClient,
+    TokenRequestDeviceCode,
+    decode_jwt_claims,
+    decode_jwt_claims_from_device_oidc,
+)
+from llama_deploy.cli.config._config import ConfigManager
 from llama_deploy.cli.config.auth_service import AuthService
 from llama_deploy.cli.config.env_service import service
-from llama_deploy.core.client.manage_client import ClientError, ControlPlaneClient
+from llama_deploy.core.client.manage_client import (
+    ControlPlaneClient,
+)
 from llama_deploy.core.schema.projects import ProjectSummary
 from rich import print as rprint
 from rich.table import Table
 from rich.text import Text
 
 from ..app import app, console
-from ..config.schema import Auth
+from ..config.schema import Auth, DeviceOIDC
 from ..options import global_options, interactive_option
 
 
@@ -61,7 +79,7 @@ def create_api_key_profile(
 
         # Interactive mode: prompt for token (masked) and validate
         token_value = api_key or _prompt_for_api_key()
-        projects = _validate_token_and_list_projects(auth_svc, token_value)
+        projects = _prompt_validate_api_key_and_list_projects(auth_svc, token_value)
 
         # Select or enter project ID
         selected_project_id = project_id or _select_or_enter_project(
@@ -81,10 +99,26 @@ def create_api_key_profile(
         raise click.Abort()
 
 
+@auth.command("login")
+@global_options
+def device_login() -> None:
+    """Login via web browser"""
+
+    try:
+        created = _create_device_profile()
+        rprint(
+            f"[green]Created login profile '{created.name}' and set as current[/green]"
+        )
+
+    except Exception as e:
+        rprint(f"[red]Error: {e}[/red]")
+        raise click.Abort()
+
+
 @auth.command("list")
 @global_options
 def list_profiles() -> None:
-    """List all logged in profiles"""
+    """List all logged in users/tokens"""
     try:
         auth_svc = service.current_auth_service()
         profiles = auth_svc.list_profiles()
@@ -119,6 +153,26 @@ def list_profiles() -> None:
         raise click.Abort()
 
 
+@auth.command("destroy", hidden=True)
+@global_options
+def destroy_database() -> None:
+    """Destroy the database"""
+    if not questionary.confirm(
+        "Are you sure you want to destroy all of your local logins? This action cannot be undone."
+    ).ask():
+        return
+    ConfigManager(init_database=False).destroy_database()
+    rprint("[green]Database destroyed[/green]")
+
+
+@auth.command("show-db", hidden=True)
+@global_options
+def config_database() -> None:
+    """Config the database"""
+    path = service.config_manager().db_path
+    rprint(f"[bold]{path}[/bold]")
+
+
 @auth.command("switch")
 @global_options
 @click.argument("name", required=False)
@@ -153,7 +207,7 @@ def delete_profile(name: str | None, interactive: bool) -> None:
             rprint("[yellow]No profile selected[/yellow]")
             return
 
-        if auth_svc.delete_profile(auth.name):
+        if asyncio.run(auth_svc.delete_profile(auth.name)):
             rprint(f"[green]Logged out from '{auth.name}'[/green]")
         else:
             rprint(f"[red]Profile '{auth.name}' not found[/red]")
@@ -163,6 +217,29 @@ def delete_profile(name: str | None, interactive: bool) -> None:
         raise click.Abort()
 
 
+# Very simple introspection: decode current token via provider JWKS
+@auth.command("me", hidden=True)
+@global_options
+def me() -> None:
+    """Print JWT claims for the current profile's token using provider JWKS.
+
+    Assumes the stored API key is a JWT (e.g., OIDC id_token).
+    """
+    try:
+        auth_svc = service.current_auth_service()
+        profile = auth_svc.get_current_profile()
+        if not profile or not profile.device_oidc:
+            raise click.ClickException(
+                "No OIDC profile selected. Run `llamactl auth login` or switch to an existing OIDC profile."
+            )
+
+        claims = asyncio.run(decode_jwt_claims_from_device_oidc(profile.device_oidc))
+        click.echo(json.dumps(claims, indent=2, sort_keys=True))
+    except Exception as e:
+        rprint(f"[red]Error: {e}[/red]")
+        raise click.Abort()
+
+
 # Projects commands
 @auth.command("project")
 @click.argument("project_id", required=False)
@@ -175,16 +252,22 @@ def change_project(project_id: str | None, interactive: bool) -> None:
         return
     auth_svc = service.current_auth_service()
     if project_id:
+        if auth_svc.env.requires_auth:
+            projects = _list_projects(auth_svc)
+            if not next(
+                (project for project in projects if project.project_id == project_id),
+                None,
+            ):
+                raise click.ClickException(f"Project {project_id} not found")
         auth_svc.set_project(profile.name, project_id)
-        rprint(f"[green]Set active project to '{project_id}'[/green]")
+        rprint(f"Set active project to [bold green]{project_id}[/]")
         return
     if not interactive:
         raise click.ClickException(
             "No --project-id provided. Run `llamactl auth project --help` for more information."
         )
     try:
-        client = get_control_plane_client()
-        projects = asyncio.run(client.list_projects())
+        projects = _list_projects(auth_svc)
 
         if not projects:
             rprint("[yellow]No projects found[/yellow]")
@@ -208,8 +291,12 @@ def change_project(project_id: str | None, interactive: bool) -> None:
             project_id = questionary.text("Enter project ID").ask()
             result = project_id
         if result:
+            selected_project = next(
+                (project for project in projects if project.project_id == result), None
+            )
+            name = selected_project.project_name if selected_project else result
             auth_svc.set_project(profile.name, result)
-            rprint(f"[green]Set active project to '{result}'[/green]")
+            rprint(f"Set active project to [bold cornflower_blue]{name}[/]")
         else:
             rprint("[yellow]No project selected[/yellow]")
     except Exception as e:
@@ -217,6 +304,160 @@ def change_project(project_id: str | None, interactive: bool) -> None:
         raise click.Abort()
 
 
+def _auto_device_name() -> str:
+    try:
+        if sys.platform == "darwin":  # macOS
+            return (
+                subprocess.check_output(["scutil", "--get", "ComputerName"])
+                .decode()
+                .strip()
+            )
+        elif sys.platform.startswith("win"):
+            return os.environ["COMPUTERNAME"]
+        else:  # Linux / Unix
+            return platform.node()
+    except Exception:
+        return platform.node()
+
+
+async def _create_or_update_agent_api_key(auth_svc: AuthService, profile: Auth) -> None:
+    """
+    Mutates and updates the profile with an agent API key if it does not exist or is invalid.
+    """
+    if profile.api_key is not None:
+        async with PlatformAuthClient(profile.api_url, profile.api_key) as client:
+            try:
+                await client.me()
+            except httpx.HTTPStatusError as e:
+                if e.response.status_code == 401:
+                    # must have been deleted
+                    profile.api_key = None
+                    profile.api_key_id = None
+                else:
+                    raise
+    if profile.api_key is None:
+        async with auth_svc.profile_client(profile) as client:
+            name = f"{profile.name} llamactl on {profile.device_oidc.device_name if profile.device_oidc else 'unknown'}"
+            api_key = await client.create_agent_api_key(name)
+        profile.api_key = api_key.token
+        profile.api_key_id = api_key.id
+        auth_svc.update_profile(profile)
+
+
+def _create_device_profile() -> Auth:
+    auth_svc = service.current_auth_service()
+    if not auth_svc.env.requires_auth:
+        raise click.ClickException("This environment does not support authentication")
+
+    base_url = auth_svc.env.api_url.rstrip("/")
+
+    oidc_device = asyncio.run(_run_device_authentication(base_url))
+
+    # Obtain or prompt for project ID and create profile
+    projects = _list_projects(auth_svc, oidc_device.device_access_token)
+    selected_project_id = _select_or_enter_project(projects, True)
+    if not selected_project_id:
+        raise click.ClickException(
+            "Project is required. Does this user have access to any projects?"
+        )
+    created = auth_svc.create_or_update_profile_from_oidc(
+        selected_project_id, oidc_device
+    )
+    asyncio.run(_create_or_update_agent_api_key(auth_svc, created))
+    return created
+
+
+async def _run_device_authentication(base_url: str) -> DeviceOIDC:
+    device_name = _auto_device_name()
+    # 1) Discover upstream and CLI client_id via client
+    async with PlatformAuthDiscoveryClient(base_url) as discovery:
+        disc = await discovery.oidc_discovery()
+    upstream = disc.discovery_url
+    client_ids = disc.client_ids or {}
+    client_ids_list = list(client_ids.values())
+    client_id = client_ids.get("cli") or (
+        client_ids_list[0] if len(client_ids_list) == 1 else None
+    )
+    if not client_id:
+        raise click.ClickException(
+            "Expected 'cli' Client ID not found from auth discovery"
+        )
+
+    # 2) Device flow via typed OIDC client
+    async with OIDCClient() as oidc:
+        provider = await oidc.fetch_provider_configuration(upstream)
+        device_endpoint = provider.device_authorization_endpoint
+        token_endpoint = provider.token_endpoint
+        if not device_endpoint or not token_endpoint:
+            raise click.ClickException("Device Authorization not supported by provider")
+
+        scope_value = " ".join(sorted({"openid", "profile", "email", "offline_access"}))
+
+        # 3) Start device authorization
+        da = await oidc.device_authorization(
+            device_endpoint,
+            DeviceAuthorizationRequest(client_id=client_id, scope=scope_value),
+        )
+
+        rprint(
+            "[bold]complete authentication by visiting the verification URI and confirming the device:[/bold]"
+        )
+        if da.verification_uri:
+            rprint(
+                f"Verification URI: {da.verification_uri} (will open in your browser if supported)"
+            )
+        if da.user_code:
+            rprint(f"User Code: {da.user_code} to confirm the device")
+        if da.verification_uri_complete:
+            try:
+                webbrowser.open(da.verification_uri_complete)
+            except Exception:
+                pass
+
+        # 4) Poll token endpoint
+        interval = int(da.interval or 5)
+        while True:
+            await asyncio.sleep(interval)
+            token = await oidc.token_with_device_code(
+                token_endpoint,
+                TokenRequestDeviceCode(
+                    device_code=da.device_code,
+                    client_id=client_id,
+                ),
+            )
+            if token.error in {"authorization_pending", "slow_down"}:
+                if token.error == "slow_down":
+                    interval += 5
+                continue
+            if token.error:
+                raise click.ClickException(
+                    f"Token polling failed: {token.error} {token.error_description or ''}"
+                )
+            if token.id_token:
+                if not token.access_token:
+                    raise click.ClickException(
+                        "Device flow failed: token response missing access_token"
+                    )
+                claims = await decode_jwt_claims(token.id_token, provider.jwks_uri)
+                email = claims.get("email")
+                if not email:
+                    raise click.ClickException(
+                        "Device flow failed: email not found in token"
+                    )
+                user_id = claims.get("sub") or email
+                return DeviceOIDC(
+                    device_name=device_name,
+                    email=email,
+                    user_id=user_id,
+                    client_id=client_id,
+                    discovery_url=upstream,
+                    device_access_token=token.access_token,
+                    device_refresh_token=token.refresh_token,
+                    device_id_token=token.id_token,
+                )
+            raise click.ClickException("Device flow failed: unexpected token response")
+
+
 def validate_authenticated_profile(interactive: bool) -> Auth:
     """Validate that the user is authenticated within the current environment.
 
@@ -259,7 +500,7 @@ def validate_authenticated_profile(interactive: bool) -> Auth:
     # No profiles exist for this env
     if current_env.requires_auth:
         # Inline token flow
-        created = _token_flow_for_env(auth_svc)
+        created = _create_device_profile()
         return created
     else:
         # No auth required: select project and create a default profile without token
@@ -276,28 +517,44 @@ def validate_authenticated_profile(interactive: bool) -> Auth:
 
 
 def _prompt_for_api_key() -> str:
-    entered = questionary.password("Enter API key token").ask()
+    entered = questionary.password("Enter API key token to login").ask()
     if entered:
         return entered.strip()
     raise click.ClickException("No API key entered")
 
 
-def _validate_token_and_list_projects(
-    auth_svc: AuthService, api_key: str
+def _list_projects(
+    auth_svc: AuthService,
+    api_key: str | None = None,
 ) -> list[ProjectSummary]:
     async def _run():
-        async with ControlPlaneClient.ctx(auth_svc.env.api_url, api_key) as client:
+        profile = auth_svc.get_current_profile()
+        async with ControlPlaneClient.ctx(
+            auth_svc.env.api_url,
+            api_key or (profile.api_key if profile else None),
+            None if api_key is not None else auth_svc.auth_middleware(profile),
+        ) as client:
             return await client.list_projects()
 
+    return asyncio.run(_run())
+
+
+def _prompt_validate_api_key_and_list_projects(
+    auth_svc: AuthService, api_key: str
+) -> list[ProjectSummary]:
     try:
-        return asyncio.run(_run())
-    except ClientError as e:
-        if getattr(e, "status_code", None) == 401:
+        return _list_projects(auth_svc, api_key)
+    except httpx.HTTPStatusError as e:
+        if e.response.status_code == 401:
             rprint("[red]Invalid API key. Please try again.[/red]")
-            return _validate_token_and_list_projects(auth_svc, _prompt_for_api_key())
-        if getattr(e, "status_code", None) == 403:
+            return _prompt_validate_api_key_and_list_projects(
+                auth_svc, _prompt_for_api_key()
+            )
+        if e.response.status_code == 403:
             rprint("[red]This environment requires a valid API key.[/red]")
-            return _validate_token_and_list_projects(auth_svc, _prompt_for_api_key())
+            return _prompt_validate_api_key_and_list_projects(
+                auth_svc, _prompt_for_api_key()
+            )
         raise
     except Exception as e:
         raise click.ClickException(f"Failed to validate API key: {e}")
@@ -327,7 +584,7 @@ def _select_or_enter_project(
 
 def _token_flow_for_env(auth_service: AuthService) -> Auth:
     token_value = _prompt_for_api_key()
-    projects = _validate_token_and_list_projects(auth_service, token_value)
+    projects = _prompt_validate_api_key_and_list_projects(auth_service, token_value)
     project_id = _select_or_enter_project(projects, auth_service.env.requires_auth)
     if not project_id:
         raise click.ClickException("No project selected")

llama_deploy/cli/commands/deployment.py

@@ -225,15 +225,15 @@ def refresh_deployment(deployment_id: str | None, interactive: bool) -> None:
     """Update the deployment, pulling the latest code from it's branch"""
     validate_authenticated_profile(interactive)
     try:
-        client = get_project_client()
-
         deployment_id = select_deployment(deployment_id)
         if not deployment_id:
             rprint("[yellow]No deployment selected[/yellow]")
             return
 
         # Get current deployment details to show what we're refreshing
-        current_deployment = asyncio.run(client.get_deployment(deployment_id))
+        current_deployment = asyncio.run(
+            get_project_client().get_deployment(deployment_id)
+        )
         deployment_name = current_deployment.name
         old_git_sha = current_deployment.git_sha or ""
 
@@ -241,7 +241,7 @@ def refresh_deployment(deployment_id: str | None, interactive: bool) -> None:
         with console.status(f"Refreshing {deployment_name}..."):
             deployment_update = DeploymentUpdate()
             updated_deployment = asyncio.run(
-                client.update_deployment(
+                get_project_client().update_deployment(
                     deployment_id,
                     deployment_update,
                 )
@@ -278,29 +278,22 @@ def select_deployment(
     if not interactive:
         return None
 
-    try:
-        client = get_project_client()
-        deployments = asyncio.run(client.list_deployments())
+    client = get_project_client()
+    deployments = asyncio.run(client.list_deployments())
 
-        if not deployments:
-            rprint(
-                f"[yellow]No deployments found for project {client.project_id}[/yellow]"
-            )
-            return None
+    if not deployments:
+        rprint(f"[yellow]No deployments found for project {client.project_id}[/yellow]")
+        return None
 
-        choices = []
-        for deployment in deployments:
-            name = deployment.name
-            deployment_id = deployment.id
-            status = deployment.status
-            choices.append(
-                questionary.Choice(
-                    title=f"{name} ({deployment_id}) - {status}", value=deployment_id
-                )
+    choices = []
+    for deployment in deployments:
+        name = deployment.name
+        deployment_id = deployment.id
+        status = deployment.status
+        choices.append(
+            questionary.Choice(
+                title=f"{name} ({deployment_id}) - {status}", value=deployment_id
             )
+        )
 
-        return questionary.select("Select deployment:", choices=choices).ask()
-
-    except Exception as e:
-        rprint(f"[red]Error loading deployments: {e}[/red]")
-        return None
+    return questionary.select("Select deployment:", choices=choices).ask()

llama_deploy/cli/commands/serve.py

@@ -1,3 +1,4 @@
+import logging
 import os
 from pathlib import Path
 
@@ -20,6 +21,8 @@ from rich import print as rprint
 
 from ..app import app
 
+logger = logging.getLogger(__name__)
+
 
 @app.command(
     "serve",
@@ -97,7 +100,7 @@ def serve(
         )
 
     except KeyboardInterrupt:
-        print("Shutting down...")
+        logger.debug("Shutting down...")
 
     except Exception as e:
         rprint(f"[red]Error: {e}[/red]")
@@ -107,17 +110,19 @@ def serve(
 def _set_env_vars_from_profile(profile: Auth):
     if profile.api_key:
         _set_env_vars(profile.api_key, profile.api_url)
+    _set_project_id_from_profile(profile)
 
 
 def _set_env_vars_from_env(env_vars: dict[str, str]):
     key = env_vars.get("LLAMA_CLOUD_API_KEY")
     url = env_vars.get("LLAMA_CLOUD_BASE_URL", "https://api.cloud.llamaindex.ai")
+    # Also propagate project id if present in the environment
+    _set_project_id_from_env(env_vars)
     if key:
         _set_env_vars(key, url)
 
 
 def _set_env_vars(key: str, url: str):
-    print(f"Setting env vars: {key}, {url}")
     os.environ["LLAMA_CLOUD_API_KEY"] = key
     os.environ["LLAMA_CLOUD_BASE_URL"] = url
     # kludge for common web servers to inject local auth key
@@ -126,6 +131,17 @@ def _set_env_vars(key: str, url: str):
         os.environ[f"{prefix}LLAMA_CLOUD_BASE_URL"] = url
 
 
+def _set_project_id_from_env(env_vars: dict[str, str]):
+    project_id = env_vars.get("LLAMA_DEPLOY_PROJECT_ID")
+    if project_id:
+        os.environ["LLAMA_DEPLOY_PROJECT_ID"] = project_id
+
+
+def _set_project_id_from_profile(profile: Auth):
+    if profile.project_id:
+        os.environ["LLAMA_DEPLOY_PROJECT_ID"] = profile.project_id
+
+
 def _maybe_inject_llama_cloud_credentials(
     deployment_file: Path, interactive: bool
 ) -> None:
@@ -155,6 +171,9 @@ def _maybe_inject_llama_cloud_credentials(
         config, deployment_file.parent if deployment_file.is_file() else deployment_file
     )
 
+    # Ensure project id is available to the app and UI processes
+    _set_project_id_from_env({**os.environ, **vars})
+
     existing = os.environ.get("LLAMA_CLOUD_API_KEY") or vars.get("LLAMA_CLOUD_API_KEY")
     if existing:
         _set_env_vars_from_env({**os.environ, **vars})
@@ -176,18 +195,14 @@ def _maybe_inject_llama_cloud_credentials(
     # No key available; consider prompting if interactive
     if interactive:
         should_login = questionary.confirm(
-            "This deployment requires Llama Cloud. Login now to inject credentials?",
+            "This deployment requires Llama Cloud. Login now to inject credentials? Otherwise the app may not work.",
             default=True,
         ).ask()
         if should_login:
-            try:
-                authed = validate_authenticated_profile(True)
-                if authed.api_key:
-                    _set_env_vars_from_profile(authed)
-                    return
-            except Exception:
-                # fall through to warning
-                pass
+            authed = validate_authenticated_profile(True)
+            if authed.api_key:
+                _set_env_vars_from_profile(authed)
+                return
         rprint(
             "[yellow]Warning: No Llama Cloud credentials configured. The app may not work.[/yellow]"
         )
@@ -195,5 +210,5 @@ def _maybe_inject_llama_cloud_credentials(
 
     # Non-interactive session
     rprint(
-        "[yellow]Warning: LLAMA_CLOUD_API_KEY is not set and no logged-in profile was found. The deployment may not work.[/yellow]"
+        "[yellow]Warning: LLAMA_CLOUD_API_KEY is not set and no logged-in profile was found. The app may not work.[/yellow]"
     )