control-tower 6.0.0rc3__py3-none-any.whl

control_tower/__init__.py

@@ -0,0 +1,3 @@
+from importlib.metadata import version
+
+__version__ = version("control-tower")

control_tower/cli/__init__.py

@@ -0,0 +1,3 @@
+from control_tower.utils.env import ensure_dotenv
+
+ensure_dotenv()

control_tower/cli/base.py

@@ -0,0 +1,215 @@
+import functools
+import json
+import re
+import sys
+import termios
+import tty
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any, cast
+
+import click
+
+from control_tower.cli.completions.completion_callbacks import (
+    completion_env_names,
+    completion_main_task_names,
+    completion_side_task_names,
+)
+from ct_setting_sdk.protocol.errors import ExpectedError
+
+
+def getch() -> str:
+    """Read a single character from stdin without requiring Enter.
+
+    Handles arrow keys by returning "left" or "right" strings.
+    """
+    fd = sys.stdin.fileno()
+    old_settings = termios.tcgetattr(fd)
+    try:
+        tty.setraw(fd)
+        ch = sys.stdin.read(1)
+        if ch == "\x1b":
+            ch2 = sys.stdin.read(1)
+            if ch2 == "[":
+                ch3 = sys.stdin.read(1)
+                if ch3 == "D":
+                    return "left"
+                elif ch3 == "C":
+                    return "right"
+        return ch
+    finally:
+        termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)
+
+
+def clear_screen() -> None:
+    """Clear the terminal screen."""
+    click.echo("\033[2J\033[H", nl=False)
+
+
+def word_wrap(text: str, width: int) -> list[str]:
+    """Wrap text to specified width, preserving line breaks."""
+    result = []
+    for paragraph in text.split("\n"):
+        if not paragraph.strip():
+            result.append("")
+            continue
+        words = paragraph.split()
+        current_line = ""
+        for word in words:
+            if len(current_line) + len(word) + 1 > width:
+                if current_line:
+                    result.append(current_line)
+                current_line = word
+            else:
+                current_line += (" " if current_line else "") + word
+        if current_line:
+            result.append(current_line)
+    return result
+
+
+def load_jsonl(path: Path) -> list[dict[str, Any]]:
+    """Load all non-empty lines from a JSONL file as dicts."""
+    results = []
+    with open(path) as f:
+        for line in f:
+            line = line.strip()
+            if line:
+                results.append(json.loads(line))
+    return results
+
+
+def select_audit_run(runs_dir: Path, title: str = "Select audit run") -> Path | None:
+    """Interactive picker for audit run directories. Returns None if no runs or user cancels."""
+    from pick import Option, pick
+
+    if not runs_dir.exists():
+        click.secho(f"No runs found in {runs_dir}/", fg="yellow")
+        return None
+
+    runs = sorted(
+        [d for d in runs_dir.iterdir() if d.is_dir()],
+        key=lambda p: p.stat().st_ctime,
+        reverse=True,
+    )
+
+    if not runs:
+        click.secho(f"No runs found in {runs_dir}/", fg="yellow")
+        return None
+
+    options: list[Option] = []
+    for run in runs:
+        results_file = run / "results.jsonl"
+        count = 0
+        if results_file.exists():
+            count = sum(1 for _ in results_file.open())
+        label = f"{run.name} ({count} results)"
+        options.append(Option(label, run))
+
+    selected, _ = pick(
+        options,
+        title=title,
+        indicator="=>",
+        default_index=0,
+    )
+
+    return selected.value if isinstance(selected, Option) else Path(str(selected))
+
+
+def _build_param_map(click_params: list[click.Parameter]) -> dict[str, str]:
+    """Build mapping from Python param names to CLI representations."""
+    param_map: dict[str, str] = {}
+
+    for param in click_params:
+        if not param.name:
+            continue
+
+        param_name = param.name
+        if not isinstance(param_name, str):
+            continue
+
+        if isinstance(param, click.Option):
+            if param.opts:
+                cli_name = cast(str, max(param.opts, key=len))
+                param_map[param_name] = cli_name
+        elif isinstance(param, click.Argument):
+            param_map[param_name] = param_name.upper()
+
+    return param_map
+
+
+def _transform_error_message(message: str, click_params: list[click.Parameter]) -> str:
+    """Transform parameter names in error message to CLI equivalents."""
+    param_map = _build_param_map(click_params)
+
+    for py_name, cli_name in param_map.items():
+        message = re.sub(rf"`{re.escape(py_name)}`", cli_name, message)
+
+        if "_" in py_name:
+            message = re.sub(rf"\b{re.escape(py_name)}\b", cli_name, message)
+
+    return message
+
+
+def handle_user_errors(f: Callable[..., object]) -> Callable[..., object]:
+    """Decorator that converts ExpectedError exceptions to click.ClickException.
+
+    Apply this decorator to Click commands to ensure that ExpectedError exceptions
+    are displayed cleanly without a stack trace, while unexpected errors still
+    show full tracebacks for debugging.
+
+    Also transforms parameter names in error messages to their CLI equivalents:
+    - `param` (backtick-wrapped) -> --param
+    - param_name (with underscores) -> --param-name
+
+    Adds a --traceback flag to show full stack traces for expected errors.
+    """
+
+    @functools.wraps(f)
+    def wrapper(*args: object, traceback: bool = False, **kwargs: object) -> object:
+        if traceback:
+            return f(*args, **kwargs)
+
+        try:
+            return f(*args, **kwargs)
+        except ExpectedError as e:
+            message = str(e)
+
+            ctx = click.get_current_context(silent=True)
+            if ctx and ctx.command:
+                message = _transform_error_message(message, ctx.command.params)
+
+            raise click.ClickException(message) from None
+
+    return click.option(
+        "--traceback",
+        is_flag=True,
+        hidden=True,
+        help="Show full traceback for expected errors (for debugging)",
+    )(wrapper)
+
+
+def task_override_options[F: Callable[..., object]](func: F) -> F:
+    """Add -e/--env, -t/--main-task and -s/--side-task options with shell completion."""
+    func = click.option(
+        "-s",
+        "--side-task",
+        type=str,
+        shell_complete=completion_side_task_names,
+        help="Side task name",
+    )(func)
+    func = click.option(
+        "-t",
+        "--main-task",
+        type=str,
+        shell_complete=completion_main_task_names,
+        help="Main task name",
+    )(func)
+    func = click.option(
+        "-e",
+        "--env",
+        "environment",
+        type=str,
+        shell_complete=completion_env_names,
+        help="Environment name",
+    )(func)
+    return func

control_tower/cli/cloud.py

@@ -0,0 +1,132 @@
+"""CLI commands for AWS cloud operations."""
+
+import click
+
+from control_tower.cli.mcp_meta import CTGroup
+
+
+@click.group(cls=CTGroup)
+def cloud_cli() -> None:
+    """AWS cloud operations for distributed evaluation."""
+
+
+@cloud_cli.command("status")
+@click.option(
+    "--region",
+    type=str,
+    default="us-east-1",
+    help="AWS region to check",
+)
+@click.option(
+    "--all-users",
+    is_flag=True,
+    help="Show resources from all users (default: only yours)",
+)
+def status_cmd(region: str, all_users: bool) -> None:
+    """Show running fleet resources (EC2 instances, SQS queues).
+
+    By default shows only resources launched by the calling AWS identity
+    (IAM username or SSO session name).
+    Use --all-users to see everything tagged ManagedBy=linuxbench.
+    """
+    from datetime import UTC, datetime
+
+    import boto3
+
+    from control_tower.fleet.aws_identity import caller_identity_name
+    from control_tower.utils.env import EnvVar
+
+    env = EnvVar()
+    if env.aws_profile and env.aws_access_key_id:
+        click.echo(
+            "WARNING: both AWS_PROFILE and AWS_ACCESS_KEY_ID are set — static keys take "
+            "precedence, so everything below reflects the static-key identity.\n"
+        )
+    ec2 = boto3.client("ec2", region_name=region)
+    sqs = boto3.client("sqs", region_name=region)
+
+    username = caller_identity_name(region)
+    click.echo(f"AWS identity: {username}")
+    click.echo(f"Region:   {region}\n")
+
+    filters = [
+        {"Name": "tag:ManagedBy", "Values": ["linuxbench"]},
+        {"Name": "instance-state-name", "Values": ["running", "pending"]},
+    ]
+    if not all_users:
+        filters.append({"Name": "tag:LaunchedBy", "Values": [username]})
+
+    resp = ec2.describe_instances(Filters=filters)
+    instances = [inst for res in resp["Reservations"] for inst in res["Instances"]]
+
+    click.echo(f"EC2 Instances: {len(instances)}")
+    by_fleet: dict[str, list] = {}
+    for inst in instances:
+        tags = {t["Key"]: t["Value"] for t in inst.get("Tags", [])}
+        fleet_id = tags.get("FleetId", "unknown")
+        by_fleet.setdefault(fleet_id, []).append((inst, tags))
+
+    if instances:
+        for fleet_id, fleet_instances in sorted(by_fleet.items()):
+            first_inst, first_tags = fleet_instances[0]
+            launch_time = first_inst["LaunchTime"]
+            age = datetime.now(UTC) - launch_time
+            hours = age.total_seconds() / 3600
+            owner = first_tags.get("LaunchedBy", "unknown")
+            itype = first_inst["InstanceType"]
+            click.echo(
+                f"  Fleet {fleet_id[:8]}:  {len(fleet_instances)} x {itype}  "
+                f"({hours:.1f}h)  owner={owner}"
+            )
+            for inst, _tags in fleet_instances:
+                click.echo(
+                    f"    {inst['InstanceId']}  "
+                    f"{inst.get('PublicIpAddress', 'no-ip')}  "
+                    f"state={inst['State']['Name']}"
+                )
+    click.echo()
+
+    # SQS list_queues doesn't support tag-based filtering; show all lb-fleet-*
+    # queues but annotate which belong to the user's EC2 fleets.
+    # None means the listing failed, which is not the same as finding no queues.
+    queue_urls: list[str] | None = None
+    try:
+        queue_resp = sqs.list_queues(QueueNamePrefix="lb-fleet-")
+    except Exception as exc:
+        click.echo(f"SQS Queues: could not be listed: {exc}", err=True)
+    else:
+        queue_urls = queue_resp.get("QueueUrls", [])
+        own_fleet_ids = set(by_fleet.keys())
+
+        click.echo(f"SQS Queues: {len(queue_urls)}")
+        seen_fleets: set[str] = set()
+        for url in sorted(queue_urls):
+            name = url.rsplit("/", 1)[-1]
+            parts = name.split("-")  # lb-fleet-{fleet_id}-{suffix}
+            if len(parts) >= 4:
+                fleet_id = parts[2]
+                if fleet_id not in seen_fleets:
+                    seen_fleets.add(fleet_id)
+                    fleet_queues = [u for u in queue_urls if f"-{fleet_id}-" in u]
+                    owner_mark = " (yours)" if fleet_id in own_fleet_ids else ""
+                    click.echo(
+                        f"  Fleet {fleet_id[:8]}:  {len(fleet_queues)} queues{owner_mark}"
+                    )
+    click.echo()
+
+    if not instances and queue_urls == []:
+        click.echo("No active fleet resources found.")
+    elif instances:
+        own_ids = [
+            inst["InstanceId"]
+            for inst in instances
+            if any(
+                t["Key"] == "LaunchedBy" and t["Value"] == username
+                for t in inst.get("Tags", [])
+            )
+        ]
+        if own_ids:
+            click.echo("To terminate your instances:")
+            click.echo(
+                f"  aws ec2 terminate-instances --instance-ids {' '.join(own_ids)}"
+            )

control_tower/cli/completions/completion_callbacks.py

@@ -0,0 +1,125 @@
+"""Click shell completion for environment / main-task / side-task values (reads ``~/.cache/ct/completions-data.json``)."""
+
+import json
+from pathlib import Path
+from typing import Any
+
+import click
+from click.shell_completion import CompletionItem
+
+_CACHE_DIR_NAME = "ct"
+_COMPLETION_JSON = "completions-data.json"
+
+_completions_cache: tuple[float, dict[str, Any]] | None = None
+
+
+def _completion_json_path() -> Path:
+    return Path.home() / ".cache" / _CACHE_DIR_NAME / _COMPLETION_JSON
+
+
+def invalidate_completion_payload_cache() -> None:
+    global _completions_cache
+    _completions_cache = None
+
+
+def load_completion_payload() -> dict[str, Any]:
+    global _completions_cache
+    path = _completion_json_path()
+    if not path.is_file():
+        return {}
+    try:
+        mtime = path.stat().st_mtime
+    except OSError:
+        return {}
+    if _completions_cache is not None and _completions_cache[0] == mtime:
+        return _completions_cache[1]
+    raw = path.read_text(encoding="utf-8")
+    lines = raw.splitlines()
+    if lines and lines[0].startswith("#"):
+        raw = "\n".join(lines[1:]).lstrip()
+    try:
+        data = json.loads(raw)
+    except json.JSONDecodeError:
+        data = {}
+    if not isinstance(data, dict):
+        data = {}
+    _completions_cache = (mtime, data)
+    return data
+
+
+def _env_ids_from_context(ctx: click.Context | None) -> list[str]:
+    if ctx is None:
+        return []
+    keys = ("environment", "environments", "env_name", "only_envs", "env")
+    collected: list[str] = []
+    c: click.Context | None = ctx
+    while c is not None:
+        params = c.params
+        for key in keys:
+            if key not in params:
+                continue
+            v = params[key]
+            if v is None:
+                continue
+            if isinstance(v, list | tuple):
+                for x in v:
+                    if x is not None and not isinstance(x, Path):
+                        collected.append(str(x))
+            elif not isinstance(v, Path):
+                collected.append(str(v))
+        c = c.parent
+    seen: set[str] = set()
+    out: list[str] = []
+    for x in collected:
+        if x not in seen:
+            seen.add(x)
+            out.append(x)
+    return out
+
+
+def _prefix_match(candidates: list[str], incomplete: str) -> list[str]:
+    if not incomplete:
+        return candidates
+    return [c for c in candidates if c.startswith(incomplete)]
+
+
+def completion_env_names(
+    ctx: click.Context, param: click.Parameter, incomplete: str
+) -> list[CompletionItem | str]:
+    payload = load_completion_payload()
+    envs = payload.get("envs", [])
+    return [CompletionItem(e) for e in _prefix_match(envs, incomplete)]
+
+
+def completion_main_task_names(
+    ctx: click.Context, param: click.Parameter, incomplete: str
+) -> list[CompletionItem | str]:
+    payload = load_completion_payload()
+    envs = _env_ids_from_context(ctx)
+    if envs:
+        by_env: dict[str, list[str]] = payload.get("main_tasks_by_env", {})
+        cands: list[str] = []
+        for name, tasks in by_env.items():
+            if any(e in (name, name.rpartition("/")[2]) for e in envs):
+                cands.extend(tasks)
+        cands = list(dict.fromkeys(cands))
+    else:
+        cands = list(payload.get("main_tasks", []))
+    return [CompletionItem(t) for t in _prefix_match(cands, incomplete)]
+
+
+def completion_side_task_names(
+    ctx: click.Context, param: click.Parameter, incomplete: str
+) -> list[CompletionItem | str]:
+    payload = load_completion_payload()
+    envs = _env_ids_from_context(ctx)
+    if envs:
+        by_env: dict[str, list[str]] = payload.get("side_tasks_by_env", {})
+        cands: list[str] = []
+        for name, tasks in by_env.items():
+            if any(e in (name, name.rpartition("/")[2]) for e in envs):
+                cands.extend(tasks)
+        cands = list(dict.fromkeys(cands))
+    else:
+        cands = list(payload.get("side_tasks", []))
+    return [CompletionItem(t) for t in _prefix_match(cands, incomplete)]

