agentcore-cli 0.31.1__py3-none-any.whl

agentcore/aws/connect_shell.py

@@ -0,0 +1,266 @@
+"""Interactive shell over WebSocket (port of ``connect-shell.ts``).
+
+Uses the ``websockets`` synchronous client. The upgrade request is SigV4-signed (service
+``bedrock-agentcore``) unless a bearer token is supplied, in which case the token travels in the
+``base64UrlBearerAuthorization`` WebSocket subprotocol and signing is skipped.
+"""
+
+from __future__ import annotations
+
+import base64
+import random
+import ssl
+import threading
+import time
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass, field
+from typing import Any
+from urllib.parse import urlencode, urlsplit
+
+from botocore.httpsession import get_cert_path
+
+from agentcore.aws._http import encode_uri_component, sign_headers
+from agentcore.aws.stage_endpoint import data_plane_endpoint
+from agentcore.lib.errors.types import ShellKickedError
+
+SHELL_ID_HEADER = "x-amzn-bedrock-agentcore-shell-id"
+SESSION_ID_HEADER = "X-Amzn-Bedrock-AgentCore-Runtime-Session-Id"
+KICKED_CLOSE_CODE = 4000
+
+PING_INTERVAL_S = 30.0
+PONG_TIMEOUT_S = 60.0
+_OUTER_DELAY_S = 30.0
+_TOTAL_WINDOW_S = 15 * 60.0
+
+# Indirection points for tests.
+_sleep = time.sleep
+_now = time.monotonic
+
+
+@dataclass(slots=True)
+class ShellReconnectOptions:
+    max_retries: int = 5
+    """Maximum number of attempts; 0 = unlimited."""
+    base_delay: float = 1
+    """Initial delay in seconds between retries; doubles each attempt."""
+    max_delay: float = 15
+    on_attempt: Callable[[int, str], None] | None = None
+    """Called at the start of each reconnect attempt."""
+    on_kicked: Callable[[], None] | None = None
+    """Called on close code 4000 (another client took the session)."""
+
+
+@dataclass(slots=True)
+class ShellConnection:
+    ws: Any
+    """The open ``websockets.sync.client.ClientConnection``."""
+    shell_id: str
+    """Server-assigned shell identifier (from the 101 upgrade response header)."""
+    session_id: str | None = None
+
+
+@dataclass(slots=True)
+class _OpenResult:
+    ws: Any
+    response_headers: Mapping[str, str] = field(default_factory=dict)
+
+
+def build_shell_url(region: str, runtime_arn: str, shell_id: str | None = None) -> str:
+    """``wss://`` URL of the shell endpoint (respects ``AGENTCORE_STAGE``)."""
+    host = urlsplit(data_plane_endpoint(region)).netloc
+    params = {"qualifier": "DEFAULT"}
+    if shell_id:
+        params["shellId"] = shell_id
+    return f"wss://{host}/runtimes/{encode_uri_component(runtime_arn)}/ws/shells?{urlencode(params)}"
+
+
+def translate_upgrade_error(err: BaseException) -> BaseException:
+    """Map WebSocket upgrade HTTP failures to actionable messages."""
+    msg = str(err)
+    if "424" in msg or "Failed Dependency" in msg:
+        return RuntimeError("Agent VM is not ready (error 424). Run `agentcore status` to check deployment.")
+    if "429" in msg or "Too Many" in msg:
+        return RuntimeError("Maximum terminal sessions reached (10). Close an existing session and retry.")
+    if "403" in msg or "Forbidden" in msg:
+        return RuntimeError(
+            "Access denied (403). Check IAM permission bedrock-agentcore:InvokeAgentRuntimeCommandWithWebSocketStream."
+        )
+    return err
+
+
+def _ws_connect(url: str, headers: Mapping[str, str], subprotocols: list[str] | None) -> _OpenResult:
+    from websockets.sync.client import connect
+
+    ssl_context = ssl.create_default_context(cafile=get_cert_path(True))
+    ws = connect(
+        url,
+        additional_headers=dict(headers),
+        subprotocols=subprotocols,  # type: ignore[arg-type]
+        ssl=ssl_context,
+        proxy=None,
+        ping_interval=None,  # keepalive is managed explicitly via start_keepalive()
+        ping_timeout=None,
+        max_size=100 * 1024 * 1024,
+    )
+    response = ws.response
+    return _OpenResult(ws=ws, response_headers=response.headers if response is not None else {})
+
+
+def _close_code(err: BaseException) -> int | None:
+    rcvd = getattr(err, "rcvd", None)
+    return getattr(rcvd, "code", None) if rcvd is not None else getattr(err, "code", None)
+
+
+def _open_websocket(
+    *,
+    region: str,
+    runtime_arn: str,
+    session_id: str | None,
+    shell_id: str | None,
+    headers: Mapping[str, str] | None,
+    bearer_token: str | None,
+) -> ShellConnection:
+    url = build_shell_url(region, runtime_arn, shell_id)
+    extra_headers = dict(headers or {})
+    if session_id:
+        extra_headers[SESSION_ID_HEADER] = session_id
+
+    subprotocols: list[str] | None = None
+    if bearer_token:
+        # CUSTOM_JWT: Sec-WebSocket-Protocol: base64UrlBearerAuthorization.<b64url>, base64UrlBearerAuthorization
+        encoded = base64.urlsafe_b64encode(bearer_token.encode("utf-8")).rstrip(b"=").decode("ascii")
+        subprotocols = [f"base64UrlBearerAuthorization.{encoded}", "base64UrlBearerAuthorization"]
+        request_headers = extra_headers
+    else:
+        https_url = "https://" + url[len("wss://") :]
+        signed = sign_headers("GET", https_url, region=region, headers=extra_headers)
+        # The websocket client writes its own Host header (identical to the signed one).
+        request_headers = {k: v for k, v in signed.items() if k.lower() != "host"}
+
+    try:
+        opened = _ws_connect(url, request_headers, subprotocols)
+    except Exception as exc:
+        code = _close_code(exc)
+        if code == KICKED_CLOSE_CODE:
+            raise ShellKickedError() from exc
+        if isinstance(code, int):
+            raise RuntimeError(f"WebSocket closed before open (code {code})") from exc  # noqa: TRY004
+        translated = translate_upgrade_error(exc)
+        if translated is exc:
+            raise
+        raise translated from exc
+
+    header_shell_id = opened.response_headers.get(SHELL_ID_HEADER) if opened.response_headers else None
+    resolved = header_shell_id or shell_id
+    if not resolved:
+        # An empty shell ID silently breaks reconnect, so treat a missing header as a failure.
+        try:
+            opened.ws.close()
+        except Exception:  # noqa: BLE001, S110
+            pass
+        raise RuntimeError(
+            "Shell session could not be established: the server did not return a shell ID. "
+            "Retry, or run `agentcore status` to check the runtime."
+        )
+    return ShellConnection(ws=opened.ws, shell_id=resolved, session_id=session_id)
+
+
+def start_keepalive(
+    ws: Any,
+    on_dead: Callable[[], None],
+    *,
+    ping_interval: float = PING_INTERVAL_S,
+    pong_timeout: float = PONG_TIMEOUT_S,
+) -> Callable[[], None]:
+    """Send RFC 6455 pings every ``ping_interval`` seconds; call ``on_dead`` if a pong is late.
+
+    Returns a ``stop`` function that cancels the keepalive.
+    """
+    stopped = threading.Event()
+    timers: list[threading.Timer] = []
+
+    def check_pong(pong: threading.Event) -> None:
+        if not stopped.is_set() and not pong.is_set():
+            on_dead()
+
+    def loop() -> None:
+        while not stopped.wait(ping_interval):
+            try:
+                pong = ws.ping()
+            except Exception:  # noqa: BLE001, S112 - socket not open; skip this tick
+                continue
+            timer = threading.Timer(pong_timeout, check_pong, args=(pong,))
+            timer.daemon = True
+            timers.append(timer)
+            timer.start()
+
+    thread = threading.Thread(target=loop, name="agentcore-shell-keepalive", daemon=True)
+    thread.start()
+
+    def stop() -> None:
+        stopped.set()
+        for timer in timers:
+            timer.cancel()
+
+    return stop
+
+
+def _jitter(seconds: float) -> float:
+    return seconds * (0.75 + random.random() * 0.5)
+
+
+def connect_shell(
+    *,
+    region: str,
+    runtime_arn: str,
+    session_id: str | None = None,
+    shell_id: str | None = None,
+    headers: Mapping[str, str] | None = None,
+    reconnect: ShellReconnectOptions | None = None,
+    bearer_token: str | None = None,
+) -> ShellConnection:
+    """Open a shell WebSocket, retrying the initial handshake when ``reconnect`` is given.
+
+    Mid-session reconnects are the caller's responsibility. Raises :class:`ShellKickedError` on
+    close code 4000 (never retried).
+    """
+
+    def attempt(current_shell_id: str | None) -> ShellConnection:
+        return _open_websocket(
+            region=region,
+            runtime_arn=runtime_arn,
+            session_id=session_id,
+            shell_id=current_shell_id,
+            headers=headers,
+            bearer_token=bearer_token,
+        )
+
+    if reconnect is None:
+        return attempt(shell_id)
+
+    start = _now()
+    attempts = 0
+    while True:
+        try:
+            return attempt(shell_id)
+        except ShellKickedError:
+            if reconnect.on_kicked:
+                reconnect.on_kicked()
+            raise
+        except Exception as exc:
+            attempts += 1
+            reason = str(exc)
+            if reconnect.on_attempt:
+                reconnect.on_attempt(attempts, reason)
+
+            # No sleep after the final attempt so the caller sees the error immediately.
+            if reconnect.max_retries > 0 and attempts >= reconnect.max_retries:
+                raise
+            if _now() - start > _TOTAL_WINDOW_S:
+                raise
+
+            # 1s -> 2s -> 4s -> 8s -> 15s (capped), +/-25% jitter; extra 30s pause every 5 attempts.
+            inner = min(reconnect.base_delay * 2 ** min(attempts - 1, 4), reconnect.max_delay)
+            _sleep(_jitter(inner))
+            if attempts % 5 == 0:
+                _sleep(_OUTER_DELAY_S)

