agentex-sdk 0.1.0a6__py3-none-any.whl

agentex/lib/cli/handlers/cleanup_handlers.py

@@ -0,0 +1,186 @@
+import asyncio
+import os
+
+from rich.console import Console
+
+from agentex import Agentex
+from agentex.lib.utils.logging import make_logger
+
+# Import Temporal client for direct workflow termination
+try:
+    from temporalio.client import Client as TemporalClient  # type: ignore
+except ImportError:
+    TemporalClient = None
+
+logger = make_logger(__name__)
+console = Console()
+
+
+def should_cleanup_on_restart() -> bool:
+    """
+    Check if cleanup should be performed on restart.
+    
+    Returns True if:
+    - ENVIRONMENT=development, OR
+    - AUTO_CLEANUP_ON_RESTART=true
+    """
+    env = os.getenv("ENVIRONMENT", "").lower()
+    auto_cleanup = os.getenv("AUTO_CLEANUP_ON_RESTART", "true").lower()
+    
+    return env == "development" or auto_cleanup == "true"
+
+
+def cleanup_agent_workflows(
+    agent_name: str,
+    force: bool = False,
+    development_only: bool = True
+) -> None:
+    """
+    Clean up all running workflows for an agent during development.
+    
+    This cancels (graceful) all running tasks for the specified agent.
+    When force=True, directly terminates workflows via Temporal client.
+    
+    Args:
+        agent_name: Name of the agent to cleanup workflows for
+        force: If True, directly terminate workflows via Temporal client
+        development_only: Only perform cleanup in development environment
+    """
+    
+    # Safety check - only run in development mode by default
+    if development_only and not force and not should_cleanup_on_restart():
+        logger.warning("Cleanup skipped - not in development mode. Use --force to override.")
+        return
+    
+    method = "terminate (direct)" if force else "cancel (via agent)"
+    console.print(f"[blue]Cleaning up workflows for agent '{agent_name}' using {method}...[/blue]")
+    
+    try:
+        client = Agentex()
+        
+        # Get all running tasks
+        all_tasks = client.tasks.list()
+        running_tasks = [task for task in all_tasks if hasattr(task, 'status') and task.status == "RUNNING"]
+        
+        if not running_tasks:
+            console.print("[yellow]No running tasks found[/yellow]")
+            return
+        
+        console.print(f"[blue]Cleaning up {len(running_tasks)} running task(s) for agent '{agent_name}'...[/blue]")
+        
+        successful_cleanups = 0
+        total_tasks = len(running_tasks)
+        
+        for task in running_tasks:
+            task_cleanup_success = False
+            
+            if force:
+                # Force mode: Do both graceful RPC cancellation AND direct Temporal termination
+                rpc_success = False
+                temporal_success = False
+                
+                try:
+                    # First: Graceful cancellation via agent RPC (handles database/agent cleanup)
+                    cleanup_single_task(client, agent_name, task.id)
+                    logger.debug(f"Completed RPC cancellation for task {task.id}")
+                    rpc_success = True
+                except Exception as e:
+                    logger.warning(f"RPC cancellation failed for task {task.id}: {e}")
+                
+                try:
+                    # Second: Direct Temporal termination (ensures workflow is forcefully stopped)
+                    asyncio.run(cleanup_single_task_direct(task.id))
+                    logger.debug(f"Completed Temporal termination for task {task.id}")
+                    temporal_success = True
+                except Exception as e:
+                    logger.warning(f"Temporal termination failed for task {task.id}: {e}")
+                
+                # Count as success if either operation succeeded
+                task_cleanup_success = rpc_success or temporal_success
+                
+            else:
+                # Normal mode: Only graceful cancellation via agent RPC
+                try:
+                    cleanup_single_task(client, agent_name, task.id)
+                    task_cleanup_success = True
+                except Exception as e:
+                    logger.error(f"Failed to cleanup task {task.id}: {e}")
+                    task_cleanup_success = False
+            
+            if task_cleanup_success:
+                successful_cleanups += 1
+                logger.debug(f"Successfully cleaned up task {task.id}")
+            else:
+                logger.error(f"Failed to cleanup task {task.id}")
+                # Don't increment successful_cleanups for actual failures
+        
+        if successful_cleanups == total_tasks:
+            console.print(f"[green]✓ Successfully cleaned up all {successful_cleanups} task(s) for agent '{agent_name}'[/green]")
+        elif successful_cleanups > 0:
+            console.print(f"[yellow]⚠ Successfully cleaned up {successful_cleanups}/{total_tasks} task(s) for agent '{agent_name}'[/yellow]")
+        else:
+            console.print(f"[red]✗ Failed to cleanup any tasks for agent '{agent_name}'[/red]")
+        
+    except Exception as e:
+        console.print(f"[red]Agent workflow cleanup failed: {str(e)}[/red]")
+        logger.exception("Agent workflow cleanup failed")
+        raise
+
+
+async def cleanup_single_task_direct(task_id: str) -> None:
+    """
+    Directly terminate a workflow using Temporal client.
+    
+    Args:
+        task_id: ID of the task (used as workflow_id)
+    """
+    if TemporalClient is None:
+        raise ImportError("temporalio package not available for direct workflow termination")
+    
+    try:
+        # Connect to Temporal server (assumes default localhost:7233)
+        client = await TemporalClient.connect("localhost:7233")  # type: ignore
+        
+        # Get workflow handle and terminate
+        handle = client.get_workflow_handle(workflow_id=task_id)  # type: ignore
+        await handle.terminate()  # type: ignore
+        
+        logger.debug(f"Successfully terminated workflow {task_id} via Temporal client")
+        
+    except Exception as e:
+        # Check if the workflow was already completed - this is actually a success case
+        if "workflow execution already completed" in str(e).lower():
+            logger.debug(f"Workflow {task_id} was already completed - no termination needed")
+            return  # Don't raise an exception for this case
+        
+        logger.error(f"Failed to terminate workflow {task_id} via Temporal client: {e}")
+        raise
+
+
+def cleanup_single_task(client: Agentex, agent_name: str, task_id: str) -> None:
+    """
+    Clean up a single task/workflow using agent RPC cancel method.
+    
+    Args:
+        client: Agentex client instance  
+        agent_name: Name of the agent that owns the task
+        task_id: ID of the task to cleanup
+    """
+    try:
+        # Use the agent RPC method to cancel the task
+        try:
+            client.agents.rpc_by_name(
+                agent_name=agent_name,
+                method="task/cancel",
+                params={"task_id": task_id}
+            )
+            logger.debug(f"Successfully cancelled task {task_id} via agent '{agent_name}'")
+        except Exception as e:
+            # If RPC cancel fails, try direct task deletion as fallback
+            logger.warning(f"RPC task/cancel failed for task {task_id}, trying direct deletion: {e}")
+            client.tasks.delete(task_id=task_id)
+            logger.debug(f"Successfully deleted task {task_id} directly")
+        
+    except Exception as e:
+        logger.warning(f"Failed to cleanup task {task_id}: {e}")
+        raise 

