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

agentstack_cli/commands/platform/base_driver.py

@@ -3,9 +3,11 @@
 
 import abc
 import importlib.resources
+import json
 import pathlib
 import shlex
 import typing
+from enum import StrEnum
 from subprocess import CompletedProcess
 from textwrap import dedent
 
@@ -17,6 +19,13 @@ from agentstack_cli.configuration import Configuration
 from agentstack_cli.utils import merge, run_command
 
 
+class ImagePullMode(StrEnum):
+    guest = "guest"
+    host = "host"
+    hybrid = "hybrid"
+    skip = "skip"
+
+
 class BaseDriver(abc.ABC):
     vm_name: str
 
@@ -54,6 +63,34 @@ class BaseDriver(abc.ABC):
     @abc.abstractmethod
     async def exec(self, command: list[str]) -> None: ...
 
+    def _canonify(self, tag: str) -> str:
+        return tag if "." in tag.split("/")[0] else f"docker.io/{tag}"
+
+    async def _grab_image_shas(
+        self,
+        *,
+        mode: typing.Literal["guest", "host"],
+    ) -> dict[str, str]:
+        return {
+            tag: sha
+            for line in (
+                await run_command(
+                    ["docker", "images", "--digests"],
+                    "Listing host images",
+                )
+                if mode == "host"
+                else await self.run_in_vm(
+                    ["k3s", "ctr", "image", "ls"],
+                    "Listing guest images",
+                )
+            )
+            .stdout.decode()
+            .splitlines()[1:]
+            if (x := line.split())
+            and (sha := x[2])
+            and ((tag := self._canonify((x[0] + ":" + x[1]) if mode == "host" else x[0])) in self.loaded_images)
+        }
+
     async def install_tools(self) -> None:
         # Configure k3s registry for local registry access
         registry_config = dedent(
@@ -103,10 +140,7 @@ class BaseDriver(abc.ABC):
         self,
         set_values_list: list[str],
         values_file: pathlib.Path | None = None,
-        import_images: list[str] | None = None,
-        pull_on_host: bool = False,
-        skip_pull: bool = False,
-        skip_restart_deployments: bool = False,
+        image_pull_mode: ImagePullMode = ImagePullMode.guest,
     ) -> None:
         _ = await self.run_in_vm(
             ["sh", "-c", "mkdir -p /tmp/agentstack && cat >/tmp/agentstack/chart.tgz"],
@@ -138,37 +172,42 @@ class BaseDriver(abc.ABC):
             input=yaml.dump(values).encode("utf-8"),
         )
 
-        images_str = (
-            await self.run_in_vm(
-                [
-                    "/bin/bash",
-                    "-c",
-                    "helm template agentstack /tmp/agentstack/chart.tgz --values=/tmp/agentstack/values.yaml "
-                    + " ".join(shlex.quote(f"--set={value}") for value in set_values_list)
-                    + " | sed -n '/^\\s*image:/{ /{{/!{ s/.*image:\\s*//p } }'",
-                ],
-                "Listing necessary images",
+        self.loaded_images = {
+            self._canonify(typing.cast(str, yaml.safe_load(line)))
+            for line in (
+                await self.run_in_vm(
+                    [
+                        "/bin/bash",
+                        "-c",
+                        "helm template agentstack /tmp/agentstack/chart.tgz --values=/tmp/agentstack/values.yaml "
+                        + " ".join(shlex.quote(f"--set={value}") for value in set_values_list)
+                        + " | sed -n '/^\\s*image:/{ /{{/!{ s/.*image:\\s*//p } }'",
+                    ],
+                    "Listing necessary images",
+                )
             )
-        ).stdout.decode()
-
-        def canonify(tag: str) -> str:
-            return tag if "." in tag.split("/")[0] else f"docker.io/{tag}"
+            .stdout.decode()
+            .splitlines()
+        }
 
-        required_images = {canonify(typing.cast(str, yaml.safe_load(line))) for line in images_str.splitlines()}
-        images_to_import = {canonify(tag) for tag in import_images or []}
-        images_to_pull = required_images - images_to_import
+        images_to_import_from_host = set[str]()
+        shas_guest_before = dict[str, str]()
 
-        if not skip_pull:
-            if pull_on_host:
+        if image_pull_mode in {ImagePullMode.host, ImagePullMode.hybrid}:
+            shas_guest_before = await self._grab_image_shas(mode="guest")
+            shas_host = await self._grab_image_shas(mode="host")
+            if image_pull_mode == ImagePullMode.host and (images_to_pull := self.loaded_images - shas_host.keys()):
                 for image in images_to_pull:
-                    await run_command(["docker", "pull", image], f"Pulling image {image} on host")
-                images_to_import = required_images
-                images_to_pull = set[str]()
+                    await run_command(
+                        ["docker", "pull", image],
+                        f"Pulling image {image} on host",
+                    )
+                shas_host = await self._grab_image_shas(mode="host")
+            images_to_import_from_host = dict(shas_host.items() - shas_guest_before.items()).keys() & self.loaded_images
+            await self.import_images(*images_to_import_from_host)
 
-            if images_to_import:
-                await self.import_images(*images_to_import)
-
-            for image in images_to_pull:
+        if image_pull_mode in {ImagePullMode.guest, ImagePullMode.hybrid}:
+            for image in self.loaded_images - images_to_import_from_host:
                 async for attempt in AsyncRetrying(stop=stop_after_attempt(5)):
                     with attempt:
                         attempt_num = attempt.retry_state.attempt_number
@@ -176,10 +215,6 @@ class BaseDriver(abc.ABC):
                             ["k3s", "ctr", "image", "pull", image],
                             f"Pulling image {image}" + (f" (attempt {attempt_num})" if attempt_num > 1 else ""),
                         )