agentcore/aws/iam.py

@@ -0,0 +1,41 @@
+"""IAM trust-policy validation (port of ``iam.ts``)."""
+
+from __future__ import annotations
+
+import json
+from typing import Any
+from urllib.parse import unquote
+
+from agentcore.lib.errors.types import ValidationError
+
+IamPolicyDocument = dict[str, Any]
+
+
+def _parse_iam_policy_document(policy: object) -> object:
+    if not isinstance(policy, str):
+        return policy
+    try:
+        return json.loads(unquote(policy, errors="strict"))
+    except (ValueError, UnicodeDecodeError):
+        return None
+
+
+def _stable_stringify(value: object) -> str:
+    # fast-json-stable-stringify: recursively sorted keys, compact separators.
+    return json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False)
+
+
+def validate_iam_role_trust_policy(
+    actual_policy: object, expected_policy: IamPolicyDocument, role_name: str, remediation: str
+) -> None:
+    """Raise unless the role's trust policy structurally matches ``expected_policy``.
+
+    IAM may return ``AssumeRolePolicyDocument`` as an object or a URL-encoded JSON string.
+    """
+    parsed = _parse_iam_policy_document(actual_policy)
+    if isinstance(parsed, dict) and _stable_stringify(parsed) == _stable_stringify(expected_policy):
+        return
+    raise ValidationError(
+        f'Refusing to reuse existing IAM role "{role_name}" because its trust policy does not match '
+        f"the policy required by AgentCore. {remediation}"
+    )

