plato-sdk-v2 2.8.7__py3-none-any.whl → 2.9.0__py3-none-any.whl

plato/cli/chronos.py

@@ -6,11 +6,13 @@ import asyncio
 import json
 import logging
 import os
+import platform as plat
 import subprocess
 import tempfile
 from pathlib import Path
 from typing import Annotated
 
+import httpx
 import typer
 
 from plato.cli.utils import console
@@ -64,8 +66,6 @@ def launch(
     The config file should contain world.package (required) and optionally world.config,
     runtime.artifact_id, and tags.
     """
-    import httpx
-
     # Set defaults
     if not chronos_url:
         chronos_url = "https://chronos.plato.so"
@@ -126,83 +126,9 @@ def launch(
         raise typer.Exit(1)
 
 
-@chronos_app.command()
-def example(
-    world: str = typer.Argument(
-        "structured-execution",
-        help="World to generate example config for",
-    ),
-    output: Path = typer.Option(
-        None,
-        "--output",
-        "-o",
-        help="Output file path (prints to stdout if not specified)",
-    ),
-):
-    """Generate an example job config file.
-
-    Creates a sample JSON configuration for launching Chronos jobs, which can be
-    customized for your use case.
-
-    Arguments:
-        world: World type to generate example for (default: "structured-execution")
-
-    Options:
-        -o, --output: Output file path. If not specified, prints to stdout.
-
-    Available worlds: structured-execution, code-world
-    """
-    examples = {
-        "structured-execution": {
-            "world_package": "plato-world-structured-execution",
-            "world_version": "latest",
-            "world_config": {
-                "sim_name": "my-sim",
-                "github_url": "https://github.com/example/repo",
-                "max_attempts": 3,
-                "use_backtrack": True,
-                "skill_runner": {
-                    "image": "claude-code:2.1.5",
-                    "config": {"model_name": "anthropic/claude-sonnet-4-20250514", "max_turns": 100},
-                },
-                "plato_api_key": "pk_xxx",
-                "anthropic_api_key": "sk-ant-xxx",
-            },
-            "_comment": "Agents and secrets are embedded directly in world_config",
-        },
-        "code-world": {
-            "world_package": "plato-world-code",
-            "world_config": {
-                "task": "Fix the bug in src/main.py",
-                "repo_url": "https://github.com/example/repo",
-                "coder": {
-                    "image": "claude-code:latest",
-                    "config": {"model_name": "anthropic/claude-sonnet-4-20250514"},
-                },
-            },
-            "_comment": "world_version is optional - uses latest if not specified",
-        },
-    }
-
-    if world not in examples:
-        console.print(f"[red]❌ Unknown world: {world}[/red]")
-        console.print(f"Available examples: {list(examples.keys())}")
-        raise typer.Exit(1)
-
-    example_config = examples[world]
-    json_output = json.dumps(example_config, indent=2)
-
-    if output:
-        with open(output, "w") as f:
-            f.write(json_output)
-        console.print(f"[green]✅ Example config written to {output}[/green]")
-    else:
-        console.print(json_output)
-
-
 def _get_world_runner_dockerfile() -> Path:
     """Get the path to the world runner Dockerfile template."""
-    return Path(__file__).parent / "templates" / "world-runner.Dockerfile"
+    return Path(__file__).parent.parent / "v1" / "cli" / "templates" / "world-runner.Dockerfile"
 
 
 def _build_world_runner_image(platform_override: str | None = None) -> str:
@@ -253,8 +179,6 @@ def _get_docker_platform(override: str | None = None) -> str:
     if override:
         return override
 
-    import platform as plat
-
     system = plat.system()
     machine = plat.machine().lower()
 
@@ -338,8 +262,22 @@ def _build_agent_image(
     return True
 
 
-def _extract_agent_images_from_config(config_data: dict) -> list[str]:
-    """Extract local agent image names from config data."""
+def _extract_agent_name_from_image(image: str) -> str:
+    """Extract agent name from image reference (local or full registry path)."""
+    # Remove tag if present (e.g., "explorer:1.0.6" -> "explorer")
+    image_without_tag = image.split(":")[0]
+    # Get the last path component (e.g., "383806609161.dkr.ecr.../agents/plato/explorer" -> "explorer")
+    return image_without_tag.split("/")[-1]
+
+
+def _extract_agent_images_from_config(config_data: dict, extract_all: bool = False) -> list[str]:
+    """Extract agent image names from config data.
+
+    Args:
+        config_data: The world configuration data
+        extract_all: If True, extract agent names from ALL images (including registry paths).
+                    If False, only extract from local images (no registry prefix).
+    """
     images = []
 
     # Check agents section
@@ -347,9 +285,14 @@ def _extract_agent_images_from_config(config_data: dict) -> list[str]:
     for agent_config in agents.values():
         if isinstance(agent_config, dict):
             image = agent_config.get("image", "")
-            # Only include local images (no registry prefix)
-            if image and "/" not in image.split(":")[0]:
-                name = image.split(":")[0]
+            if image:
+                if extract_all:
+                    name = _extract_agent_name_from_image(image)
+                elif "/" not in image.split(":")[0]:
+                    # Only include local images (no registry prefix)
+                    name = image.split(":")[0]
+                else:
+                    continue
                 if name not in images:
                     images.append(name)
 
@@ -358,14 +301,44 @@ def _extract_agent_images_from_config(config_data: dict) -> list[str]:
         agent_config = config_data.get(field, {})
         if isinstance(agent_config, dict):
             image = agent_config.get("image", "")
-            if image and "/" not in image.split(":")[0]:
-                name = image.split(":")[0]
+            if image:
+                if extract_all:
+                    name = _extract_agent_name_from_image(image)
+                elif "/" not in image.split(":")[0]:
+                    name = image.split(":")[0]
+                else:
+                    continue
                 if name not in images:
                     images.append(name)
 
     return images
 
 
+def _update_config_with_local_images(config_data: dict, built_agents: list[str]) -> dict:
+    """Update config to use local image tags for built agents."""
+    config_data = config_data.copy()
+
+    # Update agents section
+    if "agents" in config_data:
+        config_data["agents"] = config_data["agents"].copy()
+        for agent_name, agent_config in config_data["agents"].items():
+            if isinstance(agent_config, dict) and "image" in agent_config:
+                image_agent_name = _extract_agent_name_from_image(agent_config["image"])
+                if image_agent_name in built_agents:
+                    config_data["agents"][agent_name] = agent_config.copy()
+                    config_data["agents"][agent_name]["image"] = f"{image_agent_name}:latest"
+
+    # Update direct fields
+    for field in ["coder", "verifier", "skill_runner"]:
+        if field in config_data and isinstance(config_data[field], dict) and "image" in config_data[field]:
+            image_agent_name = _extract_agent_name_from_image(config_data[field]["image"])
+            if image_agent_name in built_agents:
+                config_data[field] = config_data[field].copy()
+                config_data[field]["image"] = f"{image_agent_name}:latest"
+
+    return config_data
+
+
 async def _create_chronos_session(
     chronos_url: str,
     api_key: str,
@@ -375,8 +348,6 @@ async def _create_chronos_session(
     tags: list[str] | None = None,
 ) -> dict:
     """Create a session in Chronos."""
-    import httpx
-
     url = f"{chronos_url.rstrip('/')}/api/sessions"
 
     async with httpx.AsyncClient(timeout=30.0) as client:
@@ -394,25 +365,6 @@ async def _create_chronos_session(
         return response.json()
 
 
-async def _close_chronos_session(
-    chronos_url: str,
-    api_key: str,
-    session_id: str,
-) -> None:
-    """Close a Chronos session."""
-    import httpx
-
-    url = f"{chronos_url.rstrip('/')}/api/sessions/{session_id}/close"
-
-    try:
-        async with httpx.AsyncClient(timeout=30.0) as client:
-            response = await client.post(url, headers={"x-api-key": api_key})
-            response.raise_for_status()
-            logger.info(f"Closed Chronos session: {session_id}")
-    except Exception as e:
-        logger.warning(f"Failed to close Chronos session: {e}")
-
-
 async def _complete_chronos_session(
     chronos_url: str,
     api_key: str,
@@ -422,7 +374,6 @@ async def _complete_chronos_session(
     error_message: str | None = None,
 ) -> None:
     """Complete a Chronos session with final status."""
-    import httpx
 
     url = f"{chronos_url.rstrip('/')}/api/sessions/{session_id}/complete"
 
@@ -493,15 +444,29 @@ async def _run_dev_impl(
         world_name = world_package_name or "unknown"
 
     # Build local agent images if agents_dir is provided
+    built_agents: list[str] = []
     if agents_dir:
         agents_dir = agents_dir.expanduser().resolve()
-        agent_images = _extract_agent_images_from_config(config_data)
+        # Extract all agent names (including from full registry paths) when agents_dir is provided
+        agent_images = _extract_agent_images_from_config(config_data, extract_all=True)
         if agent_images:
-            console.print(f"[blue]Building agent images: {agent_images}[/blue]")
+            console.print(f"[blue]Checking for local agent builds: {agent_images}[/blue]")
             for agent_name in agent_images:
-                success = _build_agent_image(agent_name, agents_dir, platform_override)
-                if not success:
-                    raise RuntimeError(f"Failed to build agent image: {agent_name}")
+                # Check if agent exists locally
+                agent_path = agents_dir / agent_name
+                if agent_path.exists() and (agent_path / "Dockerfile").exists():
+                    console.print(f"[blue]Building local agent: {agent_name}[/blue]")
+                    success = _build_agent_image(agent_name, agents_dir, platform_override)
+                    if not success:
+                        raise RuntimeError(f"Failed to build agent image: {agent_name}")
+                    built_agents.append(agent_name)
+                else:
+                    console.print(f"[dim]No local agent found for '{agent_name}', using remote image[/dim]")
+
+        # Update config to use local images for built agents
+        if built_agents:
+            config_data = _update_config_with_local_images(config_data, built_agents)
+            console.print(f"[green]✅ Using local builds for: {built_agents}[/green]")
 
     # Import world module to get config class for environment detection
     # We need to dynamically load the world from world_dir