-        elif images_to_import:
-            await self.import_images(*images_to_import)
-
-        self.loaded_images = required_images
 
         kubeconfig_path = anyio.Path(Configuration().lima_home) / self.vm_name / "copied-from-guest" / "kubeconfig.yaml"
         await kubeconfig_path.parent.mkdir(parents=True, exist_ok=True)
@@ -210,8 +245,35 @@ class BaseDriver(abc.ABC):
             "Deploying Agent Stack platform with Helm",
         )
 
-        if import_images and not skip_restart_deployments:
-            await self.run_in_vm(
-                ["k3s", "kubectl", "rollout", "restart", "deployment"],
-                "Restarting deployments to load imported images",
-            )
+        if shas_guest_before and (
+            replaced_digests := set(shas_guest_before.values())
+            - set((await self._grab_image_shas(mode="guest")).values())
+        ):
+            for pod in dict.get(
+                json.loads(
+                    (
+                        await self.run_in_vm(
+                            ["k3s", "kubectl", "get", "pods", "-o", "json", "--all-namespaces"],
+                            "Getting pods",
+                        )
+                    ).stdout
+                ),
+                "items",
+                [],
+            ):
+                if any(
+                    container_status.get("imageID", "") in replaced_digests
+                    for container_status in pod.get("status", {}).get("containerStatuses", [])
+                ):
+                    await self.run_in_vm(
+                        [
+                            "k3s",
+                            "kubectl",
+                            "delete",
+                            "pod",
+                            pod["metadata"]["name"],
+                            "-n",
+                            pod["metadata"]["namespace"],
+                        ],
+                        f"Removing pod with obsolete image {pod['metadata']['namespace']}/{pod['metadata']['name']}",
+                    )

agentstack_cli/commands/platform/lima_driver.py

@@ -181,6 +181,8 @@ class LimaDriver(BaseDriver):
 
     @typing.override
     async def import_images(self, *tags: str):
+        if not tags:
+            return
         image_dir = anyio.Path("/tmp/agentstack")
         await image_dir.mkdir(exist_ok=True, parents=True)
         image_file = str(uuid.uuid4())

agentstack_cli/commands/platform/wsl_driver.py

@@ -13,7 +13,7 @@ import anyio
 import pydantic
 import yaml
 
-from agentstack_cli.commands.platform.base_driver import BaseDriver
+from agentstack_cli.commands.platform.base_driver import BaseDriver, ImagePullMode
 from agentstack_cli.configuration import Configuration
 from agentstack_cli.console import console
 from agentstack_cli.utils import run_command