agentcore/aws/partition.py

@@ -0,0 +1,58 @@
+"""Partition-aware ARN/endpoint helpers (port of ``partition.ts``).
+
+Upstream resolves partitions with ``@aws-sdk/util-endpoints``; botocore ships the same
+``partitions.json`` data, so resolution follows the same rules: explicit region membership first,
+then the partition's ``regionRegex``, falling back to the ``aws`` partition.
+"""
+
+from __future__ import annotations
+
+import re
+from functools import lru_cache
+from typing import Any
+
+from botocore.loaders import Loader
+
+_CONSOLE_DOMAINS = {
+    "aws-us-gov": "console.amazonaws-us-gov.com",
+    "aws-cn": "console.amazonaws.cn",
+}
+_DEFAULT_CONSOLE_DOMAIN = "console.aws.amazon.com"
+
+
+@lru_cache(maxsize=1)
+def _partitions() -> list[dict[str, Any]]:
+    return Loader().load_data("partitions")["partitions"]
+
+
+@lru_cache(maxsize=256)
+def _partition_for(region: str) -> tuple[str, str]:
+    partitions = _partitions()
+    for p in partitions:
+        if region in p.get("regions", {}):
+            return p["id"], p["outputs"]["dnsSuffix"]
+    for p in partitions:
+        if re.match(p["regionRegex"], region):
+            return p["id"], p["outputs"]["dnsSuffix"]
+    aws = next(p for p in partitions if p["id"] == "aws")
+    return aws["id"], aws["outputs"]["dnsSuffix"]
+
+
+def get_partition(region: str) -> str:
+    return _partition_for(region)[0]
+
+
+def arn_prefix(region: str) -> str:
+    return f"arn:{get_partition(region)}"
+
+
+def dns_suffix(region: str) -> str:
+    return _partition_for(region)[1]
+
+
+def service_endpoint(service: str, region: str) -> str:
+    return f"{service}.{region}.{dns_suffix(region)}"
+
+
+def console_domain(region: str) -> str:
+    return _CONSOLE_DOMAINS.get(get_partition(region), _DEFAULT_CONSOLE_DOMAIN)

