plano/cli/planoai/core.py

import json
import subprocess
import os
import sys
import time

import yaml
from planoai.utils import convert_legacy_listeners, getLogger
from planoai.consts import (
    PLANO_DOCKER_IMAGE,
    PLANO_DOCKER_NAME,
)
import subprocess
from planoai.docker_cli import (
    docker_container_status,
    docker_remove_container,
    docker_start_plano_detached,
    docker_stop_container,
    health_check_endpoint,
    stream_gateway_logs,
)


log = getLogger(__name__)


def _get_gateway_ports(plano_config_file: str) -> list[int]:
    PROMPT_GATEWAY_DEFAULT_PORT = 10000
    LLM_GATEWAY_DEFAULT_PORT = 12000

    # parse plano_config_file yaml file and get prompt_gateway_port
    plano_config_dict = {}
    with open(plano_config_file) as f:
        plano_config_dict = yaml.safe_load(f)

    listeners, _, _ = convert_legacy_listeners(
        plano_config_dict.get("listeners"), plano_config_dict.get("llm_providers")
    )

    all_ports = [listener.get("port") for listener in listeners]

    # unique ports
    all_ports = list(set(all_ports))

    return all_ports


def start_plano(plano_config_file, env, log_timeout=120, foreground=False):
    """
    Start Docker Compose in detached mode and stream logs until services are healthy.

    Args:
        path (str): The path where the prompt_config.yml file is located.
        log_timeout (int): Time in seconds to show logs before checking for healthy state.
    """
    log.info(
        f"Starting plano gateway, image name: {PLANO_DOCKER_NAME}, tag: {PLANO_DOCKER_IMAGE}"
    )

    try:
        plano_container_status = docker_container_status(PLANO_DOCKER_NAME)
        if plano_container_status != "not found":
            log.info("plano found in docker, stopping and removing it")
            docker_stop_container(PLANO_DOCKER_NAME)
            docker_remove_container(PLANO_DOCKER_NAME)

        gateway_ports = _get_gateway_ports(plano_config_file)

        return_code, _, plano_stderr = docker_start_plano_detached(
            plano_config_file,
            env,
            gateway_ports,
        )
        if return_code != 0:
            log.info("Failed to start plano gateway: " + str(return_code))
            log.info("stderr: " + plano_stderr)
            sys.exit(1)

        start_time = time.time()
        while True:
            all_listeners_healthy = True
            for port in gateway_ports:
                health_check_status = health_check_endpoint(
                    f"http://localhost:{port}/healthz"
                )
                if not health_check_status:
                    all_listeners_healthy = False

            plano_status = docker_container_status(PLANO_DOCKER_NAME)
            current_time = time.time()
            elapsed_time = current_time - start_time

            if plano_status == "exited":
                log.info("plano container exited unexpectedly.")
                stream_gateway_logs(follow=False)
                sys.exit(1)

            # Check if timeout is reached
            if elapsed_time > log_timeout:
                log.info(f"stopping log monitoring after {log_timeout} seconds.")
                stream_gateway_logs(follow=False)
                sys.exit(1)

            if all_listeners_healthy:
                log.info("plano is running and is healthy!")
                break
            else:
                health_check_status_str = (
                    "healthy" if health_check_status else "not healthy"
                )
                log.info(
                    f"plano status: {plano_status}, health status: {health_check_status_str}"
                )
                time.sleep(1)

        if foreground:
            stream_gateway_logs(follow=True)

    except KeyboardInterrupt:
        log.info("Keyboard interrupt received, stopping plano gateway service.")
        stop_docker_container()


def stop_docker_container(service=PLANO_DOCKER_NAME):
    """
    Shutdown all Docker Compose services by running `docker-compose down`.

    Args:
        path (str): The path where the docker-compose.yml file is located.
    """
    log.info(f"Shutting down {service} service.")

    try:
        subprocess.run(
            ["docker", "stop", service],
        )
        subprocess.run(
            ["docker", "rm", service],
        )

        log.info(f"Successfully shut down {service} service.")

    except subprocess.CalledProcessError as e:
        log.info(f"Failed to shut down services: {str(e)}")