@@ -130,13 +130,10 @@ class WSLDriver(BaseDriver):
         self,
         set_values_list: list[str],
         values_file: pathlib.Path | None = None,
-        import_images: list[str] | None = None,
-        pull_on_host: bool = False,
-        skip_pull: bool = False,
-        skip_restart_deployments: bool = False,
+        image_pull_mode: ImagePullMode = ImagePullMode.guest,
     ) -> None:
-        if pull_on_host:
-            raise NotImplementedError("Pulling on host is not supported on this platform.")
+        if image_pull_mode in {ImagePullMode.host, ImagePullMode.hybrid}:
+            raise NotImplementedError("Importing host images is not supported on Windows.")
 
         host_ip = (
             (
@@ -162,7 +159,7 @@ class WSLDriver(BaseDriver):
                 }
             ).encode(),
         )
-        await super().deploy(set_values_list=set_values_list, values_file=values_file, import_images=import_images)
+        await super().deploy(set_values_list=set_values_list, values_file=values_file, image_pull_mode=image_pull_mode)
         await self.run_in_vm(
             ["sh", "-c", "cat >/etc/systemd/system/kubectl-port-forward@.service"],
             "Installing systemd unit for port-forwarding",

agentstack_cli/commands/self.py

@@ -130,7 +130,7 @@ async def install(
             ).execute_async()
         ):
             try:
-                await agentstack_cli.commands.platform.start(set_values_list=[], import_images=[], verbose=verbose)
+                await agentstack_cli.commands.platform.start(set_values_list=[], verbose=verbose)
                 already_started = True
                 console.print()
             except Exception:
@@ -190,7 +190,7 @@ async def upgrade(
             "Upgrading agentstack-cli",
             env={"PATH": _path()},
         )
-        await agentstack_cli.commands.platform.start(set_values_list=[], import_images=[], verbose=verbose)
+        await agentstack_cli.commands.platform.start(set_values_list=[], verbose=verbose)
         await version(verbose=verbose)
 
 

agentstack_cli/commands/server.py

@@ -13,6 +13,7 @@ import httpx
 import typer
 import uvicorn
 from authlib.common.security import generate_token
+from authlib.oauth2.rfc6749.errors import InvalidGrantError, OAuth2Error
 from authlib.oauth2.rfc7636 import create_s256_code_challenge
 from fastapi import FastAPI, Request
 from fastapi.responses import HTMLResponse