agentcore/aws/policy_generation.py

@@ -0,0 +1,116 @@
+"""Cedar policy generation (port of ``policy-generation.ts``)."""
+
+from __future__ import annotations
+
+import time
+from typing import TypedDict
+
+from agentcore.aws.account import create_aws_client
+
+# Mirrors the SDK waiter options upstream passes (maxWaitTime 120s, minDelay 2s, maxDelay 5s).
+_WAIT_MAX_S = 120.0
+_WAIT_MIN_DELAY_S = 2.0
+_WAIT_MAX_DELAY_S = 5.0
+_SUCCESS_STATES = frozenset({"GENERATED"})
+_FAILURE_STATES = frozenset({"GENERATE_FAILED", "DELETE_FAILED"})
+
+# Indirection points for tests.
+_sleep = time.sleep
+_now = time.monotonic
+
+
+class StartPolicyGenerationResult(TypedDict):
+    generationId: str
+
+
+class GetPolicyGenerationResult(TypedDict):
+    status: str
+    statement: str
+
+
+def start_policy_generation(
+    *, policy_engine_id: str, description: str, region: str, resource_arn: str
+) -> StartPolicyGenerationResult:
+    client = create_aws_client("bedrock-agentcore-control", region)
+    response = client.start_policy_generation(
+        policyEngineId=policy_engine_id,
+        resource={"arn": resource_arn},
+        content={"rawText": description},
+        name=f"cli_generation_{int(time.time() * 1000)}",
+    )
+    generation_id = response.get("policyGenerationId")
+    if not generation_id:
+        raise RuntimeError("No generation ID returned from StartPolicyGeneration")
+    return {"generationId": generation_id}
+
+
+def _wait_for_generation(client: object, generation_id: str, policy_engine_id: str) -> str:
+    """Poll ``GetPolicyGeneration`` until a terminal status; returns ``SUCCESS``/``FAILURE``/``TIMEOUT``.
+
+    Like the SDK waiter, request errors are treated as "retry".
+    """
+    start = _now()
+    attempt = 0
+    while True:
+        attempt += 1
+        try:
+            status = client.get_policy_generation(  # type: ignore[attr-defined]
+                policyGenerationId=generation_id, policyEngineId=policy_engine_id
+            ).get("status")
+            if status in _SUCCESS_STATES:
+                return "SUCCESS"
+            if status in _FAILURE_STATES:
+                return "FAILURE"
+        except Exception:  # noqa: BLE001, S110
+            pass
+        delay = min(_WAIT_MIN_DELAY_S * 2 ** (attempt - 1), _WAIT_MAX_DELAY_S)
+        if _now() - start + delay > _WAIT_MAX_S:
+            return "TIMEOUT"
+        _sleep(delay)
+
+
+def _finding_text(finding: dict) -> str | None:
+    description = finding.get("description")
+    return description if description is not None else finding.get("type")
+
+
+def get_policy_generation(*, generation_id: str, policy_engine_id: str, region: str) -> GetPolicyGenerationResult:
+    """Wait for a generation to finish and return its first Cedar statement."""
+    client = create_aws_client("bedrock-agentcore-control", region)
+
+    state = _wait_for_generation(client, generation_id, policy_engine_id)
+    # A FAILURE state falls through so the status check below reports the service's reasons.
+    if state == "TIMEOUT":
+        raise RuntimeError(
+            f"Policy generation did not complete within the timeout period (state: {state}). "
+            f"Generation ID: {generation_id}"
+        )
+
+    status_response = client.get_policy_generation(policyGenerationId=generation_id, policyEngineId=policy_engine_id)
+    if status_response.get("status") == "GENERATE_FAILED":
+        reasons = status_response.get("statusReasons")
+        raise RuntimeError(
+            f"Policy generation failed: {', '.join(reasons) if reasons is not None else 'Unknown reason'}"
+        )
+
+    assets = (
+        client.list_policy_generation_assets(policyGenerationId=generation_id, policyEngineId=policy_engine_id).get(
+            "policyGenerationAssets"
+        )
+        or []
+    )
+    if not assets:
+        raise RuntimeError("Policy generation completed but no assets were returned")
+
+    first = assets[0]
+    definition = first.get("definition") or {}
+    statement = (definition.get("cedar") or {}).get("statement") or (definition.get("policy") or {}).get("statement")
+    if not statement:
+        findings = first.get("findings")
+        details = ", ".join(d for d in (_finding_text(f) for f in findings) if d) if findings is not None else ""
+        raise RuntimeError(
+            "Policy generation completed but no policy statement was found in the assets"
+            + (f": {details}" if details else "")
+        )
+
+    return {"status": status_response.get("status") or "GENERATED", "statement": statement}

