llamactl 0.3.0a19__tar.gz → 0.3.0a21__tar.gz

{llamactl-0.3.0a19 → llamactl-0.3.0a21}/PKG-INFO

@@ -1,12 +1,12 @@
 Metadata-Version: 2.3
 Name: llamactl
-Version: 0.3.0a19
+Version: 0.3.0a21
 Summary: A command-line interface for managing LlamaDeploy projects and deployments
 Author: Adrian Lyjak
 Author-email: Adrian Lyjak <adrianlyjak@gmail.com>
 License: MIT
-Requires-Dist: llama-deploy-core[client]>=0.3.0a19,<0.4.0
-Requires-Dist: llama-deploy-appserver>=0.3.0a19,<0.4.0
+Requires-Dist: llama-deploy-core[client]>=0.3.0a21,<0.4.0
+Requires-Dist: llama-deploy-appserver>=0.3.0a21,<0.4.0
 Requires-Dist: httpx>=0.24.0,<1.0.0
 Requires-Dist: rich>=13.0.0
 Requires-Dist: questionary>=2.0.0
@@ -16,6 +16,8 @@ Requires-Dist: tenacity>=9.1.2
 Requires-Dist: textual>=6.0.0
 Requires-Dist: aiohttp>=3.12.14
 Requires-Dist: copier>=9.9.0
+Requires-Dist: pyjwt[crypto]>=2.10.1
+Requires-Dist: vibe-llama>=0.4.2,<0.5.0
 Requires-Python: >=3.11, <4
 Description-Content-Type: text/markdown
 

{llamactl-0.3.0a19 → llamactl-0.3.0a21}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "llamactl"
-version = "0.3.0a19"
+version = "0.3.0a21"
 description = "A command-line interface for managing LlamaDeploy projects and deployments"
 readme = "README.md"
 license = { text = "MIT" }
