vellum-ai 0.14.41__py3-none-any.whl → 1.11.2__py3-none-any.whl

vellum_cli/push.py

@@ -8,15 +8,16 @@ from uuid import UUID
 from typing import List, Optional
 
 from dotenv import load_dotenv
+from pydantic import ValidationError
 
 from vellum.client.core.api_error import ApiError
+from vellum.client.types import WorkflowPushDeploymentConfigRequest
 from vellum.resources.workflows.client import OMIT
-from vellum.types import WorkflowPushDeploymentConfigRequest
 from vellum.workflows.vellum_client import create_vellum_client
-from vellum.workflows.workflows.base import BaseWorkflow
 from vellum_cli.config import DEFAULT_WORKSPACE_CONFIG, WorkflowConfig, WorkflowDeploymentConfig, load_vellum_cli_config
-from vellum_cli.logger import load_cli_logger
-from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
+from vellum_cli.logger import handle_cli_error, load_cli_logger
+from vellum_ee.workflows.display.nodes.utils import to_kebab_case
+from vellum_ee.workflows.display.workflows.base_workflow_display import BaseWorkflowDisplay
 
 
 def push_command(
@@ -27,6 +28,7 @@ def push_command(
     deployment_name: Optional[str] = None,
     deployment_description: Optional[str] = None,
     release_tags: Optional[List[str]] = None,
+    release_description: Optional[str] = None,
     dry_run: Optional[bool] = None,
     strict: Optional[bool] = None,
     workspace: Optional[str] = None,
@@ -45,8 +47,11 @@ def push_command(
         else config.workflows
     )
 
-    if len(workflow_configs) > 1 and workspace:
-        workflow_configs = [w for w in workflow_configs if w.workspace == workspace]
+    if len(workflow_configs) > 1:
+        if workspace:
+            workflow_configs = [w for w in workflow_configs if w.workspace == workspace]
+        else:
+            workflow_configs = [w for w in workflow_configs if w.workspace == DEFAULT_WORKSPACE_CONFIG.name]
 
     if len(workflow_configs) == 0:
         if module and module_exists(module):
@@ -58,22 +63,45 @@ def push_command(
             config.workflows.append(new_config)
             workflow_configs = [new_config]
         elif not module:
-            raise ValueError("No Workflows found in project to push.")
+            handle_cli_error(
+                logger,
+                title="No workflows found in project to push",
+                message="No workflow configurations were found in vellum.lock.json.",
+                suggestion="If you pulled a workflow with a specific sandbox ID"
+                "(e.g., vellum workflows pull --workflow-sandbox-id=<id>), "
+                "you should push without specifying a different sandbox ID:\n\n"
+                "  vellum workflows push\n\n"
+                "Or, if you want to push to a specific module:\n\n"
+                "  vellum workflows push --module=<module-name>",
+            )
+            return
         else:
             raise ValueError(f"No workflow config for '{module}' found in project to push.")
 
     if len(workflow_configs) > 1:
-        raise ValueError("Multiple workflows found in project to push. Pushing only a single workflow is supported.")
+        raise ValueError(
+            "Multiple workflows found. Please specify a single workflow to push. "
+            f"Found: {', '.join([w.module for w in workflow_configs])}"
+        )
 
     workflow_config = workflow_configs[0]
 
     logger.info(f"Loading workflow from {workflow_config.module}")
     resolved_workspace = workspace or workflow_config.workspace or DEFAULT_WORKSPACE_CONFIG.name
-    workspace_config = (
-        next((w for w in config.workspaces if w.name == resolved_workspace), DEFAULT_WORKSPACE_CONFIG)
-        if workspace
-        else DEFAULT_WORKSPACE_CONFIG
-    )
+    workspace_config = next((w for w in config.workspaces if w.name == resolved_workspace), None)
+
+    if workspace_config is None:
+        if resolved_workspace == DEFAULT_WORKSPACE_CONFIG.name:
+            workspace_config = DEFAULT_WORKSPACE_CONFIG
+        else:
+            available_workspaces = sorted(set([w.name for w in config.workspaces] + [DEFAULT_WORKSPACE_CONFIG.name]))
+            handle_cli_error(
+                logger,
+                title=f"Workspace '{resolved_workspace}' not found in config",
+                message=f"Available workspaces: {', '.join(available_workspaces)}",
+            )
+            return
+
     api_key = os.getenv(workspace_config.api_key)
     if not api_key:
         raise ValueError(f"No API key value found in environment for workspace '{workspace_config.name}'.")
@@ -91,30 +119,33 @@ def push_command(
         and workflow_config.workflow_sandbox_id
         and workflow_config.workflow_sandbox_id != workflow_sandbox_id
     ):
-        raise ValueError(
-            f"Workflow sandbox id '{workflow_sandbox_id}' is already associated with '{workflow_config.module}'."
+        logger.warning(
+            f"Workflow sandbox id '{workflow_sandbox_id}' is already associated with '{workflow_config.module}'. "
+            f"Continuing with the provided workflow sandbox id '{workflow_sandbox_id}'."
         )
 
     client = create_vellum_client(
         api_key=api_key,
+        api_url=workspace_config.api_url,
     )
     sys.path.insert(0, os.getcwd())
 
     # Remove this once we could serialize using the artifact in Vembda
     # https://app.shortcut.com/vellum/story/5585
-    workflow = BaseWorkflow.load_from_module(workflow_config.module)
-    workflow_display = get_workflow_display(
-        workflow_class=workflow,
+    serialization_result = BaseWorkflowDisplay.serialize_module(
+        workflow_config.module,
+        client=client,
         dry_run=dry_run or False,
     )
-    exec_config = workflow_display.serialize()
+    exec_config = serialization_result.exec_config
 
     container_tag = workflow_config.container_image_tag
     if workflow_config.container_image_name and not workflow_config.container_image_tag:
         container_tag = "latest"
 
     exec_config["runner_config"] = {
-        "sdk_version": metadata.version("vellum-ai"),
+        "sdk_version": metadata.version("vellum-ai"),  # Deprecated in favor of codegen_version
+        "codegen_version": metadata.version("vellum-ai"),
         "container_image_tag": container_tag,
         "container_image_name": workflow_config.container_image_name,
     }
@@ -126,12 +157,28 @@ def push_command(
             workflow_config.deployments[0] if workflow_config.deployments else WorkflowDeploymentConfig()
         )
 
-        deployment_config = WorkflowPushDeploymentConfigRequest(
-            label=deployment_label or cli_deployment_config.label,
-            name=deployment_name or cli_deployment_config.name,
-            description=deployment_description or cli_deployment_config.description,
-            release_tags=release_tags or cli_deployment_config.release_tags,
-        )
+        try:
+            module_name = workflow_config.module.split(".")[-1]
+            deployment_config = WorkflowPushDeploymentConfigRequest(
+                label=deployment_label or cli_deployment_config.label,
+                name=deployment_name or cli_deployment_config.name or to_kebab_case(module_name),
+                description=deployment_description or cli_deployment_config.description,
+                release_tags=release_tags or cli_deployment_config.release_tags,
+                release_description=release_description or cli_deployment_config.release_description,
+            )
+        except ValidationError as e:
+            for error in e.errors():
+                if "release_tags" in str(error.get("loc", [])):
+                    handle_cli_error(
+                        logger,
+                        title="Invalid release tag format",
+                        message="Release tags must be provided as separate arguments. "
+                        "Use: --release-tag tag1 --release-tag tag2",
+                    )
+                    return
+
+            # Re-raise if it's not a release_tags validation error
+            raise e
 
         # We should check with fern if we could auto-serialize typed fields for us
         # https://app.shortcut.com/vellum/story/5568
@@ -159,12 +206,14 @@ def push_command(
     artifact.seek(0)
     artifact.name = f"{workflow_config.module.replace('.', '__')}.tar.gz"
 
+    provided_id = workflow_config.workflow_sandbox_id or workflow_sandbox_id
+
     try:
         response = client.workflows.push(
             # Remove this once we could serialize using the artifact in Vembda
             # https://app.shortcut.com/vellum/story/5585
             exec_config=json.dumps(exec_config),
-            workflow_sandbox_id=workflow_config.workflow_sandbox_id or workflow_sandbox_id,
+            workflow_sandbox_id=provided_id,
             artifact=artifact,
             # We should check with fern if we could auto-serialize typed object fields for us
             # https://app.shortcut.com/vellum/story/5568
@@ -173,41 +222,57 @@ def push_command(
             strict=strict,
         )
     except ApiError as e:
-        if e.status_code != 400 or not isinstance(e.body, dict) or "diffs" not in e.body:
-            raise e
-
-        diffs: dict = e.body["diffs"]
-        generated_only = diffs.get("generated_only", [])
-        generated_only_str = (
-            "\n".join(
-                ["Files that were generated but not found in the original project:"]
-                + [f"- {file}" for file in generated_only]
+        if e.status_code == 404:
+            if provided_id:
+                handle_cli_error(
+                    logger,
+                    title="Workflow Sandbox not found",
+                    message=f"Could not find Workflow Sandbox with ID '{provided_id}' "
+                    f"in workspace '{resolved_workspace}'.",
+                )
+            else:
+                error_detail = e.body.get("detail") if isinstance(e.body, dict) else None
+                default_message = "The `/workflows/push` endpoint failed with a 404 response."
+                handle_cli_error(
+                    logger,
+                    title="Workflow Sandbox not found",
+                    message=error_detail or default_message,
+                )
+            return
+
+        if e.status_code == 400 and isinstance(e.body, dict) and "diffs" in e.body:
+            diffs: dict = e.body["diffs"]
+            generated_only = diffs.get("generated_only", [])
+            generated_only_str = (
+                "\n".join(
+                    ["Files that were generated but not found in the original project:"]
+                    + [f"- {file}" for file in generated_only]
+                )
+                if generated_only
+                else ""
             )
-            if generated_only
-            else ""
-        )
 
-        original_only = diffs.get("original_only", [])
-        original_only_str = (
-            "\n".join(
-                ["Files that were found in the original project but not generated:"]
-                + [f"- {file}" for file in original_only]
+            original_only = diffs.get("original_only", [])
+            original_only_str = (
+                "\n".join(
+                    ["Files that were found in the original project but not generated:"]
+                    + [f"- {file}" for file in original_only]
+                )
+                if original_only
+                else ""
             )
-            if original_only
-            else ""
-        )
 
-        modified = diffs.get("modified", {})
-        modified_str = (
-            "\n\n".join(
-                ["Files that were different between the original project and the generated artifact:"]
-                + ["\n".join(line.strip() for line in lines) for lines in modified.values()]
+            modified = diffs.get("modified", {})
+            modified_str = (
+                "\n\n".join(
+                    ["Files that were different between the original project and the generated artifact:"]
+                    + ["\n".join(line.strip() for line in lines) for lines in modified.values()]
+                )
+                if modified
+                else ""
             )
-            if modified
-            else ""
-        )
 
-        reported_diffs = f"""\
+            reported_diffs = f"""\
 {e.body.get("detail")}
 
 {generated_only_str}
@@ -216,11 +281,17 @@ def push_command(
 
 {modified_str}
 """
-        logger.error(reported_diffs)
-        return
+            logger.error(reported_diffs)
+            return
+
+        if e.status_code == 400 and isinstance(e.body, dict) and "detail" in e.body:
+            handle_cli_error(logger, title="API request to /workflows/push failed.", message=e.body["detail"])
+            return
+
+        raise e
 
     if dry_run:
-        error_messages = [str(e) for e in workflow_display.errors]
+        error_messages = [error.message for error in serialization_result.errors]
         error_message = "\n".join(error_messages) if error_messages else "No errors found."
         logger.info(
             f"""\
@@ -233,17 +304,34 @@ def push_command(
 {json.dumps(response.proposed_diffs, indent=2)}
 """
         )  # type: ignore[attr-defined]
+
+        error_list = serialization_result.errors
+        has_errors = len(error_list) > 0
+        has_diffs = response.proposed_diffs is not None and response.proposed_diffs
+
+        if has_errors or has_diffs:
+            exit(1)
     else:
+        default_api_url = client._client_wrapper._environment.default
+        base_url = default_api_url.split("/v1")[0].replace("//api.", "//app.")
         logger.info(
             f"""Successfully pushed {workflow_config.module} to Vellum!
-Visit at: https://app.vellum.ai/workflow-sandboxes/{response.workflow_sandbox_id}"""
+Visit at: {base_url}/workflow-sandboxes/{response.workflow_sandbox_id}"""
         )
 
     if not workflow_config.workflow_sandbox_id:
         workflow_config.workflow_sandbox_id = response.workflow_sandbox_id
 
     if not workflow_config.deployments and response.workflow_deployment_id:
-        workflow_config.deployments.append(WorkflowDeploymentConfig(id=UUID(response.workflow_deployment_id)))
+        stored_deployment_config = WorkflowDeploymentConfig(
+            id=UUID(response.workflow_deployment_id),
+            label=deployment_config.label if deploy else None,
+            name=deployment_config.name if deploy else None,
+            description=deployment_config.description if deploy else None,
+            release_tags=deployment_config.release_tags if deploy else None,
+            release_description=getattr(deployment_config, "release_description", None) if deploy else None,
+        )
+        workflow_config.deployments.append(stored_deployment_config)
 
     config.save()
     logger.info("Updated vellum.lock.json file.")

vellum_cli/tests/conftest.py

@@ -58,3 +58,9 @@ def mock_module(request) -> Generator[MockModuleResult, None, None]:
 
     os.chdir(current_dir)
     shutil.rmtree(temp_dir)
+
+
+@pytest.fixture
+def info_log_level(monkeypatch):
+    """Set log level to INFO for tests that request this fixture"""
+    monkeypatch.setenv("LOG_LEVEL", "INFO")

vellum_cli/tests/test_image_push.py

@@ -1,22 +1,54 @@
+import pytest
+import json
+import os
+import shutil
 import subprocess
-from unittest.mock import MagicMock, patch
+import tempfile
+from unittest.mock import MagicMock
+from uuid import uuid4
+from typing import Generator
 
 from click.testing import CliRunner
+from httpx import Response
 
+from vellum.client.types.docker_service_token import DockerServiceToken
 from vellum_cli import main as cli_main
 
 
-@patch("subprocess.run")
-@patch("docker.from_env")
-def test_image_push__self_hosted_happy_path(mock_docker_from_env, mock_run, vellum_client, monkeypatch):
+@pytest.fixture
+def mock_temp_dir() -> Generator[str, None, None]:
+    current_dir = os.getcwd()
+    temp_dir = tempfile.mkdtemp()
+    os.chdir(temp_dir)
+
+    yield temp_dir
+
+    os.chdir(current_dir)
+    shutil.rmtree(temp_dir)
+
+
+@pytest.fixture
+def mock_docker_from_env(mocker):
+    return mocker.patch("docker.from_env")
+
+
+@pytest.fixture
+def mock_subprocess_run(mocker):
+    return mocker.patch("subprocess.run")
+
+
+@pytest.mark.usefixtures("vellum_client", "info_log_level")
+def test_image_push__self_hosted_happy_path(mock_docker_from_env, mock_subprocess_run, monkeypatch):
     # GIVEN a self hosted vellum api URL env var
     monkeypatch.setenv("VELLUM_API_URL", "mycompany.api.com")
+    monkeypatch.setenv("VELLUM_API_KEY", "123456abcdef")
 
     # Mock Docker client
     mock_docker_client = MagicMock()
     mock_docker_from_env.return_value = mock_docker_client
 
-    mock_run.side_effect = [
+    mock_subprocess_run.side_effect = [
+        subprocess.CompletedProcess(args="", returncode=0, stdout=b"Pruning successful"),
         subprocess.CompletedProcess(
             args="", returncode=0, stdout=b'{"manifests": [{"platform": {"architecture": "amd64"}}]}'
         ),
@@ -35,9 +67,96 @@ def test_image_push__self_hosted_happy_path(mock_docker_from_env, mock_run, vell
     assert "Image successfully pushed" in result.output
 
 
-@patch("subprocess.run")
-@patch("docker.from_env")
-def test_image_push__self_hosted_blocks_repo(mock_docker_from_env, mock_run, vellum_client, monkeypatch):
+@pytest.mark.usefixtures("info_log_level")
+def test_image_push__self_hosted_happy_path__workspace_option(
+    mock_docker_from_env, mock_subprocess_run, mock_httpx_transport, mock_temp_dir
+):
+    # GIVEN a workspace config with a new env for url
+    with open(os.path.join(mock_temp_dir, "vellum.lock.json"), "w") as f:
+        f.write(
+            json.dumps(
+                {
+                    "workspaces": [
+                        {
+                            "name": "my_workspace",
+                            "api_url": "MY_WORKSPACE_VELLUM_API_URL",
+                            "api_key": "MY_WORKSPACE_VELLUM_API_KEY",
+                        }
+                    ]
+                }
+            )
+        )
+
+    # AND a .env file with the workspace api key and url
+    with open(os.path.join(mock_temp_dir, ".env"), "w") as f:
+        f.write(
+            "VELLUM_API_KEY=123456abcdef\n"
+            "VELLUM_API_URL=https://api.vellum.ai\n"
+            "MY_WORKSPACE_VELLUM_API_KEY=789012ghijkl\n"
+            "MY_WORKSPACE_VELLUM_API_URL=https://api.vellum.mycompany.ai\n"
+        )
+
+    # AND the Docker client returns the correct response
+    mock_docker_client = MagicMock()
+    mock_docker_from_env.return_value = mock_docker_client
+
+    mock_subprocess_run.side_effect = [
+        subprocess.CompletedProcess(args="", returncode=0, stdout=b"Pruning successful"),
+        subprocess.CompletedProcess(
+            args="", returncode=0, stdout=b'{"manifests": [{"platform": {"architecture": "amd64"}}]}'
+        ),
+        subprocess.CompletedProcess(args="", returncode=0, stdout=b"sha256:hellosha"),
+    ]
+
+    # AND the vellum client returns the correct response for
+    mock_httpx_transport.handle_request.side_effect = [
+        # First call to get the docker service token
+        Response(
+            status_code=200,
+            text=json.dumps(
+                {
+                    "access_token": "345678mnopqr",
+                    "organization_id": str(uuid4()),
+                    "repository": "myrepo.net",
+                }
+            ),
+        ),
+        # Second call to push the image
+        Response(
+            status_code=200,
+            text=json.dumps(
+                {
+                    "id": str(uuid4()),
+                    "name": "myrepo.net/myimage",
+                    "visibility": "PRIVATE",
+                    "created": "2021-01-01T00:00:00Z",
+                    "modified": "2021-01-01T00:00:00Z",
+                    "repository": "myrepo.net",
+                    "sha": "sha256:hellosha",
+                    "tags": [],
+                }
+            ),
+        ),
+    ]
+
+    # WHEN the user runs the image push command
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["image", "push", "myrepo.net/myimage:latest", "--workspace", "my_workspace"])
+
+    # THEN the command exits successfully
+    assert result.exit_code == 0, (result.output, str(result.exception))
+
+    # AND gives the success message
+    assert "Image successfully pushed" in result.output
+
+    # AND the vellum client was called with the correct api key and url
+    request = mock_httpx_transport.handle_request.call_args[0][0]
+    assert request.headers["X-API-KEY"] == "789012ghijkl", result.stdout
+    assert str(request.url) == "https://api.vellum.mycompany.ai/v1/container-images/push"
+
+
+@pytest.mark.usefixtures("vellum_client", "mock_subprocess_run")
+def test_image_push__self_hosted_blocks_repo(mock_docker_from_env, monkeypatch):
     # GIVEN a self hosted vellum api URL env var
     monkeypatch.setenv("VELLUM_API_URL", "mycompany.api.com")
 
@@ -54,3 +173,171 @@ def test_image_push__self_hosted_blocks_repo(mock_docker_from_env, mock_run, vel
 
     # AND gives the error message for self hosted installs not including the repo
     assert "For adding images to your self hosted install you must include" in result.output
+
+
+@pytest.mark.usefixtures("info_log_level")
+def test_image_push_with_source_success(
+    mock_docker_from_env, mock_subprocess_run, vellum_client, monkeypatch, mock_temp_dir
+):
+    monkeypatch.setenv("VELLUM_API_URL", "https://api.vellum.ai")
+    monkeypatch.setenv("VELLUM_API_KEY", "123456abcdef")
+
+    dockerfile_path = os.path.join(mock_temp_dir, "Dockerfile")
+    with open(dockerfile_path, "w") as f:
+        f.write("FROM alpine:latest\n")
+
+    mock_docker_client = MagicMock()
+    mock_docker_from_env.return_value = mock_docker_client
+    mock_docker_client.images.push.return_value = [b'{"status": "Pushed"}']
+
+    mock_subprocess_run.side_effect = [
+        subprocess.CompletedProcess(args="", returncode=0, stdout=b"Pruning successful"),
+        subprocess.CompletedProcess(args="", returncode=0, stdout=b"Build successful"),
+        subprocess.CompletedProcess(
+            args="", returncode=0, stdout=b'{"manifests": [{"platform": {"architecture": "amd64"}}]}'
+        ),
+        subprocess.CompletedProcess(args="", returncode=0, stdout=b"sha256:hellosha"),
+    ]
+
+    vellum_client.container_images.docker_service_token.return_value = DockerServiceToken(
+        access_token="345678mnopqr", organization_id="test-org", repository="myrepo.net"
+    )
+
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["image", "push", "myimage:latest", "--source", dockerfile_path])
+
+    assert result.exit_code == 0, result.output
+
+    build_call = mock_subprocess_run.call_args_list[1]
+    assert build_call[0][0] == [
+        "docker",
+        "buildx",
+        "build",
+        "-f",
+        dockerfile_path,
+        "--platform=linux/amd64",
+        "-t",
+        "myimage:latest",
+        ".",
+    ]
+
+    assert "Docker build completed successfully" in result.output
+    assert "Image successfully pushed" in result.output
+
+
+@pytest.mark.usefixtures("mock_docker_from_env", "mock_subprocess_run", "vellum_client")
+def test_image_push_with_source_dockerfile_not_exists(monkeypatch, mock_temp_dir):
+    monkeypatch.setenv("VELLUM_API_URL", "https://api.vellum.ai")
+    monkeypatch.setenv("VELLUM_API_KEY", "123456abcdef")
+
+    nonexistent_dockerfile = os.path.join(mock_temp_dir, "nonexistent_dockerfile")
+
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["image", "push", "myimage:latest", "--source", nonexistent_dockerfile])
+
+    assert result.exit_code == 1
+    assert "Dockerfile does not exist" in result.output
+
+
+@pytest.mark.usefixtures("mock_docker_from_env", "vellum_client")
+def test_image_push_with_source_build_fails(mock_subprocess_run, monkeypatch, mock_temp_dir):
+    monkeypatch.setenv("VELLUM_API_URL", "https://api.vellum.ai")
+    monkeypatch.setenv("VELLUM_API_KEY", "123456abcdef")
+
+    dockerfile_path = os.path.join(mock_temp_dir, "Dockerfile")
+    with open(dockerfile_path, "w") as f:
+        f.write("FROM alpine:latest\n")
+
+    mock_subprocess_run.side_effect = [
+        subprocess.CompletedProcess(args="", returncode=0, stdout=b"Pruning successful"),
+        subprocess.CompletedProcess(args="", returncode=1, stderr=b"Build failed: missing dependency"),
+    ]
+
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["image", "push", "myimage:latest", "--source", dockerfile_path])
+
+    assert result.exit_code == 1
+    assert "Docker build failed" in result.output
+    assert "Build failed: missing dependency" in result.output
+
+
+@pytest.mark.usefixtures("vellum_client", "info_log_level")
+def test_image_push_includes_docker_prune(mock_docker_from_env, mock_subprocess_run, monkeypatch):
+    """
+    Tests that image_push_command calls docker image prune before validation.
+    """
+    # GIVEN a self hosted vellum api URL env var
+    monkeypatch.setenv("VELLUM_API_URL", "mycompany.api.com")
+    monkeypatch.setenv("VELLUM_API_KEY", "123456abcdef")
+
+    # AND a mock Docker client
+    mock_docker_client = MagicMock()
+    mock_docker_from_env.return_value = mock_docker_client
+
+    mock_subprocess_run.side_effect = [
+        subprocess.CompletedProcess(args="", returncode=0, stdout=b"Pruning successful"),
+        subprocess.CompletedProcess(
+            args="", returncode=0, stdout=b'{"manifests": [{"platform": {"architecture": "amd64"}}]}'
+        ),
+        subprocess.CompletedProcess(args="", returncode=0, stdout=b"manifest"),
+        subprocess.CompletedProcess(args="", returncode=0, stdout=b"sha256:hellosha"),
+    ]
+
+    # WHEN the user runs the image push command with any image
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["image", "push", "myrepo.net/myimage:latest"])
+
+    # THEN the command exits successfully
+    assert result.exit_code == 0, result.output
+
+    prune_call = mock_subprocess_run.call_args_list[0]
+    assert prune_call[0][0] == [
+        "docker",
+        "image",
+        "prune",
+        "--all",
+        "-f",
+        "--filter",
+        "label=image-type=python-workflow-runtime",
+    ]
+
+    # AND the success message includes pruning completion
+    assert "Docker image pruning completed successfully" in result.output
+
+
+@pytest.mark.usefixtures("vellum_client", "info_log_level")
+def test_image_push_continues_if_prune_fails(mock_docker_from_env, mock_subprocess_run, monkeypatch):
+    """
+    Tests that image_push_command continues if docker image prune fails.
+    """
+    # GIVEN a self hosted vellum api URL env var
+    monkeypatch.setenv("VELLUM_API_URL", "mycompany.api.com")
+    monkeypatch.setenv("VELLUM_API_KEY", "123456abcdef")
+
+    # AND a mock Docker client
+    mock_docker_client = MagicMock()
+    mock_docker_from_env.return_value = mock_docker_client
+
+    mock_subprocess_run.side_effect = [
+        subprocess.CalledProcessError(
+            1, ["docker", "image", "prune", "--all", "-f", "--filter", "label=image-type=python-workflow-runtime"]
+        ),
+        subprocess.CompletedProcess(
+            args="", returncode=0, stdout=b'{"manifests": [{"platform": {"architecture": "amd64"}}]}'
+        ),
+        subprocess.CompletedProcess(args="", returncode=0, stdout=b"manifest"),
+        subprocess.CompletedProcess(args="", returncode=0, stdout=b"sha256:hellosha"),
+    ]
+
+    # WHEN the user runs the image push command with any image
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["image", "push", "myrepo.net/myimage:latest"])
+
+    # THEN the command exits successfully despite pruning failure
+    assert result.exit_code == 0, result.output
+
+    # AND the warning message about pruning failure is shown
+    assert "Docker image pruning failed" in result.output
+
+    # AND the success message is still shown
+    assert "Image successfully pushed" in result.output