agentex/lib/cli/handlers/deploy_handlers.py

@@ -0,0 +1,351 @@
+import os
+import subprocess
+import tempfile
+from pathlib import Path
+from typing import Any
+
+import yaml
+from pydantic import BaseModel, Field
+from rich.console import Console
+
+from agentex.lib.cli.utils.exceptions import DeploymentError, HelmError
+from agentex.lib.cli.utils.kubectl_utils import check_and_switch_cluster_context
+from agentex.lib.sdk.config.agent_config import AgentConfig
+from agentex.lib.sdk.config.agent_manifest import AgentManifest
+from agentex.lib.sdk.config.deployment_config import ClusterConfig
+from agentex.lib.utils.logging import make_logger
+
+logger = make_logger(__name__)
+console = Console()
+
+TEMPORAL_WORKER_KEY = "temporal-worker"
+AGENTEX_AGENTS_HELM_CHART_VERSION = "0.1.2-v2-beta"
+
+
+class InputDeployOverrides(BaseModel):
+    repository: str | None = Field(
+        default=None, description="Override the repository for deployment"
+    )
+    image_tag: str | None = Field(
+        default=None, description="Override the image tag for deployment"
+    )
+
+
+def check_helm_installed() -> bool:
+    """Check if helm is installed and available"""
+    try:
+        result = subprocess.run(
+            ["helm", "version", "--short"], capture_output=True, text=True, check=True
+        )
+        logger.info(f"Helm version: {result.stdout.strip()}")
+        return True
+    except (subprocess.CalledProcessError, FileNotFoundError):
+        return False
+
+
+def add_helm_repo() -> None:
+    """Add the agentex helm repository if not already added"""
+    try:
+        # Check if repo already exists
+        result = subprocess.run(
+            ["helm", "repo", "list"], capture_output=True, text=True, check=True
+        )
+
+        if "scale-egp" not in result.stdout:
+            console.print("Adding agentex helm repository...")
+            subprocess.run(
+                [
+                    "helm",
+                    "repo",
+                    "add",
+                    "scale-egp",
+                    "https://scale-egp-helm-charts-us-west-2.s3.amazonaws.com/charts",
+                ],
+                check=True,
+            )
+        else:
+            logger.info("Helm repository already exists. Running update...")
+
+        subprocess.run(["helm", "repo", "update"], check=True)
+        console.print("[green]✓[/green] Helm repository update successfully")
+
+    except subprocess.CalledProcessError as e:
+        raise HelmError(f"Failed to add helm repository: {e}") from e
+
+
+def load_override_config(override_file_path: str | None = None) -> ClusterConfig | None:
+    """Load override configuration from specified file path"""
+    if not override_file_path:
+        return None
+
+    override_path = Path(override_file_path)
+    if not override_path.exists():
+        raise DeploymentError(f"Override file not found: {override_file_path}")
+
+    try:
+        with open(override_path) as f:
+            config_data = yaml.safe_load(f)
+            return ClusterConfig(**config_data) if config_data else None
+    except Exception as e:
+        raise DeploymentError(
+            f"Failed to load override config from {override_file_path}: {e}"
+        ) from e
+
+
+def merge_deployment_configs(
+    manifest: AgentManifest,
+    cluster_config: ClusterConfig | None,
+    deploy_overrides: InputDeployOverrides,
+) -> dict[str, Any]:
+    agent_config: AgentConfig = manifest.agent
+
+    """Merge global deployment config with cluster-specific overrides into helm values"""
+    if not manifest.deployment:
+        raise DeploymentError("No deployment configuration found in manifest")
+
+    repository = deploy_overrides.repository or manifest.deployment.image.repository
+    image_tag = deploy_overrides.image_tag or manifest.deployment.image.tag
+
+    if not repository or not image_tag:
+        raise DeploymentError("Repository and image tag are required")
+
+    # Start with global configuration
+    helm_values = {
+        "global": {
+            "image": {
+                "repository": repository,
+                "tag": image_tag,
+                "pullPolicy": "IfNotPresent",
+            },
+            "agent": {
+                "name": manifest.agent.name,
+                "description": manifest.agent.description,
+                "acp_type": manifest.agent.acp_type,
+            },
+        },
+        "replicaCount": manifest.deployment.global_config.replicaCount,
+        "resources": {
+            "requests": {
+                "cpu": manifest.deployment.global_config.resources.requests.cpu,
+                "memory": manifest.deployment.global_config.resources.requests.memory,
+            },
+            "limits": {
+                "cpu": manifest.deployment.global_config.resources.limits.cpu,
+                "memory": manifest.deployment.global_config.resources.limits.memory,
+            },
+        },
+    }
+
+    # Handle temporal configuration using new helper methods
+    if agent_config.is_temporal_agent():
+        temporal_config = agent_config.get_temporal_workflow_config()
+        if temporal_config:
+            helm_values[TEMPORAL_WORKER_KEY] = {}
+            helm_values["global"]["workflow"] = {
+                "name": temporal_config.name,
+                "taskQueue": temporal_config.queue_name,
+            }
+            helm_values[TEMPORAL_WORKER_KEY]["enabled"] = True
+
+    secret_env_vars = []
+    if agent_config.credentials:
+        for credential in agent_config.credentials:
+            secret_env_vars.append(
+                {
+                    "name": credential.env_var_name,
+                    "secretName": credential.secret_name,
+                    "secretKey": credential.secret_key,
+                }
+            )
+
+        helm_values["secretEnvVars"] = secret_env_vars
+        if TEMPORAL_WORKER_KEY in helm_values:
+            helm_values[TEMPORAL_WORKER_KEY]["secretEnvVars"] = secret_env_vars
+
+    # Set the agent_config env vars first to the helm values and so then it can be overriden by the cluster config
+    if agent_config.env:
+        helm_values["env"] = agent_config.env
+        if TEMPORAL_WORKER_KEY in helm_values:
+            helm_values[TEMPORAL_WORKER_KEY]["env"] = agent_config.env
+
+    if manifest.deployment and manifest.deployment.imagePullSecrets:
+        pull_secrets = [
+            pull_secret.to_dict()
+            for pull_secret in manifest.deployment.imagePullSecrets
+        ]
+        helm_values["global"]["imagePullSecrets"] = pull_secrets
+        # TODO: Remove this once i bump the chart version again
+        helm_values["imagePullSecrets"] = pull_secrets
+
+    # Apply cluster-specific overrides
+    if cluster_config:
+        if cluster_config.image:
+            if cluster_config.image.repository:
+                helm_values["global"]["image"]["repository"] = (
+                    cluster_config.image.repository
+                )
+            if cluster_config.image.tag:
+                helm_values["global"]["image"]["tag"] = cluster_config.image.tag
+
+        if cluster_config.replicaCount is not None:
+            helm_values["replicaCount"] = cluster_config.replicaCount
+
+        if cluster_config.resources:
+            if cluster_config.resources.requests:
+                helm_values["resources"]["requests"].update(
+                    {
+                        "cpu": cluster_config.resources.requests.cpu,
+                        "memory": cluster_config.resources.requests.memory,
+                    }
+                )
+            if cluster_config.resources.limits:
+                helm_values["resources"]["limits"].update(
+                    {
+                        "cpu": cluster_config.resources.limits.cpu,
+                        "memory": cluster_config.resources.limits.memory,
+                    }
+                )
+
+        if cluster_config.env:
+            helm_values["env"] = cluster_config.env
+
+        # Apply additional arbitrary overrides
+        if cluster_config.additional_overrides:
+            _deep_merge(helm_values, cluster_config.additional_overrides)
+
+    return helm_values
+
+
+def _deep_merge(base_dict: dict[str, Any], override_dict: dict[str, Any]) -> None:
+    """Deep merge override_dict into base_dict"""
+    for key, value in override_dict.items():
+        if (
+            key in base_dict
+            and isinstance(base_dict[key], dict)
+            and isinstance(value, dict)
+        ):
+            _deep_merge(base_dict[key], value)
+        else:
+            base_dict[key] = value
+
+
+def create_helm_values_file(helm_values: dict[str, Any]) -> str:
+    """Create a temporary helm values file"""
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
+        yaml.dump(helm_values, f, default_flow_style=False)
+        return f.name
+
+
+def deploy_agent(
+    manifest_path: str,
+    cluster_name: str,
+    namespace: str,
+    deploy_overrides: InputDeployOverrides,
+    override_file_path: str | None = None,
+) -> None:
+    """Deploy an agent using helm"""
+
+    # Validate prerequisites
+    if not check_helm_installed():
+        raise DeploymentError("Helm is not installed. Please install helm first.")
+
+    # Switch to the specified cluster context
+    check_and_switch_cluster_context(cluster_name)
+
+    manifest = AgentManifest.from_yaml(file_path=manifest_path)
+    override_config = load_override_config(override_file_path)
+
+    # Provide feedback about override configuration
+    if override_config:
+        console.print(f"[green]✓[/green] Using override config: {override_file_path}")
+    else:
+        console.print(
+            "[yellow]ℹ[/yellow] No override config specified, using global defaults"
+        )
+
+    # Add helm repository/update
+    add_helm_repo()
+
+    # Merge configurations
+    helm_values = merge_deployment_configs(manifest, override_config, deploy_overrides)
+
+    # Create values file
+    values_file = create_helm_values_file(helm_values)
+
+    try:
+        agent_name = manifest.agent.name
+        release_name = agent_name
+
+        console.print(
+            f"Deploying agent [bold]{agent_name}[/bold] to cluster [bold]{cluster_name}[/bold] in namespace [bold]{namespace}[/bold]"
+        )
+
+        # Check if release exists
+        try:
+            subprocess.run(
+                ["helm", "status", release_name, "-n", namespace],
+                capture_output=True,
+                check=True,
+            )
+
+            # Release exists, do upgrade
+            console.print("Existing deployment found, upgrading...")
+            command = [
+                "helm",
+                "upgrade",
+                release_name,
+                "scale-egp/agentex-agent",
+                "--version",
+                AGENTEX_AGENTS_HELM_CHART_VERSION,
+                "-f",
+                values_file,
+                "-n",
+                namespace,
+                "--atomic",
+                "--timeout",
+                "10m",
+            ]
+            console.print(f"[blue]ℹ[/blue] Running command: {' '.join(command)}")
+            subprocess.run(command, check=True)
+            console.print("[green]✓[/green] Agent upgraded successfully")
+
+        except subprocess.CalledProcessError:
+            # Release doesn't exist, do install
+            console.print("Installing new deployment...")
+            command = [
+                "helm",
+                "install",
+                release_name,
+                "scale-egp/agentex-agent",
+                "--version",
+                AGENTEX_AGENTS_HELM_CHART_VERSION,
+                "-f",
+                values_file,
+                "-n",
+                namespace,
+                "--create-namespace",
+                "--atomic",
+                "--timeout",
+                "10m",
+            ]
+            console.print(f"[blue]ℹ[/blue] Running command: {' '.join(command)}")
+            subprocess.run(command, check=True)
+            console.print("[green]✓[/green] Agent deployed successfully")
+
+        # Show success message with helpful commands
+        console.print("\n[green]🎉 Deployment completed successfully![/green]")
+        console.print(
+            f"[blue]Check deployment status:[/blue] helm status {release_name} -n {namespace}"
+        )
+        console.print(
+            f"[blue]View logs:[/blue] kubectl logs -l app.kubernetes.io/name=agentex-agent -n {namespace}"
+        )
+
+    except subprocess.CalledProcessError as e:
+        raise HelmError(
+            f"Helm deployment failed: {e}\n"
+            f"Note: Due to --atomic flag, any partial deployment has been automatically rolled back."
+        ) from e
+    finally:
+        # Clean up values file
+        os.unlink(values_file)