agentcore/aws/region.py

@@ -0,0 +1,54 @@
+"""Ambient region detection (port of ``region.ts``)."""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from typing import Literal, cast
+
+import botocore.session
+
+from agentcore.schema.schemas.aws_targets import AGENTCORE_REGIONS, AgentCoreRegion
+
+DEFAULT_REGION: AgentCoreRegion = "us-east-1"
+
+RegionSource = Literal["env", "config", "default"]
+
+
+@dataclass(frozen=True, slots=True)
+class RegionDetectionResult:
+    region: AgentCoreRegion
+    source: RegionSource
+
+
+def _is_agentcore_region(region: str) -> bool:
+    return region in AGENTCORE_REGIONS
+
+
+def _load_profile_config(profile: str) -> dict[str, object]:
+    # Same view as @smithy/shared-ini-file-loader's configFile: ~/.aws/config (or AWS_CONFIG_FILE)
+    # with "profile " prefixes stripped.
+    profiles = botocore.session.Session().full_config.get("profiles", {})
+    return profiles.get(profile, {}) or {}
+
+
+def detect_region() -> RegionDetectionResult:
+    """Priority: ``AWS_REGION`` > ``AWS_DEFAULT_REGION`` > profile config > ``us-east-1``.
+
+    Only AgentCore-supported regions are accepted from each source.
+    """
+    env_region = os.environ.get("AWS_REGION")
+    if env_region is None:
+        env_region = os.environ.get("AWS_DEFAULT_REGION")
+    if env_region and _is_agentcore_region(env_region):
+        return RegionDetectionResult(cast(AgentCoreRegion, env_region), "env")
+
+    try:
+        profile = os.environ.get("AWS_PROFILE") or "default"
+        config_region = _load_profile_config(profile).get("region")
+        if isinstance(config_region, str) and _is_agentcore_region(config_region):
+            return RegionDetectionResult(cast(AgentCoreRegion, config_region), "config")
+    except Exception:  # noqa: BLE001, S110 - config file missing or unparsable
+        pass
+
+    return RegionDetectionResult(DEFAULT_REGION, "default")

agentcore/aws/retry.py

@@ -0,0 +1,27 @@
+"""AWS error-retryability helpers (port of ``retry.ts``).
+
+Mirrors the signals the AWS SDK's retry middleware uses: name-based throttling/transient sets
+plus an HTTP status fallback. No message matching.
+"""
+
+from __future__ import annotations
+
+import re
+
+from agentcore.aws._aws_errors import aws_error_name, aws_http_status
+
+_THROTTLING_NAME = re.compile(
+    r"^(Throttling|TooManyRequests|RequestLimitExceeded|LimitExceeded)(Exception)?$", re.IGNORECASE
+)
+_TRANSIENT_NAME = re.compile(r"^(ServiceUnavailable|InternalServer|InternalFailure)(Exception)?$", re.IGNORECASE)
+
+
+def is_retryable_aws_error(err: object) -> bool:
+    """True if ``err`` is a transient AWS error worth retrying."""
+    name = aws_error_name(err) or ""
+    if _THROTTLING_NAME.match(name) or _TRANSIENT_NAME.match(name):
+        return True
+    status = aws_http_status(err)
+    if status == 429:
+        return True
+    return status is not None and 500 <= status < 600