control_tower/cli/completions/completions_build.py

@@ -0,0 +1,59 @@
+"""Static completion data for ``ct`` (written under ``~/.cache/ct``)."""
+
+import json
+from pathlib import Path
+from typing import Any
+
+from control_tower.cli.completions.completion_callbacks import (
+    invalidate_completion_payload_cache,
+)
+
+COMPLETION_GENERATED_LINE = "# Generated by LinuxArena."
+COMPLETION_DATA_JSON = "completions-data.json"
+
+
+def default_completion_data_dir() -> Path:
+    return Path.home() / ".cache" / "ct"
+
+
+def scan_environments() -> dict[str, Any]:
+    """Return env and task ids for shell completion.
+
+    Ids come from ``task_space()`` rather than from task folders, so an env
+    that builds its tasks in code completes the same as a folder-backed one.
+    """
+    from control_tower.environments.registry import environments, map_by_package
+
+    spaces = map_by_package(environments.all(), lambda _, env: env.task_space())
+    main_tasks_by_env = {env_id: sorted(s.mains()) for env_id, s in spaces.items()}
+    side_tasks_by_env = {env_id: sorted(s.sides()) for env_id, s in spaces.items()}
+
+    return {
+        "envs": sorted(main_tasks_by_env),
+        "main_tasks": sorted({t for ts in main_tasks_by_env.values() for t in ts}),
+        "main_tasks_by_env": {
+            k: main_tasks_by_env[k] for k in sorted(main_tasks_by_env)
+        },
+        "side_tasks": sorted({t for ts in side_tasks_by_env.values() for t in ts}),
+        "side_tasks_by_env": {
+            k: side_tasks_by_env[k] for k in sorted(side_tasks_by_env)
+        },
+    }
+
+
+def build_completion_data(*, data_dir: Path | None = None) -> dict[str, Any]:
+    """Write ``completions-data.json`` with env/task lists for Click shell completion."""
+    payload = scan_environments()
+    data_dir = data_dir or default_completion_data_dir()
+    data_dir.mkdir(parents=True, exist_ok=True)
+
+    out_path = data_dir / COMPLETION_DATA_JSON
+    text = json.dumps(payload, indent=2, sort_keys=True)
+    out_path.write_text(
+        "\n".join([COMPLETION_GENERATED_LINE, text]) + "\n",
+        encoding="utf-8",
+    )
+
+    invalidate_completion_payload_cache()
+
+    return payload