def start_cli_agent(plano_config_file=None, settings_json="{}", agent_type="claude"):
    """Start a CLI client connected to Plano."""

    with open(plano_config_file, "r") as file:
        plano_config = file.read()
        plano_config_yaml = yaml.safe_load(plano_config)

    # Get egress listener configuration (supports both legacy dict and list formats)
    host = "127.0.0.1"
    port = 12000
    listeners = plano_config_yaml.get("listeners")
    if isinstance(listeners, dict):
        egress_config = listeners.get("egress_traffic", {})
        host = egress_config.get("host", host)
        port = egress_config.get("port", port)
    elif isinstance(listeners, list):
        model_listener = next(
            (
                listener
                for listener in listeners
                if listener.get("type") in ("model", "model_listener")
            ),
            {},
        )
        host = model_listener.get("host", host)
        port = model_listener.get("port", port)

    # Parse additional settings from command line
    try:
        additional_settings = json.loads(settings_json) if settings_json else {}
    except json.JSONDecodeError:
        log.error("Settings must be valid JSON")
        sys.exit(1)

    # Set up environment variables
    env = os.environ.copy()
    env.update(
        {
            "NO_PROXY": host,
            "DISABLE_TELEMETRY": "true",
            "API_TIMEOUT_MS": "600000",
        }
    )

    model_aliases = plano_config_yaml.get("model_aliases", {})
    command_path = None
    command_args = []
    handled_settings = {"NON_INTERACTIVE_MODE"}

    if agent_type == "claude":
        env.update(
            {
                "ANTHROPIC_AUTH_TOKEN": "test",  # Use test token for plano
                "ANTHROPIC_API_KEY": "",
                "ANTHROPIC_BASE_URL": f"http://{host}:{port}",
                "DISABLE_COST_WARNINGS": "true",
            }
        )
        command_path = "claude"

        # Set ANTHROPIC_SMALL_FAST_MODEL from additional_settings or model alias
        if "ANTHROPIC_SMALL_FAST_MODEL" in additional_settings:
            env["ANTHROPIC_SMALL_FAST_MODEL"] = additional_settings[
                "ANTHROPIC_SMALL_FAST_MODEL"
            ]
        elif "arch.claude.code.small.fast" in model_aliases:
            env["ANTHROPIC_SMALL_FAST_MODEL"] = "arch.claude.code.small.fast"
        else:
            log.info(
                "Tip: Set an alias 'arch.claude.code.small.fast' in your model_aliases config to set a small fast model Claude Code"
            )
            log.info("Or provide ANTHROPIC_SMALL_FAST_MODEL in --settings JSON")

        handled_settings.add("ANTHROPIC_SMALL_FAST_MODEL")

    elif agent_type == "codex":
        env.update(
            {
                # Codex uses OpenAI-compatible auth/base URL when routing through Plano.
                "OPENAI_API_KEY": env.get("OPENAI_API_KEY", "test"),
                "OPENAI_BASE_URL": f"http://{host}:{port}/v1",
            }
        )
        command_path = "codex"

        codex_model = additional_settings.get("CODEX_MODEL")
        if codex_model is None and "arch.codex.default" in model_aliases:
            # Codex expects known model metadata. Resolve alias to concrete target by default
            # to avoid metadata fallback warnings for custom alias names.
            codex_model = model_aliases["arch.codex.default"].get(
                "target", "arch.codex.default"
            )
        if codex_model:
            command_args.extend(["-m", codex_model])

        handled_settings.add("CODEX_MODEL")
    elif agent_type == "opencode":
        env.update(
            {
                # OpenCode uses OpenAI-compatible auth/base URL when routing through Plano.
                "OPENAI_API_KEY": env.get("OPENAI_API_KEY", "test"),
                "OPENAI_BASE_URL": f"http://{host}:{port}/v1",
            }
        )
        command_path = "opencode"

        opencode_model = additional_settings.get("OPENCODE_MODEL")
        if opencode_model is None and "arch.opencode.default" in model_aliases:
            opencode_model = model_aliases["arch.opencode.default"].get(
                "target", "arch.opencode.default"
            )

        if opencode_model:
            # Set both generic and client-specific model env vars for compatibility.
            env["OPENAI_MODEL"] = opencode_model
            env["OPENCODE_MODEL"] = opencode_model

        handled_settings.add("OPENCODE_MODEL")
    else:
        raise ValueError(
            f"Unsupported cli-agent type '{agent_type}'. Supported values: claude, codex, opencode"
        )

    # Non-interactive mode configuration from additional_settings only
    if additional_settings.get("NON_INTERACTIVE_MODE", False):
        env.update(
            {
                "CI": "true",
                "FORCE_COLOR": "0",
                "NODE_NO_READLINE": "1",
                "TERM": "dumb",
            }
        )

    # Add passthrough settings for supported agents.
    if additional_settings:
        passthrough_settings = {
            k: v for k, v in additional_settings.items() if k not in handled_settings
        }
        if agent_type == "claude" and passthrough_settings:
            command_args.append(f"--settings={json.dumps(passthrough_settings)}")

    log.info(f"Connecting {agent_type} CLI agent to Plano at {host}:{port}")

    try:
        subprocess.run([command_path] + command_args, env=env, check=True)
    except subprocess.CalledProcessError as e:
        log.error(f"Error starting {agent_type}: {e}")
        sys.exit(1)
    except FileNotFoundError:
        if agent_type == "claude":
            log.error(
                "claude not found. Make sure Claude Code is installed: npm install -g @anthropic-ai/claude-code"
            )
        elif agent_type == "codex":
            log.error(
                "codex not found. Make sure Codex CLI is installed: npm install -g @openai/codex"
            )
        else:
            log.error(
                "opencode not found. Make sure OpenCode CLI is installed and available in PATH"
            )
        sys.exit(1)