agentstack-cli 0.6.0rc4__py3-none-any.whl → 0.6.1rc1__py3-none-any.whl

agentstack_cli/auth_manager.py

@@ -7,8 +7,13 @@ from collections import defaultdict
 from typing import Any
 
 import httpx
+from authlib.common.errors import AuthlibBaseError
+from authlib.integrations.httpx_client import AsyncOAuth2Client
+from authlib.oauth2.rfc6749.errors import InvalidGrantError, OAuth2Error
 from pydantic import BaseModel, Field
 
+TOKEN_EXPIRY_LEEWAY = 60  # seconds
+
 
 class AuthToken(BaseModel):
     access_token: str
@@ -44,6 +49,7 @@ class AuthManager:
     def __init__(self, config_path: pathlib.Path):
         self._auth_path = config_path
         self._auth = self._load()
+        self._oidc_cache: dict[str, dict[str, Any]] = {}
 
     def _load(self) -> Auth:
         if not self._auth_path.exists():
@@ -53,7 +59,63 @@ class AuthManager:
     def _save(self) -> None:
         self._auth_path.write_text(self._auth.model_dump_json(indent=2))
 
-    def save_auth_token(
+    async def _get_oidc_metadata(self, auth_server: str) -> dict[str, Any]:
+        """Fetch and cache OIDC metadata."""
+        if auth_server in self._oidc_cache:
+            return self._oidc_cache[auth_server]
+
+        async with httpx.AsyncClient() as client:
+            try:
+                resp = await client.get(f"{auth_server}/.well-known/openid-configuration")
+                resp.raise_for_status()
+                metadata = resp.json()
+                self._oidc_cache[auth_server] = metadata
+                return metadata
+            except Exception as e:
+                raise RuntimeError(f"OIDC discovery failed: {e}") from e
+
+    def _create_token_update_callback(self, server: str, auth_server: str):
+        """Create a callback that saves tokens when they're refreshed."""
+
+        def update_token(token: dict[str, Any]):
+            # Authlib calls this automatically when tokens are refreshed
+            # kwargs may include refresh_token and access_token but we don't need them
+            auth_config = self._auth.servers[server].authorization_servers[auth_server]
+            self.save_auth_info(
+                server=server,
+                auth_server=auth_server,
+                client_id=auth_config.client_id,
+                client_secret=auth_config.client_secret,
+                token=token,
+                registration_token=auth_config.registration_token,
+            )
+
+        return update_token
+
+    async def _get_oauth_client(self, server: str, auth_server: str) -> AsyncOAuth2Client:
+        """Create an OAuth2 client configured with current credentials."""
+        auth_config = self._auth.servers[server].authorization_servers[auth_server]
+
+        if not auth_config or not auth_config.token:
+            raise ValueError(f"No token found for {auth_server}")
+
+        metadata = await self._get_oidc_metadata(auth_server)
+
+        # Convert AuthToken to dict format authlib expects
+        token_dict = auth_config.token.model_dump(exclude_none=True)
+
+        client = AsyncOAuth2Client(
+            client_id=auth_config.client_id,
+            client_secret=auth_config.client_secret,
+            token_endpoint=metadata["token_endpoint"],
+            token=token_dict,
+            scope=token_dict.get("scope"),
+            update_token=self._create_token_update_callback(server, auth_server),
+        )
+
+        return client
+
+    def save_auth_info(
         self,
         server: str,
         auth_server: str | None = None,
@@ -63,7 +125,7 @@ class AuthManager:
         registration_token: str | None = None,
     ) -> None:
         if auth_server is not None and client_id is not None and token is not None:
-            if token["access_token"]:
+            if token["access_token"] and token.get("expires_in") is not None:
                 usetimestamp = int(time.time()) + int(token["expires_in"])
                 token["expires_at"] = usetimestamp
             self._auth.servers[server].authorization_servers[auth_server] = AuthServer(
@@ -78,57 +140,56 @@ class AuthManager:
 
     async def exchange_refresh_token(self, auth_server: str, token: AuthToken) -> dict[str, Any] | None:
         """
-        This method exchanges a refresh token for a new access token.
-        """
+        Exchange a refresh token for a new access token using authlib.
 
-        async with httpx.AsyncClient(headers={"Accept": "application/json"}) as client:
-            resp = None
-            try:
-                resp = await client.get(f"{auth_server}/.well-known/openid-configuration")
-                resp.raise_for_status()
-                oidc = resp.json()
-            except Exception as e:
-                if resp:
-                    error_details = resp.json()
-                    print(f"error: {error_details['error']} error description: {error_details['error_description']}")
-                raise RuntimeError(f"OIDC discovery failed: {e}") from e
-
-            token_endpoint = oidc["token_endpoint"]
-            try:
-                client_id = (
-                    self._auth.servers[self._auth.active_server or ""].authorization_servers[auth_server].client_id
-                )
-                client_secret = (
-                    self._auth.servers[self._auth.active_server or ""].authorization_servers[auth_server].client_secret
-                )
-                resp = await client.post(
-                    f"{token_endpoint}",
-                    data={
-                        "grant_type": "refresh_token",
-                        "refresh_token": token.refresh_token,
-                        "scope": token.scope,
-                        "client_id": client_id,
-                    }
-                    | ({"client_secret": client_secret} if client_secret else {}),
+        Raises:
+            InvalidGrantError: If the refresh token is invalid or expired (4xx auth errors)
+            OAuth2Error: For other OAuth2 protocol errors
+            RuntimeError: For network errors or OIDC discovery failures
+        """
+        if not self._auth.active_server:
+            raise ValueError("No active server configured")
+
+        try:
+            metadata = await self._get_oidc_metadata(auth_server)
+            token_endpoint = metadata["token_endpoint"]
+
+            async with await self._get_oauth_client(self._auth.active_server, auth_server) as client:
+                # Authlib's fetch_token with refresh_token grant automatically handles the refresh
+                # and calls update_token callback to save the new token
+                new_token = await client.fetch_token(
+                    url=token_endpoint,
+                    grant_type="refresh_token",
+                    refresh_token=token.refresh_token,
                 )
-                resp.raise_for_status()
-                new_token = resp.json()
-            except Exception as e:
-                if resp:
-                    error_details = resp.json()
-                    print(f"error: {error_details['error']} error description: {error_details['error_description']}")
-                raise RuntimeError(f"Failed to refresh token: {e}") from e
-            self.save_auth_token(
-                self._auth.active_server or "",
-                self._auth.active_auth_server or "",
-                self._auth.servers[self._auth.active_server or ""].authorization_servers[auth_server].client_id or "",
-                self._auth.servers[self._auth.active_server or ""].authorization_servers[auth_server].client_secret
-                or "",
-                token=new_token,
-            )
-            return new_token
+                return new_token
+        except InvalidGrantError as e:
+            # 400-level OAuth errors: invalid/expired refresh token
+            raise InvalidGrantError(
+                description=f"Token refresh failed - invalid or expired refresh token: {e.description}"
+            ) from e
+        except OAuth2Error as e:
+            # Other OAuth2 protocol errors
+            raise OAuth2Error(description=f"OAuth2 error during token refresh: {e.description}") from e
+        except AuthlibBaseError as e:
+            # Other authlib errors
+            raise RuntimeError(f"Token refresh failed: {e}") from e
+        except Exception as e:
+            # Network errors, OIDC discovery failures, etc.
+            raise RuntimeError(f"Failed to refresh token: {e}") from e
 
     async def load_auth_token(self) -> str | None:
+        """
+        Load and refresh auth token if needed using authlib.
+
+        Returns:
+            Access token string, or None if no auth configured
+
+        Raises:
+            InvalidGrantError: If token is expired and refresh fails due to auth issues (4xx)
+            OAuth2Error: For other OAuth2 protocol errors
+            RuntimeError: For network or other errors
+        """
         active_res = self._auth.active_server
         active_auth_server = self._auth.active_auth_server
         if not active_res or not active_auth_server:
@@ -136,12 +197,12 @@ class AuthManager:
         server = self._auth.servers.get(active_res)
         if not server:
             return None
-
         auth_server = server.authorization_servers.get(active_auth_server)
         if not auth_server or not auth_server.token:
             return None
 
-        if (auth_server.token.expires_at or 0) - 60 < time.time():
+        if (auth_server.token.expires_at or 0) - TOKEN_EXPIRY_LEEWAY < time.time():
+            # Token expired, try to refresh - this may raise TokenRefreshError
             new_token = await self.exchange_refresh_token(active_auth_server, auth_server.token)
             if new_token:
                 return new_token["access_token"]
@@ -149,62 +210,53 @@ class AuthManager:
 
         return auth_server.token.access_token
 
-    async def deregister_client(self, auth_server, client_id, registration_token) -> None:
-        async with httpx.AsyncClient(headers={"Accept": "application/json"}) as client:
-            resp = None
-            try:
-                resp = await client.get(f"{auth_server}/.well-known/openid-configuration")
-                resp.raise_for_status()
-                oidc = resp.json()
-                registration_endpoint = oidc["registration_endpoint"]
-            except Exception as e:
-                if resp:
-                    error_details = resp.json()
-                    print(f"error: {error_details['error']} error description: {error_details['error_description']}")
-                raise RuntimeError(f"OIDC discovery failed: {e}") from e
+    async def deregister_client(self, auth_server: str, client_id: str | None, registration_token: str | None) -> None:
+        """Deregister a dynamically registered OAuth2 client."""
+        if not client_id or not registration_token:
+            return  # Nothing to deregister
 
-            try:
-                if client_id is not None and client_id != "" and registration_token is not None:
-                    headers = {"authorization": f"bearer {registration_token}"}
-                    resp = await client.delete(f"{registration_endpoint}/{client_id}", headers=headers)
-                    resp.raise_for_status()
+        try:
+            metadata = await self._get_oidc_metadata(auth_server)
+            registration_endpoint = metadata.get("registration_endpoint")
 
-            except Exception as e:
-                if resp:
-                    error_details = resp.json()
-                    print(f"error: {error_details['error']} error description: {error_details['error_description']}")
-                raise RuntimeError(f"Dynamic client de-registration failed. {e}") from e
+            if not registration_endpoint:
+                raise RuntimeError("Registration endpoint not found in OIDC metadata")
+
+            async with AsyncOAuth2Client() as client:
+                headers = {"Authorization": f"Bearer {registration_token}"}
+                resp = await client.delete(f"{registration_endpoint}/{client_id}", headers=headers)
+                resp.raise_for_status()
+
+        except Exception as e:
+            raise RuntimeError(f"Dynamic client de-registration failed: {e}") from e
 
     async def clear_auth_token(self, all: bool = False) -> None:
         if all:
             for server in self._auth.servers:
                 for auth_server in self._auth.servers[server].authorization_servers:
+                    auth_config = self._auth.servers[server].authorization_servers[auth_server]
                     await self.deregister_client(
                         auth_server,
-                        self._auth.servers[server].authorization_servers[auth_server].client_id,
-                        self._auth.servers[server].authorization_servers[auth_server].registration_token,
+                        auth_config.client_id,
+                        auth_config.registration_token,
                     )
 
             self._auth.servers = defaultdict(Server)
         else:
             if self._auth.active_server and self._auth.active_auth_server:
-                if (
-                    self._auth.servers[self._auth.active_server]
-                    .authorization_servers[self._auth.active_auth_server]
-                    .client_id
-                ):
-                    await self.deregister_client(
-                        self._auth.active_auth_server,
-                        self._auth.servers[self._auth.active_server]
-                        .authorization_servers[self._auth.active_auth_server]
-                        .client_id,
-                        self._auth.servers[self._auth.active_server]
-                        .authorization_servers[self._auth.active_auth_server]
-                        .registration_token,
-                    )
+                auth_config = self._auth.servers[self._auth.active_server].authorization_servers[
+                    self._auth.active_auth_server
+                ]
+                await self.deregister_client(
+                    self._auth.active_auth_server,
+                    auth_config.client_id,
+                    auth_config.registration_token,
+                )
                 del self._auth.servers[self._auth.active_server].authorization_servers[self._auth.active_auth_server]
+
             if self._auth.active_server and not self._auth.servers[self._auth.active_server].authorization_servers:
                 del self._auth.servers[self._auth.active_server]
+
         self._auth.active_server = None
         self._auth.active_auth_server = None
         self._save()

agentstack_cli/commands/agent.py

@@ -112,9 +112,8 @@ from rich.table import Column
 
 from agentstack_cli.api import a2a_client
 from agentstack_cli.async_typer import AsyncTyper, console, create_table, err_console
+from agentstack_cli.server_utils import announce_server_action, confirm_server_action
 from agentstack_cli.utils import (
-    announce_server_action,
-    confirm_server_action,
     generate_schema_example,
     is_github_url,
     parse_env_var,
@@ -181,6 +180,36 @@ processing_messages = [
 
 configuration = Configuration()
 
+DISCOVERY_TIMEOUT_SEC = 180
+DISCOVERY_POLL_INTERVAL_SEC = 2
+
+
+async def _discover_agent_card(docker_image: str) -> AgentCard:
+    from agentstack_sdk.platform.provider_discovery import DiscoveryState, ProviderDiscovery
+
+    console.info("Image missing agent card label, starting discovery...")
+
+    async with configuration.use_platform_client():
+        with status("Creating discovery task"):
+            discovery = await ProviderDiscovery.create(docker_image=docker_image)
+
+        start = asyncio.get_event_loop().time()
+        with status("Discovering agent card (this may take a while)"):
+            while discovery.status in (DiscoveryState.PENDING, DiscoveryState.IN_PROGRESS):
+                if asyncio.get_event_loop().time() - start > DISCOVERY_TIMEOUT_SEC:
+                    raise RuntimeError("Discovery timed out after 3 minutes")
+                await asyncio.sleep(DISCOVERY_POLL_INTERVAL_SEC)
+                await discovery.get()
+
+        if discovery.status == DiscoveryState.FAILED:
+            raise RuntimeError(f"Discovery failed: {discovery.error_message}")
+
+        card = discovery.agent_card
+        if not card:
+            raise RuntimeError("Discovery completed but no agent card was returned")
+
+        return card
+
 
 @app.command("add")
 async def add_agent(
@@ -257,9 +286,18 @@ async def add_agent(
             if dockerfile:
                 raise ValueError("Dockerfile can be specified only if location is a GitHub url")
             console.info(f"Assuming public docker image or network address, attempting to add {location}")
-            with status("Registering agent to platform"):
-                async with configuration.use_platform_client():
-                    await Provider.create(location=location)
+            try:
+                with status("Registering agent to platform"):
+                    async with configuration.use_platform_client():
+                        await Provider.create(location=location)
+            except httpx.HTTPStatusError as e:
+                if e.response.status_code == 422:
+                    agent_card = await _discover_agent_card(location)
+                    with status("Registering agent with discovered card"):
+                        async with configuration.use_platform_client():
+                            await Provider.create(location=location, agent_card=agent_card)
+                else:
+                    raise
         console.success(f"Agent [bold]{location}[/bold] added to platform")
         await list_agents()
 

agentstack_cli/commands/build.py

@@ -25,10 +25,9 @@ from tenacity import AsyncRetrying, retry_if_exception_type, stop_after_delay, w
 
 from agentstack_cli.async_typer import AsyncTyper
 from agentstack_cli.console import console, err_console
+from agentstack_cli.server_utils import announce_server_action, confirm_server_action
 from agentstack_cli.utils import (
-    announce_server_action,
     capture_output,
-    confirm_server_action,
     extract_messages,
     print_log,
     run_command,

agentstack_cli/commands/connector.py

@@ -1,6 +1,7 @@
 # Copyright 2025 © BeeAI a Series of LF Projects, LLC
 # SPDX-License-Identifier: Apache-2.0
 import asyncio
+import sys
 import typing
 
 import pydantic
@@ -14,7 +15,7 @@ from agentstack_cli import configuration
 from agentstack_cli.async_typer import AsyncTyper
 from agentstack_cli.configuration import Configuration
 from agentstack_cli.console import console
-from agentstack_cli.utils import (
+from agentstack_cli.server_utils import (
     announce_server_action,
     confirm_server_action,
 )
@@ -99,7 +100,7 @@ async def remove_connector(
         console.error(
             "[red]Cannot specify both --all and a search path. Use --all to remove all connectors, or provide a search path for specific connectors.[/red]"
         )
-        raise typer.Exit(1)
+        sys.exit(1)
 
     async with configuration.use_platform_client():
         connectors_list = await Connector.list()
@@ -191,7 +192,7 @@ async def select_connector(search_path: str) -> Connector | None:
     except ValueError as e:
         console.error(e.__str__())
         console.hint("Please refine your input to match exactly one connector id or url.")
-        raise typer.Exit(code=1) from None
+        sys.exit(1)
 
 
 @app.command("get")
@@ -258,7 +259,7 @@ async def disconnect(
         console.error(
             "[red]Cannot specify both --all and a search path. Use --all to remove all connectors, or provide a search path for specific connectors.[/red]"
         )
-        raise typer.Exit(1)
+        sys.exit(1)
 
     async with configuration.use_platform_client():
         connectors_list = await Connector.list()

agentstack_cli/commands/model.py

@@ -25,7 +25,8 @@ from rich.table import Column
 from agentstack_cli.api import openai_client
 from agentstack_cli.async_typer import AsyncTyper, console, create_table
 from agentstack_cli.configuration import Configuration
-from agentstack_cli.utils import announce_server_action, confirm_server_action, run_command, verbosity
+from agentstack_cli.server_utils import announce_server_action, confirm_server_action
+from agentstack_cli.utils import run_command, verbosity
 
 app = AsyncTyper()
 configuration = Configuration()

agentstack_cli/commands/platform/__init__.py

@@ -17,7 +17,7 @@ import typer
 from tenacity import AsyncRetrying, retry_if_exception_type, stop_after_delay, wait_fixed
 
 from agentstack_cli.async_typer import AsyncTyper
-from agentstack_cli.commands.platform.base_driver import BaseDriver
+from agentstack_cli.commands.platform.base_driver import BaseDriver, ImagePullMode
 from agentstack_cli.commands.platform.lima_driver import LimaDriver
 from agentstack_cli.commands.platform.wsl_driver import WSLDriver
 from agentstack_cli.configuration import Configuration
@@ -64,19 +64,20 @@ async def start(
     set_values_list: typing.Annotated[
         list[str], typer.Option("--set", help="Set Helm chart values using <key>=<value> syntax", default_factory=list)
     ],
-    import_images: typing.Annotated[
-        list[str],
+    image_pull_mode: typing.Annotated[
+        ImagePullMode,
         typer.Option(
-            "--import", help="Import an image from a local Docker CLI into Agent Stack platform", default_factory=list
+            "--image-pull-mode",
+            help=textwrap.dedent(
+                """\
+                guest = pull all images inside VM
+                host = pull unavailable images on host, then import all
+                hybrid = import available images from host, pull the rest in VM
+                skip = skip explicit pull step (Kubernetes will attempt to pull missing images)
+                """
+            ),
         ),
-    ],
-    pull_on_host: typing.Annotated[
-        bool,
-        typer.Option(
-            "--pull-on-host",
-            help="Pull images on host Docker daemon and import them instead of pulling inside the VM. Acts as a pull cache layer.",
-        ),
-    ] = False,
+    ] = ImagePullMode.guest,
     values_file: typing.Annotated[
         pathlib.Path | None, typer.Option("-f", help="Set Helm chart values using yaml values file")
     ] = None,
@@ -101,10 +102,7 @@ async def start(
         await driver.deploy(
             set_values_list=set_values_list,
             values_file=values_file_path,
-            import_images=import_images,
-            pull_on_host=pull_on_host,
-            skip_pull=skip_pull,
-            skip_restart_deployments=skip_restart_deployments,
+            image_pull_mode=image_pull_mode,
         )
 
         if not no_wait_for_platform: