agentcore-cli 0.31.1__py3-none-any.whl

agentcore/primitives/capacity_provider.py

@@ -0,0 +1,501 @@
+"""Capacity provider add/remove (port of src/cli/primitives/CapacityProviderPrimitive.ts).
+
+A capacity provider is a customer-managed EC2 compute pool stored in ``agentcore.json``
+``capacityProviders`` and synthesized by the vended CDK project.
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass, field
+from typing import Any
+
+import typer
+from pydantic import ValidationError as PydanticValidationError
+
+from agentcore.cli import interactive
+from agentcore.cli.errors import get_error_message
+from agentcore.cli.run_command import run_cli_command
+from agentcore.lib.errors.types import ConflictError, ResourceNotFoundError, ValidationError, to_error
+from agentcore.lib.result import Err, Ok, serialize_result
+from agentcore.lib.schemas.io.config_io import to_jsonable
+from agentcore.lib.schemas.io.path_resolver import find_config_root
+from agentcore.primitives.base import BasePrimitive, RemovableResource, RemovalPreview, SchemaChange
+
+
+@dataclass
+class AddCapacityProviderOptions:
+    name: str | None = None
+    subnets: str | None = None
+    security_groups: str | None = None
+    instance_types: str | None = None
+    operator_role_arn: str | None = None
+    description: str | None = None
+    os: str | None = None
+    volume_name: list[str] = field(default_factory=list)
+    volume_size: list[str] = field(default_factory=list)
+    volume_encrypted: bool | None = None
+    volume_kms_key: str | None = None
+    instance_profile_arn: str | None = None
+    idle_instance_timeout: str | None = None
+    max_lifetime: str | None = None
+
+
+def _split_list(value: str | None) -> list[str]:
+    return [s.strip() for s in (value or "").split(",") if s.strip()]
+
+
+def _js_number(raw: str) -> float | int | None:
+    from agentcore.primitives._harness_cli import _js_number as js_number
+
+    n = js_number(raw)
+    if n is None or n != n:
+        return float("nan")
+    return int(n) if abs(n) != float("inf") and float(n).is_integer() else n
+
+
+def build_capacity_provider(options: AddCapacityProviderOptions) -> dict[str, Any]:
+    """Build and schema-validate the capacity provider config (rejects bad input at add time)."""
+    from agentcore.primitives.knowledge_base import zod_error_message
+    from agentcore.schema import CapacityProvider
+
+    names, sizes = options.volume_name or [], options.volume_size or []
+    if len(names) != len(sizes):
+        raise ValidationError(
+            f"--volume-name and --volume-size must be provided in matching pairs (got {len(names)} name(s) and "
+            f"{len(sizes)} size(s))."
+        )
+    volumes = []
+    for vol_name, size_raw in zip(names, sizes, strict=True):
+        if not vol_name.strip() or not re.fullmatch(r"[0-9]+", size_raw):
+            raise ValidationError(
+                f'Invalid volume "{vol_name}:{size_raw}". --volume-name must be non-empty and --volume-size a whole '
+                "number of GiB (e.g. --volume-name data --volume-size 20)."
+            )
+        ebs: dict[str, Any] = {"name": vol_name, "sizeGiB": int(size_raw)}
+        if options.volume_encrypted is not None:
+            ebs["encrypted"] = options.volume_encrypted
+        if options.volume_kms_key:
+            ebs["kmsKeyId"] = options.volume_kms_key
+        volumes.append({"ebsConfiguration": ebs})
+
+    lifecycle: dict[str, Any] = {}
+    if options.idle_instance_timeout is not None:
+        lifecycle["idleInstanceTimeout"] = _js_number(options.idle_instance_timeout)
+    if options.max_lifetime is not None:
+        lifecycle["maxLifetime"] = _js_number(options.max_lifetime)
+
+    launch: dict[str, Any] = {
+        "operatingSystem": options.os or "LINUX_X86_64",
+        "instanceRequirements": {"allowedInstanceTypes": _split_list(options.instance_types)},
+    }
+    if options.instance_profile_arn:
+        launch["instanceProfileArn"] = options.instance_profile_arn
+    ec2: dict[str, Any] = {
+        "launchTemplateSource": {"launchParameters": launch},
+        "vpcConfiguration": {
+            "subnets": _split_list(options.subnets),
+            "securityGroups": _split_list(options.security_groups),
+        },
+    }
+    if volumes:
+        ec2["volumes"] = volumes
+    if lifecycle:
+        ec2["lifecycleConfiguration"] = lifecycle
+
+    candidate: dict[str, Any] = {"name": options.name}
+    if options.description:
+        candidate["description"] = options.description
+    if options.operator_role_arn:
+        candidate["operatorRoleArn"] = options.operator_role_arn
+    candidate["computeConfiguration"] = {"ec2Configuration": ec2}
+    try:
+        return to_jsonable(CapacityProvider.model_validate(candidate))
+    except PydanticValidationError as exc:
+        raise Exception(zod_error_message(exc)) from None
+
+
+class CapacityProviderPrimitive(BasePrimitive):
+    kind = "capacity-provider"
+    label = "Capacity Provider"
+
+    def _project(self) -> dict[str, Any]:
+        return to_jsonable(self.read_project_spec())
+
+    def add(self, options: AddCapacityProviderOptions) -> Ok[dict[str, Any]] | Err[Exception]:
+        try:
+            capacity_provider = build_capacity_provider(options)
+            project = self._project()
+            existing = project.get("capacityProviders") or []
+            if any(cp["name"] == capacity_provider["name"] for cp in existing):
+                raise ValueError(f'{self.label} "{capacity_provider["name"]}" already exists.')
+            project["capacityProviders"] = [*existing, capacity_provider]
+            self.write_project_spec(project)
+            return Ok({"capacityProviderName": capacity_provider["name"]})
+        except Exception as exc:  # noqa: BLE001
+            return Err(to_error(exc))
+
+    @staticmethod
+    def _referencing_runtimes(project: dict[str, Any], name: str) -> list[str]:
+        return [
+            r["name"]
+            for r in project.get("runtimes") or []
+            if (r.get("capacityProviderConfiguration") or {}).get("capacityProviderName") == name
+        ]
+
+    def remove(self, name: str) -> Ok[None] | Err[Exception]:
+        try:
+            project = self._project()
+            existing = project.get("capacityProviders") or []
+            if not any(cp["name"] == name for cp in existing):
+                return Err(ResourceNotFoundError(f'Capacity provider "{name}" not found.'))
+            refs = self._referencing_runtimes(project, name)
+            if refs:
+                return Err(
+                    ConflictError(
+                        f'Capacity provider "{name}" is referenced by agent(s): {", ".join(refs)}. Remove those '
+                        "references first."
+                    )
+                )
+            remaining = [cp for cp in existing if cp["name"] != name]
+            updated = {**project}
+            if remaining:
+                updated["capacityProviders"] = remaining
+            else:
+                updated.pop("capacityProviders", None)
+            self.write_project_spec(updated)
+            return Ok(None)
+        except Exception as exc:  # noqa: BLE001
+            return Err(to_error(exc))
+
+    def preview_remove(self, name: str) -> RemovalPreview:
+        project = self._project()
+        existing = project.get("capacityProviders") or []
+        if not any(cp["name"] == name for cp in existing):
+            raise ResourceNotFoundError(f'Capacity provider "{name}" not found.')
+        refs = self._referencing_runtimes(project, name)
+        if refs:
+            raise ConflictError(
+                f'Capacity provider "{name}" is referenced by agent(s): {", ".join(refs)}. Remove those references '
+                "first."
+            )
+        remaining = [cp for cp in existing if cp["name"] != name]
+        after = {**project}
+        if remaining:
+            after["capacityProviders"] = remaining
+        else:
+            after.pop("capacityProviders", None)
+        return RemovalPreview(
+            [f"Removing capacity provider: {name}"], [], [SchemaChange("agentcore/agentcore.json", project, after)]
+        )
+
+    def get_removable(self) -> list[RemovableResource]:
+        try:
+            return [RemovableResource(cp["name"]) for cp in self._project().get("capacityProviders") or []]
+        except Exception:  # noqa: BLE001
+            return []
+
+    def get_all_names(self) -> list[str]:
+        return [r.name for r in self.get_removable()]
+
+    # --- CLI ----------------------------------------------------------------------------
+
+    def register_commands(self, add_app: typer.Typer, remove_app: typer.Typer) -> None:
+        @add_app.command(
+            "capacity-provider",
+            help="Add a capacity provider (customer-managed EC2 compute pool for agent runtimes)",
+        )
+        def add_command(
+            name: str | None = typer.Option(None, "--name", help="Capacity provider name [non-interactive]"),
+            operator_role_arn: str | None = typer.Option(
+                None,
+                "--operator-role-arn",
+                help="IAM role ARN AgentCore assumes to manage the capacity provider. Optional — omit to have one "
+                "created automatically [non-interactive]",
+            ),
+            description: str | None = typer.Option(
+                None, "--description", help="Capacity provider description [non-interactive]"
+            ),
+            subnets: str | None = typer.Option(
+                None, "--subnets", help="Comma-separated subnet IDs (1-16) [non-interactive]"
+            ),
+            security_groups: str | None = typer.Option(
+                None, "--security-groups", help="Comma-separated security group IDs (1-16) [non-interactive]"
+            ),
+            os_: str | None = typer.Option(
+                None,
+                "--os",
+                help="Operating system: LINUX_X86_64 or LINUX_ARM64 (default: LINUX_X86_64) [non-interactive]",
+            ),
+            instance_types: str | None = typer.Option(
+                None, "--instance-types", help="Comma-separated allowed EC2 instance types (1-30) [non-interactive]"
+            ),
+            volume_name: list[str] | None = typer.Option(
+                None,
+                "--volume-name",
+                help="Named EBS volume name (repeatable, max 5, paired with --volume-size) [non-interactive]",
+            ),
+            volume_size: list[str] | None = typer.Option(
+                None,
+                "--volume-size",
+                help="EBS volume size in GiB (repeatable, paired with --volume-name) [non-interactive]",
+            ),
+            volume_encrypted: bool = typer.Option(
+                False, "--volume-encrypted", help="Encrypt EBS volumes [non-interactive]"
+            ),
+            volume_kms_key: str | None = typer.Option(
+                None, "--volume-kms-key", help="KMS key ARN for EBS volume encryption [non-interactive]"
+            ),
+            instance_profile_arn: str | None = typer.Option(
+                None, "--instance-profile-arn", help="IAM instance profile ARN for launched instances [non-interactive]"
+            ),
+            idle_instance_timeout: str | None = typer.Option(
+                None, "--idle-instance-timeout", help="Idle instance timeout in seconds (60-1209600) [non-interactive]"
+            ),
+            max_lifetime: str | None = typer.Option(
+                None, "--max-lifetime", help="Maximum instance lifetime in seconds (60-1209600) [non-interactive]"
+            ),
+            as_json: bool = typer.Option(False, "--json", help="Output as JSON [non-interactive]"),
+        ) -> None:
+            if not find_config_root():
+                typer.echo("No agentcore project found. Run `agentcore create` first.", err=True)
+                raise typer.Exit(1)
+            options = AddCapacityProviderOptions(
+                name=name,
+                operator_role_arn=operator_role_arn,
+                description=description,
+                subnets=subnets,
+                security_groups=security_groups,
+                os=os_,
+                instance_types=instance_types,
+                volume_name=list(volume_name or []),
+                volume_size=list(volume_size or []),
+                volume_encrypted=True if volume_encrypted else None,
+                volume_kms_key=volume_kms_key,
+                instance_profile_arn=instance_profile_arn,
+                idle_instance_timeout=idle_instance_timeout,
+                max_lifetime=max_lifetime,
+            )
+
+            def action() -> None:
+                missing = [
+                    flag
+                    for flag, value in (
+                        ("--name", name),
+                        ("--subnets", subnets),
+                        ("--security-groups", security_groups),
+                        ("--instance-types", instance_types),
+                    )
+                    if not value
+                ]
+                if missing:
+                    raise ValidationError(f"Missing required option(s): {', '.join(missing)}")
+                result = self.add(options)
+                if not result.success:
+                    raise result.error
+                if as_json:
+                    typer.echo(_dumps(serialize_result(result)))
+                else:
+                    typer.echo(f"Added capacity provider '{result.value['capacityProviderName']}'")
+
+            run_cli_command(as_json, action)
+
+        self.register_remove_subcommand(remove_app)
+
+    # --- interactive (port of tui/screens/capacity-provider) ------------------------------
+
+    def add_interactive(self) -> None:
+        from agentcore.cli.commands.add._common import first_issue
+        from agentcore.cli.naming import generate_unique_name
+        from agentcore.schema import CapacityProviderName, is_valid_operator_role_arn
+
+        existing = self.get_all_names()
+        C = interactive.Choice
+        name = interactive.text(
+            "Capacity provider name",
+            default=generate_unique_name("MyCapacityProvider", existing),
+            validate=lambda v: (
+                first_issue(CapacityProviderName, v)
+                or ("Capacity provider name already exists" if v in existing else None)
+            ),
+        )
+        operator_role = interactive.text(
+            "Operator role ARN (optional, press Enter to auto-create)",
+            validate=lambda v: (
+                None if v.strip() == "" or is_valid_operator_role_arn(v) else "Must be a valid IAM role ARN"
+            ),
+        )
+
+        def id_list(pattern: str, plural: str, invalid: str):  # type: ignore[no-untyped-def]
+            def validate(value: str) -> str | None:
+                ids = _split_list(value)
+                if not 1 <= len(ids) <= 16:
+                    return f"Provide 1-16 {plural}"
+                return None if all(re.fullmatch(pattern, i) for i in ids) else invalid
+
+            return validate
+
+        subnets = interactive.text(
+            "Subnet IDs (comma-separated, 1-16)",
+            validate=id_list(r"subnet-[0-9a-zA-Z]{8,17}", "subnet IDs", "Each must be a valid subnet ID (subnet-...)"),
+        )
+        security_groups = interactive.text(
+            "Security group IDs (comma-separated, 1-16)",
+            validate=id_list(
+                r"sg-[0-9a-zA-Z]{8,17}", "security group IDs", "Each must be a valid security group ID (sg-...)"
+            ),
+        )
+        os_choice = interactive.select(
+            "Operating system (CPU architecture for the capacity provider instances)",
+            [
+                C("Linux x86_64", "LINUX_X86_64", "Intel/AMD 64-bit Linux instances"),
+                C("Linux ARM64", "LINUX_ARM64", "Graviton/ARM 64-bit Linux instances"),
+            ],
+        )
+        instance_types = interactive.text(
+            "Allowed EC2 instance types (comma-separated, 1-30)",
+            validate=lambda v: None if 1 <= len(_split_list(v)) <= 30 else "Provide 1-30 instance types",
+        )
+        instance_profile = interactive.text(
+            "Instance profile ARN for launched instances (optional, press Enter to skip)",
+            validate=lambda v: (
+                None
+                if v.strip() == "" or re.fullmatch(r"arn:[^:]+:iam::[0-9]{12}:instance-profile/.+", v.strip())
+                else "Must be a valid IAM instance profile ARN (arn:...:instance-profile/...)"
+            ),
+        ).strip()
+
+        volumes: list[dict[str, Any]] = []
+        while len(volumes) < 5:
+            taken = [v["name"] for v in volumes]
+            vol_name = interactive.text(
+                "Volume name (optional — press Enter to finish adding volumes)",
+                validate=lambda v, taken=taken: (
+                    None if v.strip() == "" or v.strip() not in taken else f'Volume "{v.strip()}" already added'
+                ),
+            ).strip()
+            if not vol_name:
+                break
+
+            def validate_size(value: str) -> str | None:
+                if not re.fullmatch(r"[0-9]+", value.strip()):
+                    return "Enter a whole number of GiB"
+                return None if 1 <= int(value.strip()) <= 65536 else "Size must be between 1 and 65536 GiB"
+
+            size = int(interactive.text(f'Size in GiB for "{vol_name}" (1-65536)', validate=validate_size).strip())
+            volumes.append({"name": vol_name, "sizeGiB": size})
+            if len(volumes) >= 5:
+                break
+            if (
+                interactive.select(
+                    "Add another volume?", [C("Yes", "yes", "Add another volume"), C("No", "no", "Continue")]
+                )
+                != "yes"
+            ):
+                break
+
+        encrypted = False
+        kms_key = ""
+        if volumes:
+            encrypted = (
+                interactive.select(
+                    "Encrypt EBS volumes? (Applies to all volumes on this capacity provider)",
+                    [
+                        C("No encryption", "no", "EBS volumes are not encrypted"),
+                        C("Encrypt EBS volumes", "yes", "Optionally with your own KMS key"),
+                    ],
+                )
+                == "yes"
+            )
+            if encrypted:
+                kms_key = interactive.text(
+                    "KMS key ARN for volume encryption (optional, press Enter for an AWS-managed key)",
+                    validate=lambda v: (
+                        None
+                        if v.strip() == ""
+                        or re.fullmatch(
+                            r"arn:[^:]+:kms:[a-z0-9-]+:[0-9]{12}:key/[a-f0-9]{8}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{12}",
+                            v.strip(),
+                        )
+                        else "Must be a valid KMS key ARN"
+                    ),
+                ).strip()
+
+        def validate_lifetime(value: str) -> str | None:
+            trimmed = value.strip()
+            if trimmed == "":
+                return None
+            if not re.fullmatch(r"[0-9]+", trimmed):
+                return "Enter a whole number of seconds"
+            return None if 60 <= int(trimmed) <= 1209600 else "Must be between 60 and 1209600 seconds"
+
+        idle = interactive.text(
+            "Idle instance timeout in seconds (60-1209600, press Enter to skip)", validate=validate_lifetime
+        ).strip()
+        lifetime = interactive.text(
+            "Maximum instance lifetime in seconds (60-1209600, press Enter to skip)", validate=validate_lifetime
+        ).strip()
+        description = interactive.text("Description (optional, press Enter to skip)").strip()
+
+        for label, value in (
+            ("Name", name),
+            ("Operator Role", operator_role or "(auto-create)"),
+            ("Subnets", subnets),
+            ("Security Groups", security_groups),
+            ("OS", os_choice),
+            ("Instance Types", instance_types),
+            ("Instance Profile", instance_profile),
+            ("Volumes", ", ".join(f"{v['name']} ({v['sizeGiB']} GiB)" for v in volumes)),
+            ("Encryption", ("Yes" + (f" ({kms_key})" if kms_key else "")) if volumes and encrypted else ""),
+            ("Idle Timeout", idle),
+            ("Max Lifetime", lifetime),
+            ("Description", description),
+        ):
+            if value:
+                typer.echo(f"  {label}: {value}")
+        if not interactive.confirm("Create this capacity provider?", default=True):
+            typer.echo("Cancelled.")
+            return
+
+        result = self.add(
+            AddCapacityProviderOptions(
+                name=name,
+                operator_role_arn=operator_role or None,
+                description=description or None,
+                subnets=subnets,
+                security_groups=security_groups,
+                os=os_choice,
+                instance_types=instance_types,
+                instance_profile_arn=instance_profile or None,
+                volume_name=[v["name"] for v in volumes],
+                volume_size=[str(v["sizeGiB"]) for v in volumes],
+                volume_encrypted=encrypted if volumes else None,
+                volume_kms_key=kms_key if encrypted and kms_key else None,
+                idle_instance_timeout=idle or None,
+                max_lifetime=lifetime or None,
+            )
+        )
+        if not result.success:
+            typer.echo("Failed to add capacity provider", err=True)
+            typer.echo(get_error_message(result.error), err=True)
+            raise typer.Exit(1)
+        typer.echo(f"Added capacity provider: {result.value['capacityProviderName']}")
+        typer.echo("Deploy with `agentcore deploy`.")
+        typer.echo(f" OS: {os_choice}")
+        typer.echo(f" Instance types: {instance_types}")
+        if volumes:
+            summary = ", ".join(f"{v['name']} ({v['sizeGiB']} GiB)" for v in volumes)
+            typer.echo(f" Volumes: {summary}")
+        if description:
+            typer.echo(f" Desc: {description}")
+
+
+def _dumps(payload: dict[str, Any]) -> str:
+    import json
+
+    return json.dumps(payload, separators=(",", ":"))
+
+
+__all__ = ["AddCapacityProviderOptions", "CapacityProviderPrimitive", "build_capacity_provider"]
+
+PRIMITIVE = CapacityProviderPrimitive()