@@ -104,11 +105,19 @@ async def server_login(server: typing.Annotated[str | None, typer.Argument()] =
 
     server = server.rstrip("/")
 
+    log_in_message = "No authentication tokens found for this server. Proceeding to log in."
+
     if server_data := config.auth_manager.get_server(server):
-        console.info("Switching to an already logged in server.")
+        console.info("Logging in to an already logged in server.")
         auth_server = None
         auth_servers = list(server_data.authorization_servers.keys())
-        if len(auth_servers) == 1:
+        if not auth_servers:
+            # Known server with no auth servers - save and exit
+            config.auth_manager.active_server = server
+            config.auth_manager.active_auth_server = auth_server
+            console.success(f"Logged in to [cyan]{server}[/cyan].")
+            return
+        elif len(auth_servers) == 1:
             auth_server = auth_servers[0]
         elif len(auth_servers) > 1:
             auth_server = await inquirer.select(  #  type: ignore
@@ -127,129 +136,106 @@ async def server_login(server: typing.Annotated[str | None, typer.Argument()] =
             if not auth_server:
                 console.info("Action cancelled.")
                 sys.exit(1)
-    else:
-        console.info("No authentication tokens found for this server. Proceeding to log in.")
-        async with httpx.AsyncClient() as client:
-            resp = await client.get(f"{server}/.well-known/oauth-protected-resource/", follow_redirects=True)
-            if resp.is_error:
-                console.error("This server does not appear to run a compatible version of Agent Stack Platform.")
-                sys.exit(1)
-            metadata = resp.json()
 
-        auth_servers = metadata.get("authorization_servers", [])
-        auth_server = None
-        token = None
+        # Validate that the token is still valid by attempting to load it
+        # Keep the original active server/auth server in case of failure
+        previous_server = config.auth_manager.active_server
+        previous_auth_server = config.auth_manager.active_auth_server
 
-        client_id = config.client_id
-        client_secret = config.client_secret
-        registration_token = None
+        config.auth_manager.active_server = server
+        config.auth_manager.active_auth_server = auth_server
 
-        if auth_servers:
-            if len(auth_servers) == 1:
-                auth_server = auth_servers[0]
+        try:
+            token = await config.auth_manager.load_auth_token()
+            if not token:
+                # No token available, need to log in
+                # Restore previous state until login completes
+                config.auth_manager.active_server = previous_server
+                config.auth_manager.active_auth_server = previous_auth_server
+                # Fall through to login flow below
             else:
-                auth_server = await inquirer.select(  # type: ignore
-                    message="Select an authorization server:",
-                    choices=auth_servers,
-                ).execute_async()
-
-            if not auth_server:
-                raise RuntimeError("No authorization server selected.")
-
-            async with httpx.AsyncClient() as client:
-                try:
-                    resp = await client.get(f"{auth_server}/.well-known/openid-configuration")
-                    resp.raise_for_status()
-                    oidc = resp.json()
-                except Exception as e:
-                    raise RuntimeError(f"OIDC discovery failed: {e}") from e
-
-            registration_endpoint = oidc["registration_endpoint"]
-            if not client_id and registration_endpoint:
-                async with httpx.AsyncClient() as client:
-                    resp = None
-                    try:
-                        app_name = get_unique_app_name()
-                        resp = await client.post(
-                            registration_endpoint,
-                            json={
-                                "client_name": app_name,
-                                "grant_types": ["authorization_code", "refresh_token"],
-                                "enforce_pkce": True,
-                                "all_users_entitled": True,
-                                "redirect_uris": [REDIRECT_URI],
-                            },
-                        )
-                        resp.raise_for_status()
-                        data = resp.json()
-                        client_id = data["client_id"]
-                        client_secret = data["client_secret"]
-                        registration_token = data["registration_access_token"]
-                    except Exception as e:
-                        if resp:
-                            try:
-                                error_details = resp.json()
-                                console.warning(
-                                    f"error: {error_details['error']} error description: {error_details['error_description']}"
-                                )
-
-                            except Exception:
-                                console.info("no parsable json response.")
-                        console.warning(f" Dynamic client registration failed. Proceed with manual input.  {e!s}")
-
-            if not client_id:
-                client_id = (
-                    await inquirer.text(  #  type: ignore
-                        message="Enter Client ID (default agentstack-cli):",
-                        instruction=f"(Redirect URI: {REDIRECT_URI})",
-                    ).execute_async()
-                    or "agentstack-cli"
-                )
-                if not client_id:
-                    raise RuntimeError("Client ID is mandatory. Action cancelled.")
-                client_secret = (
-                    await inquirer.text(  #  type: ignore
-                        message="Enter Client Secret (optional):"
-                    ).execute_async()
-                    or None
-                )
-
-            code_verifier = generate_token(64)
-
-            auth_url = f"{oidc['authorization_endpoint']}?{
-                urlencode(
-                    {
-                        'client_id': client_id,
-                        'response_type': 'code',
-                        'redirect_uri': REDIRECT_URI,
-                        'scope': ' '.join(metadata.get('scopes_supported', ['openid', 'email', 'profile'])),
-                        'code_challenge': typing.cast(str, create_s256_code_challenge(code_verifier)),
-                        'code_challenge_method': 'S256',
-                    }
-                )
-            }"
+                console.success(f"Logged in to [cyan]{server}[/cyan].")
+                return
+        except InvalidGrantError:
+            # Token refresh failed due to invalid/expired refresh token
+            log_in_message = "Your session has expired. Please log in again."
+            # Restore previous state until login completes
+            config.auth_manager.active_server = previous_server
+            config.auth_manager.active_auth_server = previous_auth_server
+            # Fall through to login flow below
+        except OAuth2Error as e:
+            # Other OAuth2 protocol errors - report but don't continue
+            console.error(f"OAuth2 error: {e.description}")
+            config.auth_manager.active_server = previous_server
+            config.auth_manager.active_auth_server = previous_auth_server
+            sys.exit(1)
+        except RuntimeError as e:
+            # Network or OIDC discovery errors - report but don't continue
+            console.error(f"Failed to validate authentication: {e}")
+            console.hint("Check your network connection and try again.")
+            config.auth_manager.active_server = previous_server
+            config.auth_manager.active_auth_server = previous_auth_server
+            sys.exit(1)
+
+    # Starting the login flow
+    console.info(log_in_message)
+
+    async with httpx.AsyncClient() as client:
+        resp = await client.get(f"{server}/.well-known/oauth-protected-resource/", follow_redirects=True)
+        if resp.is_error:
+            console.error("This server does not appear to run a compatible version of Agent Stack Platform.")
+            sys.exit(1)
+        metadata = resp.json()
+
+    auth_servers = metadata.get("authorization_servers", [])
+    auth_server = None
+    token = None
+
+    client_id = config.client_id
+    client_secret = config.client_secret
+    registration_token = None
+
+    if auth_servers:
+        if len(auth_servers) == 1:
+            auth_server = auth_servers[0]
+        else:
+            auth_server = await inquirer.select(  # type: ignore
+                message="Select an authorization server:",
+                choices=auth_servers,
+            ).execute_async()
 
-            console.info(f"Opening browser for login: [cyan]{auth_url}[/cyan]")
-            if not webbrowser.open(auth_url):
-                console.warning("Could not open browser. Please visit the above URL manually.")
+        if not auth_server:
+            raise RuntimeError("No authorization server selected.")
 
-            code = await _wait_for_auth_code()
+        async with httpx.AsyncClient() as client:
+            try:
+                resp = await client.get(f"{auth_server}/.well-known/openid-configuration")
+                resp.raise_for_status()
+                oidc = resp.json()
+            except Exception as e:
+                raise RuntimeError(f"OIDC discovery failed: {e}") from e
+
+        registration_endpoint = oidc["registration_endpoint"]
+        if not client_id and registration_endpoint:
             async with httpx.AsyncClient() as client:
-                token_resp = None
+                resp = None
                 try:
-                    token_resp = await client.post(
-                        oidc["token_endpoint"],
-                        data={
-                            "grant_type": "authorization_code",
-                            "code": code,
-                            "redirect_uri": REDIRECT_URI,
-                            "client_id": client_id,
-                            "client_secret": client_secret,
-                            "code_verifier": code_verifier,
+                    app_name = get_unique_app_name()
+                    resp = await client.post(
+                        registration_endpoint,
+                        json={
+                            "client_name": app_name,
+                            "grant_types": ["authorization_code", "refresh_token"],
+                            "enforce_pkce": True,
+                            "all_users_entitled": True,
+                            "redirect_uris": [REDIRECT_URI],
                         },
                     )
-                    token_resp.raise_for_status()
-                    token = token_resp.json()
+                    resp.raise_for_status()
+                    data = resp.json()
+                    client_id = data["client_id"]
+                    client_secret = data["client_secret"]
+                    registration_token = data["registration_access_token"]
                 except Exception as e:
                     if resp:
                         try:
@@ -257,22 +243,87 @@ async def server_login(server: typing.Annotated[str | None, typer.Argument()] =
                             console.warning(
                                 f"error: {error_details['error']} error description: {error_details['error_description']}"
                             )
+
                         except Exception:
                             console.info("no parsable json response.")
+                    console.warning(f" Dynamic client registration failed. Proceed with manual input.  {e!s}")
 
-                    raise RuntimeError(f"Token request failed: {e}") from e
+        if not client_id:
+            client_id = (
+                await inquirer.text(  #  type: ignore
+                    message="Enter Client ID:",
+                    instruction=f"(Redirect URI: {REDIRECT_URI})",
+                ).execute_async()
+                or "agentstack-cli"
+            )
+            if not client_id:
+                raise RuntimeError("Client ID is mandatory. Action cancelled.")
+            client_secret = (
+                await inquirer.text(  #  type: ignore
+                    message="Enter Client Secret (optional):"
+                ).execute_async()
+                or None
+            )
+
+        code_verifier = generate_token(64)
+
+        auth_url = f"{oidc['authorization_endpoint']}?{
+            urlencode(
+                {
+                    'client_id': client_id,
+                    'response_type': 'code',
+                    'redirect_uri': REDIRECT_URI,
+                    'scope': ' '.join(metadata.get('scopes_supported', ['openid', 'email', 'profile'])),
+                    'code_challenge': typing.cast(str, create_s256_code_challenge(code_verifier)),
+                    'code_challenge_method': 'S256',
+                }
+            )
+        }"
+
+        console.info(f"Opening browser for login: [cyan]{auth_url}[/cyan]")
+        if not webbrowser.open(auth_url):
+            console.warning("Could not open browser. Please visit the above URL manually.")
+
+        code = await _wait_for_auth_code()
+        async with httpx.AsyncClient() as client:
+            token_resp = None
+            try:
+                token_resp = await client.post(
+                    oidc["token_endpoint"],
+                    data={
+                        "grant_type": "authorization_code",
+                        "code": code,
+                        "redirect_uri": REDIRECT_URI,
+                        "client_id": client_id,
+                        "client_secret": client_secret,
+                        "code_verifier": code_verifier,
+                    },
+                )
+                token_resp.raise_for_status()
+                token = token_resp.json()
+            except Exception as e:
+                if token_resp:
+                    try:
+                        error_details = token_resp.json()
+                        console.warning(
+                            f"error: {error_details['error']} error description: {error_details['error_description']}"
+                        )
+                    except Exception:
+                        console.info("no parsable json response.")
 
-            if not token:
-                raise RuntimeError("Login timed out or not successful.")
-
-        config.auth_manager.save_auth_token(
-            server=server,
-            auth_server=auth_server,
-            client_id=client_id,
-            client_secret=client_secret,
-            token=token,
-            registration_token=registration_token,
-        )
+                raise RuntimeError(f"Token request failed: {e}") from e
+
+        if not token:
+            raise RuntimeError("Login timed out or not successful.")
+
+    config.auth_manager.save_auth_info(
+        server=server,
+        auth_server=auth_server,
+        client_id=client_id,
+        client_secret=client_secret,
+        token=token,
+        registration_token=registration_token,
+    )
 
     config.auth_manager.active_server = server
     config.auth_manager.active_auth_server = auth_server

agentstack_cli/commands/user.py

@@ -11,7 +11,7 @@ from rich.table import Column
 
 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
+from agentstack_cli.server_utils import announce_server_action, confirm_server_action
 
 app = AsyncTyper()
 configuration = Configuration()

agentstack_cli/configuration.py

@@ -13,6 +13,7 @@ from contextlib import asynccontextmanager
 import pydantic
 import pydantic_settings
 from agentstack_sdk.platform import PlatformClient, use_platform_client
+from authlib.oauth2.rfc6749.errors import InvalidGrantError, OAuth2Error
 from pydantic import HttpUrl, SecretStr
 
 from agentstack_cli.auth_manager import AuthManager
@@ -64,9 +65,28 @@ class Configuration(pydantic_settings.BaseSettings):
                 "Run [green]agentstack platform start[/green] to start a local server, or [green]agentstack server login[/green] to connect to a remote one."
             )
             sys.exit(1)
+
+        re_auth_message = (
+            f"Run [green]agentstack server login {self.auth_manager.active_server}[/green] to re-authenticate."
+        )
+        try:
+            auth_token = await self.auth_manager.load_auth_token()
+        except InvalidGrantError:
+            console.error("Your session has expired.")
+            console.hint(re_auth_message)
+            sys.exit(1)
+        except OAuth2Error as e:
+            console.error(f"OAuth2 error: {e.description}")
+            console.hint(re_auth_message)
+            sys.exit(1)
+        except RuntimeError as e:
+            console.error(f"Failed to load authentication: {e}")
+            console.hint("Check your network connection and try again.")
+            sys.exit(1)
+
         async with use_platform_client(
             auth=(self.username, self.password.get_secret_value()) if self.password else None,
-            auth_token=await self.auth_manager.load_auth_token(),
+            auth_token=auth_token,
             base_url=self.auth_manager.active_server + "/",
         ) as client:
             yield client