Adding all changes for temporal to autokill tasks on restart

Author: jasonyang101

src/agentex/lib/cli/commands/agents.py

@@ -11,6 +11,7 @@
     build_agent,
     run_agent,
 )
+from agentex.lib.cli.handlers.cleanup_handlers import cleanup_agent_workflows
 from agentex.lib.cli.handlers.deploy_handlers import (
     DeploymentError,
     HelmError,
@@ -71,6 +72,35 @@ def delete(
     logger.info(f"Agent deleted: {agent_name}")
 
 
+@agents.command()
+def cleanup_workflows(
+    agent_name: str = typer.Argument(..., help="Name of the agent to cleanup workflows for"),
+    force: bool = typer.Option(False, help="Force cleanup using direct Temporal termination (bypasses development check)"),
+):
+    """
+    Clean up all running workflows for an agent.
+    
+    By default, uses graceful cancellation via agent RPC.
+    With --force, directly terminates workflows via Temporal client.
+    This is a convenience command that does the same thing as 'agentex tasks cleanup'.
+    """
+    try:
+        console.print(f"[blue]Cleaning up workflows for agent '{agent_name}'...[/blue]")
+        
+        cleanup_agent_workflows(
+            agent_name=agent_name,
+            force=force,
+            development_only=True
+        )
+        
+        console.print(f"[green]✓ Workflow cleanup completed for agent '{agent_name}'[/green]")
+        
+    except Exception as e:
+        console.print(f"[red]Cleanup failed: {str(e)}[/red]")
+        logger.exception("Agent workflow cleanup failed")
+        raise typer.Exit(1) from e
+
+
 @agents.command()
 def build(
     manifest: str = typer.Option(..., help="Path to the manifest you want to use"),
@@ -101,23 +131,33 @@ def build(
     """
     typer.echo(f"Building agent image from manifest: {manifest}")
 
+    # Validate required parameters for building
+    if push and not registry:
+        typer.echo("Error: --registry is required when --push is enabled", err=True)
+        raise typer.Exit(1)
+    
+    # Only proceed with build if we have a registry (for now, to match existing behavior)
+    if not registry:
+        typer.echo("No registry provided, skipping image build")
+        return
+
     platform_list = platforms.split(",") if platforms else []
 
     try:
         image_url = build_agent(
             manifest_path=manifest,
-            registry_url=registry,
-            repository_name=repository_name,
+            registry_url=registry,  # Now guaranteed to be non-None
+            repository_name=repository_name or "default-repo",  # Provide default
             platforms=platform_list,
             push=push,
-            secret=secret,
-            tag=tag,
-            build_args=build_arg,
+            secret=secret or "",  # Provide default empty string
+            tag=tag or "latest",  # Provide default
+            build_args=build_arg or [],  # Provide default empty list
         )
         if image_url:
             typer.echo(f"Successfully built image: {image_url}")
         else:
-            typer.echo("No registry provided, image was not built")
+            typer.echo("Image build completed but no URL returned")
     except Exception as e:
         typer.echo(f"Error building agent image: {str(e)}", err=True)
         logger.exception("Error building agent image")
@@ -127,11 +167,35 @@ def build(
 @agents.command()
 def run(
     manifest: str = typer.Option(..., help="Path to the manifest you want to use"),
+    cleanup_on_start: bool = typer.Option(
+        False, 
+        help="Clean up existing workflows for this agent before starting"
+    ),
 ):
     """
     Run an agent locally from the given manifest.
     """
     typer.echo(f"Running agent from manifest: {manifest}")
+    
+    # Optionally cleanup existing workflows before starting
+    if cleanup_on_start:
+        try:
+            # Parse manifest to get agent name
+            manifest_obj = AgentManifest.from_yaml(file_path=manifest)
+            agent_name = manifest_obj.agent.name
+            
+            console.print(f"[yellow]Cleaning up existing workflows for agent '{agent_name}'...[/yellow]")
+            cleanup_agent_workflows(
+                agent_name=agent_name,
+                force=False,
+                development_only=True
+            )
+            console.print("[green]✓ Pre-run cleanup completed[/green]")
+            
+        except Exception as e:
+            console.print(f"[yellow]⚠ Pre-run cleanup failed: {str(e)}[/yellow]")
+            logger.warning(f"Pre-run cleanup failed: {e}")
+    
     try:
         run_agent(manifest_path=manifest)
     except Exception as e:

src/agentex/lib/cli/commands/tasks.py

@@ -1,10 +1,13 @@
 import typer
 from rich import print_json
+from rich.console import Console
 
 from agentex import Agentex
+from agentex.lib.cli.handlers.cleanup_handlers import cleanup_agent_workflows
 from agentex.lib.utils.logging import make_logger
 
 logger = make_logger(__name__)
+console = Console()
 
 tasks = typer.Typer()
 
@@ -33,6 +36,47 @@ def list():
     print_json(data=[task.to_dict() for task in tasks])
 
 
+@tasks.command()
+def list_running(
+    agent_name: str = typer.Option(..., help="Name of the agent to list running tasks for"),
+):
+    """
+    List all currently running tasks for a specific agent.
+    """
+    client = Agentex()
+    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(f"[yellow]No running tasks found for agent '{agent_name}'[/yellow]")
+        return
+        
+    console.print(f"[green]Found {len(running_tasks)} running task(s) for agent '{agent_name}':[/green]")
+    
+    # Convert to dict with proper datetime serialization
+    serializable_tasks = []
+    for task in running_tasks:
+        try:
+            # Use model_dump with mode='json' for proper datetime handling
+            if hasattr(task, 'model_dump'):
+                serializable_tasks.append(task.model_dump(mode='json'))
+            else:
+                # Fallback for non-Pydantic objects
+                serializable_tasks.append({
+                    "id": getattr(task, 'id', 'unknown'), 
+                    "status": getattr(task, 'status', 'unknown')
+                })
+        except Exception as e:
+            logger.warning(f"Failed to serialize task: {e}")
+            # Minimal fallback
+            serializable_tasks.append({
+                "id": getattr(task, 'id', 'unknown'), 
+                "status": getattr(task, 'status', 'unknown')
+            })
+    
+    print_json(data=serializable_tasks)
+
+
 @tasks.command()
 def delete(
     task_id: str = typer.Argument(..., help="ID of the task to delete"),
@@ -44,3 +88,31 @@ def delete(
     client = Agentex()
     client.tasks.delete(task_id=task_id)
     logger.info(f"Task deleted: {task_id}")
+
+
+@tasks.command()
+def cleanup(
+    agent_name: str = typer.Option(..., help="Name of the agent to cleanup tasks for"),
+    force: bool = typer.Option(False, help="Force cleanup using direct Temporal termination (bypasses development check)"),
+):
+    """
+    Clean up all running tasks/workflows for an agent.
+    
+    By default, uses graceful cancellation via agent RPC.
+    With --force, directly terminates workflows via Temporal client.
+    """
+    try:
+        console.print(f"[blue]Starting cleanup for agent '{agent_name}'...[/blue]")
+        
+        cleanup_agent_workflows(
+            agent_name=agent_name,
+            force=force,
+            development_only=True
+        )
+        
+        console.print(f"[green]✓ Cleanup completed for agent '{agent_name}'[/green]")
+        
+    except Exception as e:
+        console.print(f"[red]Cleanup failed: {str(e)}[/red]")
+        logger.exception("Task cleanup failed")
+        raise typer.Exit(1) from e