@@ -9,8 +9,8 @@ authors = [
 ]
 requires-python = ">=3.11, <4"
 dependencies = [
-    "llama-deploy-core[client]>=0.3.0a19,<0.4.0",
-    "llama-deploy-appserver>=0.3.0a19,<0.4.0",
+    "llama-deploy-core[client]>=0.3.0a21,<0.4.0",
+    "llama-deploy-appserver>=0.3.0a21,<0.4.0",
     "httpx>=0.24.0,<1.0.0",
     "rich>=13.0.0",
     "questionary>=2.0.0",
@@ -20,6 +20,8 @@ dependencies = [
     "textual>=6.0.0",
     "aiohttp>=3.12.14",
     "copier>=9.9.0",
+    "pyjwt[crypto]>=2.10.1",
+    "vibe-llama>=0.4.2,<0.5.0",
 ]
 
 [project.scripts]

llamactl-0.3.0a21/src/llama_deploy/cli/auth/client.py

@@ -0,0 +1,362 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+from types import TracebackType
+from typing import Any, AsyncContextManager, AsyncGenerator, Awaitable, Callable, Self
+
+import httpx
+import jwt
+from jwt.algorithms import RSAAlgorithm  # type: ignore[possibly-unbound-import]
+from llama_deploy.cli.config.schema import DeviceOIDC
+from pydantic import BaseModel
+
+logger = logging.getLogger(__name__)
+
+
+class OidcDiscoveryResponse(BaseModel):
+    discovery_url: str
+    client_ids: dict[str, str] | None = None
+
+
+class OidcProviderConfiguration(BaseModel):
+    device_authorization_endpoint: str | None = None
+    token_endpoint: str | None = None
+    scopes_supported: list[str] | None = None
+    jwks_uri: str | None = None
+
+
+class JsonWebKey(BaseModel):
+    kty: str
+    kid: str | None = None
+    use: str | None = None
+    alg: str | None = None
+    n: str | None = None
+    e: str | None = None
+    x5c: list[str] | None = None
+    x5t: str | None = None
+    x5t_s256: str | None = None
+
+
+class JsonWebKeySet(BaseModel):
+    keys: list[JsonWebKey]
+
+
+class AuthMeResponse(BaseModel):
+    id: str
+    email: str | None = None
+    last_login_provider: str | None = None
+    name: str | None = None
+    first_name: str | None = None
+    last_name: str | None = None
+    claims: dict[str, Any] | None = None
+    restrict: Any | None = None
+    created_at: str | None = None
+
+
+class ClientContextManager(AsyncContextManager):
+    def __init__(self, base_url: str | None, auth: httpx.Auth | None = None) -> None:
+        self.base_url = base_url.rstrip("/") if base_url else None
+        if self.base_url:
+            self.client = httpx.AsyncClient(base_url=self.base_url, auth=auth)
+        else:
+            self.client = httpx.AsyncClient(auth=auth)
+
+    async def close(self) -> None:
+        try:
+            await self.client.aclose()
+        except Exception:
+            pass
+
+    async def __aenter__(self) -> Self:
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        await self.close()
+
+
+class PlatformAuthDiscoveryClient(ClientContextManager):
+    """Client for ad hoc auth endpoints under /api/v1/auth."""
+
+    def __init__(self, base_url: str) -> None:
+        super().__init__(base_url)
+
+    async def oidc_discovery(self) -> OidcDiscoveryResponse:
+        resp = await self.client.get("/api/v1/auth/oidc/discovery", timeout=10.0)
+        resp.raise_for_status()
+        return OidcDiscoveryResponse.model_validate(resp.json())
+
+
+class APIToken(BaseModel):
+    token: str
+    id: str
+
+
+class PlatformAuthClient(ClientContextManager):
+    """Client for user introspection under /api/v1/auth/me."""
+
+    def __init__(
+        self, base_url: str, id_token: str | None = None, auth: httpx.Auth | None = None
+    ) -> None:
+        self.id_token = id_token
+        super().__init__(base_url, auth=auth)
+
+    async def me(self) -> AuthMeResponse:
+        headers = (
+            {"Authorization": f"Bearer {self.id_token}"} if self.id_token else None
+        )
+        resp = await self.client.get("/api/v1/auth/me", headers=headers, timeout=10.0)
+        resp.raise_for_status()
+        return AuthMeResponse.model_validate(resp.json())
+
+    async def create_agent_api_key(self, name: str) -> APIToken:
+        resp = await self.client.post(
+            "/api/v1/api-keys",
+            json={"name": name, "project_id": None},
+        )
+        resp.raise_for_status()
+        json = resp.json()
+        token = json["redacted_api_key"]
+        id = json["id"]
+        return APIToken(token=token, id=id)
+
+    async def delete_api_key(self, id: str) -> None:
+        response = await self.client.delete(f"/api/v1/api-keys/{id}")
+        response.raise_for_status()
+
+
+class RefreshMiddleware(httpx.Auth):
+    def __init__(
+        self,
+        device_oidc: DeviceOIDC,
+        on_refresh: Callable[[DeviceOIDC], Awaitable[None]],
+    ) -> None:
+        self.device_oidc = device_oidc
+        self.on_refresh = on_refresh
+        self.lock = asyncio.Lock()
+
+    async def _refresh_and_update(self) -> None:
+        new_device_oidc = await refresh(self.device_oidc)
+        self.device_oidc = new_device_oidc
+        try:
+            await self.on_refresh(new_device_oidc)
+        except Exception:
+            logger.exception("Error in on_refresh callback")
+
+    async def async_auth_flow(
+        self, request: httpx.Request
+    ) -> AsyncGenerator[httpx.Request, httpx.Response]:
+        token = self.device_oidc.device_access_token
+        request.headers["Authorization"] = f"Bearer {token}"
+
+        response = yield request
+        if response.status_code == 401:
+            async with self.lock:
+                if token == self.device_oidc.device_access_token:
+                    await self._refresh_and_update()
+                    request.headers["Authorization"] = (
+                        f"Bearer {self.device_oidc.device_access_token}"
+                    )
+                yield request
+
+
+class DeviceAuthorizationRequest(BaseModel):
+    client_id: str
+    scope: str
+
+
+class DeviceAuthorizationResponse(BaseModel):
+    device_code: str
+    user_code: str
+    verification_uri: str
+    verification_uri_complete: str | None = None
+    expires_in: int
+    interval: int | None = None
+
+
+class TokenRequestDeviceCode(BaseModel):
+    grant_type: str = "urn:ietf:params:oauth:grant-type:device_code"
+    device_code: str
+    client_id: str
+
+
+class TokenResponse(BaseModel):
+    # Success fields
+    id_token: str | None = None
+    access_token: str | None = None
+    refresh_token: str | None = None
+    expires_in: int | None = None
+    token_type: str | None = None
+    scope: str | None = None
+    # Error fields
+    error: str | None = None
+    error_description: str | None = None
+
+
+class TokenRequestRefresh(BaseModel):
+    grant_type: str = "refresh_token"
+    refresh_token: str
+    client_id: str
+
+
+class OIDCClient(ClientContextManager):
+    def __init__(self) -> None:
+        super().__init__(None)
+
+    async def fetch_provider_configuration(
+        self, discovery_url: str
+    ) -> OidcProviderConfiguration:
+        resp = await self.client.get(discovery_url, timeout=10.0)
+        resp.raise_for_status()
+        return OidcProviderConfiguration.model_validate(resp.json())
+
+    async def device_authorization(
+        self, device_endpoint: str, request: DeviceAuthorizationRequest
+    ) -> DeviceAuthorizationResponse:
+        resp = await self.client.post(
+            device_endpoint,
+            data=request.model_dump(),
+            headers={
+                "Accept": "application/json",
+                "Content-Type": "application/x-www-form-urlencoded",
+            },
+            timeout=10.0,
+        )
+        resp.raise_for_status()
+        return DeviceAuthorizationResponse.model_validate(resp.json())
+
+    async def token_with_device_code(
+        self, token_endpoint: str, request: TokenRequestDeviceCode
+    ) -> TokenResponse:
+        resp = await self.client.post(
+            token_endpoint,
+            data=request.model_dump(),
+            headers={
+                "Accept": "application/json",
+                "Content-Type": "application/x-www-form-urlencoded",
+            },
+            timeout=10.0,
+        )
+        # Do not raise for status; callers inspect error payloads during polling
+        try:
+            payload = resp.json()
+        except Exception:
+            # Fall back to minimal error information
+            return TokenResponse(error="invalid_response", error_description=resp.text)
+        return TokenResponse.model_validate(payload)
+
+    async def token_with_refresh(
+        self, token_endpoint: str, request: TokenRequestRefresh
+    ) -> TokenResponse:
+        resp = await self.client.post(
+            token_endpoint,
+            data=request.model_dump(),
+            headers={
+                "Accept": "application/json",
+                "Content-Type": "application/x-www-form-urlencoded",
+            },
+            timeout=10.0,
+        )
+        try:
+            payload = resp.json()
+        except Exception:
+            return TokenResponse(error="invalid_response", error_description=resp.text)
+        return TokenResponse.model_validate(payload)
+
+    async def get_jwks(self, jwks_uri: str) -> JsonWebKeySet:
+        resp = await self.client.get(jwks_uri, timeout=10.0)
+        resp.raise_for_status()
+        return JsonWebKeySet.model_validate(resp.json())
+
+
+async def decode_jwt_claims_from_device_oidc(
+    oidc_device: DeviceOIDC,
+    verify_audience: bool = False,
+    verify_expiration: bool = False,
+    audience: str | None = None,
+) -> dict[str, Any]:
+    """Decode JWT claims by discovering provider and verifying via JWKS.
+
+    Assumes RSA signing. Audience verification can be toggled and, when enabled,
+    an audience value can be provided.
+    """
+    if not oidc_device.device_id_token:
+        raise ValueError("Device ID token is missing. Cannot decode claims.")
+    async with OIDCClient() as oidc:
+        provider = await oidc.fetch_provider_configuration(oidc_device.discovery_url)
+        jwks_uri = provider.jwks_uri
+        if not jwks_uri:
+            raise ValueError("Provider does not expose jwks_uri")
+    return await decode_jwt_claims(
+        oidc_device.device_id_token,
+        jwks_uri,
+        verify_audience,
+        verify_expiration,
+        audience,
+    )
+
+
+async def decode_jwt_claims(
+    token: str,
+    jwks_uri: str,
+    verify_audience: bool = False,
+    verify_expiration: bool = False,
+    audience: str | None = None,
+) -> dict[str, Any]:
+    async with OIDCClient() as oidc:
+        jwks = await oidc.get_jwks(jwks_uri)
+
+    # Select key
+    header = jwt.get_unverified_header(token)
+    kid = header.get("kid")
+    alg = header.get("alg", "RS256")
+    keys = jwks.keys
+    key = next((k for k in keys if k.kid == kid), None) or next(iter(keys), None)
+    if not key:
+        raise ValueError("Signing key not found in JWKS")
+
+    # Build public key (RSA-only)
+    if key.kty != "RSA":
+        raise ValueError("Unsupported JWK kty; only RSA is supported")
+    key_json = key.model_dump_json()
+    public_key = RSAAlgorithm.from_jwk(key_json)
+
+    return jwt.decode(
+        token,
+        public_key,
+        algorithms=[alg],
+        options={"verify_aud": verify_audience, "verify_exp": verify_expiration},
+        audience=audience,
+    )
+
+
+async def refresh(device_oidc: DeviceOIDC) -> DeviceOIDC:
+    """
+    Run a refresh on the access token, storing updated tokens in a new DeviceOIDC.
+    """
+    async with OIDCClient() as oidc:
+        provider = await oidc.fetch_provider_configuration(device_oidc.discovery_url)
+        token_endpoint = provider.token_endpoint
+        if not token_endpoint:
+            raise ValueError("Provider does not expose token_endpoint")
+        if not device_oidc.device_refresh_token:
+            raise ValueError("Device refresh token is missing. Cannot refresh.")
+        token = await oidc.token_with_refresh(
+            token_endpoint,
+            TokenRequestRefresh(
+                refresh_token=device_oidc.device_refresh_token,
+                client_id=device_oidc.client_id,
+            ),
+        )
+        copy = device_oidc.model_copy()
+        if not token.access_token:
+            raise ValueError("Refresh failed: token response missing access_token")
+        copy.device_access_token = token.access_token
+        copy.device_refresh_token = token.refresh_token or copy.device_refresh_token
+        copy.device_id_token = token.id_token or copy.device_id_token
+        return copy

{llamactl-0.3.0a19 → llamactl-0.3.0a21}/src/llama_deploy/cli/client.py

@@ -7,26 +7,35 @@ from rich import print as rprint
 
 
 def get_control_plane_client() -> ControlPlaneClient:
+    auth_svc = service.current_auth_service()
     profile = service.current_auth_service().get_current_profile()
     if profile:
         resolved_base_url = profile.api_url.rstrip("/")
         resolved_api_key = profile.api_key
-        return ControlPlaneClient(resolved_base_url, resolved_api_key)
+        return ControlPlaneClient(
+            resolved_base_url, resolved_api_key, auth_svc.auth_middleware()
+        )
 
     # Fallback: allow env-scoped client construction for env operations
     env = service.get_current_environment()
     resolved_base_url = env.api_url.rstrip("/")
-    return ControlPlaneClient(resolved_base_url, None)
+    return ControlPlaneClient(resolved_base_url)
 
 
 def get_project_client() -> ProjectClient:
-    profile = service.current_auth_service().get_current_profile()
+    auth_svc = service.current_auth_service()
+    profile = auth_svc.get_current_profile()
     if not profile:
         rprint("\n[bold red]No profile configured![/bold red]")
         rprint("\nTo get started, create a profile with:")
-        rprint("[cyan]llamactl auth token[/cyan]")
+        if auth_svc.env.requires_auth:
+            rprint("[cyan]llamactl auth login[/cyan]")
+        else:
+            rprint("[cyan]llamactl auth token[/cyan]")
         raise SystemExit(1)
-    return ProjectClient(profile.api_url, profile.project_id, profile.api_key)
+    return ProjectClient(
+        profile.api_url, profile.project_id, profile.api_key, auth_svc.auth_middleware()
+    )
 
 
 @asynccontextmanager