agentcore-cli 0.31.1__py3-none-any.whl

agentcore/aws/agentcore.py

@@ -0,0 +1,971 @@
+"""AgentCore data-plane runtime operations (port of ``agentcore.ts``).
+
+Covers HTTP/SSE invocation, MCP and A2A JSON-RPC over ``InvokeAgentRuntime``, AG-UI event
+streaming, evaluation, session control, and ``InvokeAgentRuntimeCommand`` (exec). SigV4 calls go
+through boto3's ``bedrock-agentcore`` client; ``CUSTOM_JWT`` runtimes are called with a bearer
+token over plain HTTPS, as upstream does with ``fetch``.
+"""
+
+from __future__ import annotations
+
+import codecs
+import itertools
+import json
+import time
+from collections.abc import Iterable, Iterator, Mapping
+from dataclasses import dataclass, field, replace
+from typing import Any, Literal, NamedTuple, NotRequired, Protocol, TypedDict
+
+from botocore.client import BaseClient
+from botocore.config import Config
+
+from agentcore.aws._http import encode_uri_component, http_request, iter_stream_chunks
+from agentcore.aws._util import drop_none, js_string, js_template
+from agentcore.aws.account import create_aws_client
+from agentcore.aws.agui_parser import parse_agui_sse_stream
+from agentcore.aws.agui_types import AguiEvent, AguiRunInput
+from agentcore.aws.partition import service_endpoint
+from agentcore.aws.stage_endpoint import data_plane_endpoint
+from agentcore.lib.utils.json_format import to_json_string
+from agentcore.lib.utils.json_rpc import parse_json_rpc_response
+
+DEFAULT_RUNTIME_USER_ID = "default-user"
+"""Default user ID sent with invocations. Container agents need it to obtain workload access tokens."""
+
+MCP_PROTOCOL_VERSION = "2025-03-26"
+SESSION_ID_HEADER = "X-Amzn-Bedrock-AgentCore-Runtime-Session-Id"
+USER_ID_HEADER = "X-Amzn-Bedrock-AgentCore-Runtime-User-Id"
+_ACCEPT_JSON_OR_SSE = "application/json, text/event-stream"
+
+# Runtime invocations can stream for many minutes; botocore's 60s default read timeout is too short.
+_INVOKE_READ_TIMEOUT_S = 900
+
+# Indirection point for tests.
+_sleep = time.sleep
+
+
+class SSELogger(Protocol):
+    def log_sse_event(self, raw_line: str) -> None: ...
+
+
+class EvaluationReferenceInput(TypedDict):
+    context: dict[str, Any]
+    """``{"spanContext": {"sessionId": str, "traceId"?: str}}``"""
+    expectedTrajectory: NotRequired[dict[str, list[str]]]
+    assertions: NotRequired[list[dict[str, str]]]
+    expectedResponse: NotRequired[dict[str, str]]
+
+
+def _create_agentcore_client(region: str, headers: Mapping[str, str] | None = None) -> BaseClient:
+    """``bedrock-agentcore`` client on the stage-aware data-plane endpoint, injecting ``headers`` pre-signing."""
+    client = create_aws_client(
+        "bedrock-agentcore",
+        region,
+        endpoint_url=data_plane_endpoint(region),
+        config=Config(read_timeout=_INVOKE_READ_TIMEOUT_S),
+    )
+    if headers:
+        extra = dict(headers)
+
+        def add_custom_headers(request: Any, **_: Any) -> None:
+            for name, value in extra.items():
+                request.headers[name] = value
+
+        client.meta.events.register("before-sign.bedrock-agentcore", add_custom_headers)
+    return client
+
+
+# ---------------------------------------------------------------------------
+# Options / results
+# ---------------------------------------------------------------------------
+
+
+@dataclass(slots=True)
+class InvokeAgentRuntimeOptions:
+    region: str
+    runtime_arn: str
+    payload: str
+    session_id: str | None = None
+    user_id: str | None = None
+    """Runtime user ID header (defaults to ``default-user``); scopes OAuth tokens for Container agents."""
+    logger: SSELogger | None = None
+    headers: dict[str, str] | None = None
+    """Custom headers forwarded to the agent runtime."""
+    bearer_token: str | None = None
+    """CUSTOM_JWT token; when set, raw HTTPS with ``Authorization: Bearer`` replaces SigV4."""
+    baggage: str | None = None
+    """W3C baggage header value (e.g. a config bundle ref)."""
+    endpoint: str | None = None
+    """Runtime endpoint qualifier (defaults to ``DEFAULT``)."""
+    payment_instrument_id: str | None = None
+    payment_session_id: str | None = None
+    payment_user_id: str | None = None
+    """Payments end-user identity, sent in the body as ``user_id`` (distinct from ``user_id``)."""
+
+
+@dataclass(slots=True)
+class InvokeAgentRuntimeResult:
+    content: str
+    session_id: str | None = None
+
+
+@dataclass(slots=True)
+class StreamingInvokeResult:
+    stream: Iterator[str]
+    session_id: str | None
+
+
+class StopRuntimeSessionResult(TypedDict):
+    sessionId: str | None
+    statusCode: int | None
+
+
+class SSELineResult(NamedTuple):
+    content: str | None
+    error: str | None
+
+
+# ---------------------------------------------------------------------------
+# SSE / response parsing
+# ---------------------------------------------------------------------------
+
+
+def parse_sse_line(line: str) -> SSELineResult:
+    """Content or error carried by one SSE ``data:`` line (both ``None`` for other lines)."""
+    if not line.startswith("data: "):
+        return SSELineResult(None, None)
+    raw = line[6:]
+    try:
+        parsed = json.loads(raw)
+    except ValueError:
+        return SSELineResult(raw, None)
+    if isinstance(parsed, str):
+        return SSELineResult(parsed, None)
+    if isinstance(parsed, dict) and "error" in parsed:
+        return SSELineResult(None, js_string(parsed["error"]))
+    # ConverseStream-shaped event: extract the text delta.
+    if isinstance(parsed, dict):
+        event = parsed.get("event")
+        delta = event.get("contentBlockDelta") if isinstance(event, dict) else None
+        delta = delta.get("delta") if isinstance(delta, dict) else None
+        text = delta.get("text") if isinstance(delta, dict) else None
+        if isinstance(text, str):
+            return SSELineResult(text, None)
+    return SSELineResult(None, None)
+
+
+def parse_sse(text: str) -> str:
+    """Concatenate the content of an SSE body; stops at the first error (``"Error: ..."``)."""
+    parts: list[str] = []
+    for line in text.split("\n"):
+        content, error = parse_sse_line(line)
+        if error:
+            return f"Error: {error}"
+        if content:
+            parts.append(content)
+    return "".join(parts)
+
+
+def extract_result(text: str) -> str:
+    """``{"result": ...}`` payload text, a JSON string's value, pretty JSON, or the raw text."""
+    try:
+        parsed = json.loads(text)
+    except ValueError:
+        return text
+    if isinstance(parsed, dict) and "result" in parsed:
+        result = parsed["result"]
+        return result if isinstance(result, str) else to_json_string(result)
+    return parsed if isinstance(parsed, str) else to_json_string(parsed)
+
+
+def build_invoke_payload(options: InvokeAgentRuntimeOptions) -> str:
+    """JSON invoke body; payment fields are included only when set."""
+    body: dict[str, str] = {"prompt": options.payload}
+    # The agent reads payload.user_id to scope the payment wallet/budget and applies its own
+    # "default-user" fallback, so it is never baked into the wire when unset.
+    if options.payment_user_id:
+        body["user_id"] = options.payment_user_id
+    if options.payment_instrument_id:
+        body["payment_instrument_id"] = options.payment_instrument_id
+    if options.payment_session_id:
+        body["payment_session_id"] = options.payment_session_id
+    return to_json_string(body, indent=None)
+
+
+def _stream_sse_text(chunks: Iterable[bytes], logger: SSELogger | None) -> Iterator[str]:
+    """Incrementally parse an SSE byte stream into text content, falling back to plain JSON."""
+    decoder = codecs.getincrementaldecoder("utf-8")(errors="replace")
+    buffer = ""
+    full_response = ""
+    yielded_content = False
+    source = iter(chunks)
+    try:
+        for chunk in source:
+            decoded = decoder.decode(chunk)
+            buffer += decoded
+            full_response += decoded
+            lines = buffer.split("\n")
+            buffer = lines.pop()
+            for line in lines:
+                if logger is not None and line.strip():
+                    logger.log_sse_event(line)
+                content, error = parse_sse_line(line)
+                if error:
+                    yield f"Error: {error}"
+                    return
+                if content:
+                    yield content
+                    yielded_content = True
+
+        if buffer:
+            if logger is not None and buffer.strip():
+                logger.log_sse_event(buffer)
+            content, error = parse_sse_line(buffer)
+            if error:
+                yield f"Error: {error}"
+            elif content:
+                yield content
+                yielded_content = True
+
+        # Plain JSON (non-SSE) response.
+        if not yielded_content and full_response.strip():
+            yield extract_result(full_response.strip())
+    finally:
+        close = getattr(source, "close", None)
+        if close is not None:
+            close()
+
+
+# ---------------------------------------------------------------------------
+# Bearer token (CUSTOM_JWT) thin HTTP client
+# ---------------------------------------------------------------------------
+
+
+def _build_invoke_url(region: str, runtime_arn: str, endpoint: str | None = None) -> str:
+    qualifier = endpoint if endpoint is not None else "DEFAULT"
+    host = service_endpoint("bedrock-agentcore", region)
+    return f"https://{host}/runtimes/{encode_uri_component(runtime_arn)}/invocations?qualifier={qualifier}"
+
+
+def build_bearer_invoke_headers(options: Any, accept: str) -> dict[str, str]:
+    """Headers for bearer-token invocations (``options``: ``bearer_token``/``session_id``/``user_id``/``headers``/``baggage``)."""
+    headers = {
+        "Authorization": f"Bearer {options.bearer_token}",
+        "Content-Type": "application/json",
+        "Accept": accept,
+    }
+    if options.session_id:
+        headers[SESSION_ID_HEADER] = options.session_id
+    headers[USER_ID_HEADER] = options.user_id if options.user_id is not None else DEFAULT_RUNTIME_USER_ID
+    baggage = getattr(options, "baggage", None)
+    if baggage:
+        headers["baggage"] = baggage
+    if options.headers:
+        headers.update(options.headers)
+    return headers
+
+
+def _raise_invoke_failed(response: Any, prefix: str = "Invoke failed") -> None:
+    try:
+        body = response.text()
+    except Exception:  # noqa: BLE001
+        body = ""
+    raise RuntimeError(f"{prefix} ({response.status}): {body or response.reason}")
+
+
+def _invoke_with_bearer_token_streaming(options: InvokeAgentRuntimeOptions) -> StreamingInvokeResult:
+    url = _build_invoke_url(options.region, options.runtime_arn, options.endpoint)
+    headers = build_bearer_invoke_headers(options, _ACCEPT_JSON_OR_SSE)
+    response = http_request("POST", url, headers=headers, body=build_invoke_payload(options), stream=True)
+    if not response.ok:
+        _raise_invoke_failed(response)
+    session_id = response.header(SESSION_ID_HEADER)
+    return StreamingInvokeResult(stream=_stream_sse_text(response.iter_bytes(), options.logger), session_id=session_id)
+
+
+def _invoke_with_bearer_token(options: InvokeAgentRuntimeOptions) -> InvokeAgentRuntimeResult:
+    url = _build_invoke_url(options.region, options.runtime_arn, options.endpoint)
+    headers = build_bearer_invoke_headers(options, "application/json")
+    response = http_request("POST", url, headers=headers, body=build_invoke_payload(options))
+    if not response.ok:
+        _raise_invoke_failed(response)
+    session_id = response.header(SESSION_ID_HEADER)
+    text = response.text()
+    content = parse_sse(text) if "data: " in text else extract_result(text)
+    return InvokeAgentRuntimeResult(content=content, session_id=session_id)
+
+
+# ---------------------------------------------------------------------------
+# SDK-based invoke (SigV4)
+# ---------------------------------------------------------------------------
+
+
+def _invoke_params(options: InvokeAgentRuntimeOptions) -> dict[str, Any]:
+    params: dict[str, Any] = {
+        "agentRuntimeArn": options.runtime_arn,
+        "payload": build_invoke_payload(options).encode("utf-8"),
+        "contentType": "application/json",
+        "accept": _ACCEPT_JSON_OR_SSE,
+        "runtimeUserId": options.user_id if options.user_id is not None else DEFAULT_RUNTIME_USER_ID,
+    }
+    if options.session_id is not None:
+        params["runtimeSessionId"] = options.session_id
+    if options.baggage:
+        params["baggage"] = options.baggage
+    return params
+
+
+def invoke_agent_runtime_streaming(options: InvokeAgentRuntimeOptions) -> StreamingInvokeResult:
+    """Invoke a runtime and stream text chunks as they arrive."""
+    if options.bearer_token:
+        return _invoke_with_bearer_token_streaming(options)
+
+    client = _create_agentcore_client(options.region, options.headers)
+    response = client.invoke_agent_runtime(**_invoke_params(options))
+    body = response.get("response")
+    if body is None:
+        raise RuntimeError("No response from AgentCore Runtime")
+    return StreamingInvokeResult(
+        stream=_stream_sse_text(iter_stream_chunks(body), options.logger),
+        session_id=response.get("runtimeSessionId"),
+    )
+
+
+def invoke_agent_runtime(options: InvokeAgentRuntimeOptions) -> InvokeAgentRuntimeResult:
+    """Invoke a runtime and return the full response text."""
+    if options.bearer_token:
+        return _invoke_with_bearer_token(options)
+
+    client = _create_agentcore_client(options.region, options.headers)
+    response = client.invoke_agent_runtime(**_invoke_params(options))
+    body = response.get("response")
+    if body is None:
+        raise RuntimeError("No response from AgentCore Runtime")
+    text = body.read().decode("utf-8", errors="replace")
+    content = parse_sse(text) if "data: " in text else extract_result(text)
+    return InvokeAgentRuntimeResult(content=content, session_id=response.get("runtimeSessionId"))
+
+
+# ============================================================================
+# Evaluate
+# ============================================================================
+
+
+class EvaluationResultContext(TypedDict):
+    sessionId: str | None
+    traceId: str | None
+    spanId: str | None
+
+
+class EvaluationResultTokenUsage(TypedDict):
+    inputTokens: int
+    outputTokens: int
+    totalTokens: int
+
+
+class EvaluationResult(TypedDict):
+    evaluatorArn: str | None
+    evaluatorId: str | None
+    evaluatorName: str | None
+    explanation: str | None
+    value: float | None
+    label: str | None
+    errorMessage: str | None
+    errorCode: str | None
+    context: EvaluationResultContext | None
+    tokenUsage: EvaluationResultTokenUsage | None
+
+
+class EvaluateResult(TypedDict):
+    evaluationResults: list[EvaluationResult]
+
+
+def evaluate(
+    *,
+    region: str,
+    evaluator_id: str,
+    session_spans: list[Any],
+    target_span_ids: list[str] | None = None,
+    target_trace_ids: list[str] | None = None,
+    evaluation_reference_inputs: list[EvaluationReferenceInput] | None = None,
+) -> EvaluateResult:
+    """Run an on-demand evaluation of agent spans with one evaluator."""
+    client = create_aws_client("bedrock-agentcore", region)
+
+    params: dict[str, Any] = {"evaluatorId": evaluator_id, "evaluationInput": {"sessionSpans": session_spans}}
+    if target_span_ids is not None:
+        params["evaluationTarget"] = {"spanIds": target_span_ids}
+    elif target_trace_ids is not None:
+        params["evaluationTarget"] = {"traceIds": target_trace_ids}
+    if evaluation_reference_inputs is not None:
+        params["evaluationReferenceInputs"] = evaluation_reference_inputs
+
+    response = client.evaluate(**params)
+    results = response.get("evaluationResults")
+    if results is None:
+        raise RuntimeError("No evaluation results returned")
+
+    def map_result(r: dict[str, Any]) -> EvaluationResult:
+        ctx = r.get("context")
+        span_context = ctx.get("spanContext") if isinstance(ctx, dict) else None
+        usage = r.get("tokenUsage")
+        return {
+            "evaluatorArn": r.get("evaluatorArn"),
+            "evaluatorId": r.get("evaluatorId"),
+            "evaluatorName": r.get("evaluatorName"),
+            "explanation": r.get("explanation"),
+            "value": r.get("value"),
+            "label": r.get("label"),
+            "errorMessage": r.get("errorMessage"),
+            "errorCode": r.get("errorCode"),
+            "context": {
+                "sessionId": span_context.get("sessionId"),
+                "traceId": span_context.get("traceId"),
+                "spanId": span_context.get("spanId"),
+            }
+            if span_context
+            else None,
+            "tokenUsage": {
+                "inputTokens": usage.get("inputTokens") or 0,
+                "outputTokens": usage.get("outputTokens") or 0,
+                "totalTokens": usage.get("totalTokens") or 0,
+            }
+            if usage
+            else None,
+        }
+
+    return {"evaluationResults": [map_result(r) for r in results]}
+
+
+# ---------------------------------------------------------------------------
+# MCP: JSON-RPC over InvokeAgentRuntime
+# ---------------------------------------------------------------------------
+
+
+@dataclass(slots=True)
+class McpInvokeOptions:
+    region: str
+    runtime_arn: str
+    user_id: str | None = None
+    mcp_session_id: str | None = None
+    logger: SSELogger | None = None
+    headers: dict[str, str] | None = None
+    bearer_token: str | None = None
+
+
+class McpToolDef(TypedDict):
+    name: str
+    description: NotRequired[str]
+    inputSchema: NotRequired[dict[str, Any]]
+
+
+@dataclass(slots=True)
+class McpListToolsResult:
+    tools: list[McpToolDef]
+    mcp_session_id: str | None = None
+
+
+@dataclass(slots=True)
+class _McpRpcResult:
+    result: dict[str, Any]
+    mcp_session_id: str | None = None
+    error: dict[str, Any] | None = None
+
+
+_mcp_request_ids = itertools.count(1)
+_a2a_request_ids = itertools.count(1)
+
+_TRANSIENT_STATUS_CODES = frozenset({429, 500, 502, 503, 504})
+_MAX_FETCH_RETRIES = 3
+_FETCH_RETRY_DELAY_S = 1.0
+_COLD_START_RETRIES = 3
+_COLD_START_DELAY_S = 2.0
+
+
+def _log(logger: SSELogger | None, line: str) -> None:
+    if logger is not None:
+        logger.log_sse_event(line)
+
+
+def _fetch_with_retry(url: str, headers: Mapping[str, str], body: str, logger: SSELogger | None) -> Any:
+    """POST with retries on transient failures (5xx, 429, network errors)."""
+    for attempt in range(_MAX_FETCH_RETRIES):
+        try:
+            response = http_request("POST", url, headers=headers, body=body)
+            if response.ok or response.status not in _TRANSIENT_STATUS_CODES or attempt == _MAX_FETCH_RETRIES - 1:
+                return response
+            _log(logger, f"Transient failure ({response.status}), retrying ({attempt + 1}/{_MAX_FETCH_RETRIES})...")
+        except Exception:
+            if attempt == _MAX_FETCH_RETRIES - 1:
+                raise
+            _log(logger, f"Network error, retrying ({attempt + 1}/{_MAX_FETCH_RETRIES})...")
+        _sleep(_FETCH_RETRY_DELAY_S)
+    raise RuntimeError("fetchWithRetry: exhausted retries")
+
+
+def _build_mcp_bearer_headers(options: McpInvokeOptions) -> dict[str, str]:
+    headers = {
+        "Authorization": f"Bearer {options.bearer_token}",
+        "Content-Type": "application/json",
+        "Accept": _ACCEPT_JSON_OR_SSE,
+        "Mcp-Protocol-Version": MCP_PROTOCOL_VERSION,
+        USER_ID_HEADER: options.user_id if options.user_id is not None else DEFAULT_RUNTIME_USER_ID,
+    }
+    if options.mcp_session_id:
+        headers["Mcp-Session-Id"] = options.mcp_session_id
+    if options.headers:
+        headers.update(options.headers)
+    return headers
+
+
+def _rpc_result(parsed: dict[str, Any], session_id: str | None) -> _McpRpcResult:
+    result = parsed.get("result")
+    return _McpRpcResult(
+        result=result if result is not None else {},
+        mcp_session_id=session_id,
+        error=parsed.get("error"),
+    )
+
+
+def _mcp_rpc_call_with_bearer(options: McpInvokeOptions, body: dict[str, Any]) -> _McpRpcResult:
+    url = _build_invoke_url(options.region, options.runtime_arn)
+    headers = _build_mcp_bearer_headers(options)
+    serialized = to_json_string(body, indent=None)
+    _log(options.logger, f"MCP request: {serialized}")
+
+    response = _fetch_with_retry(url, headers, serialized, options.logger)
+    if not response.ok:
+        _raise_invoke_failed(response, "MCP call failed")
+
+    text = response.text()
+    _log(options.logger, f"MCP response: {text}")
+    return _rpc_result(parse_json_rpc_response(text), response.header("Mcp-Session-Id"))
+
+
+def _mcp_rpc_notify_with_bearer(options: McpInvokeOptions, body: dict[str, Any]) -> None:
+    url = _build_invoke_url(options.region, options.runtime_arn)
+    headers = _build_mcp_bearer_headers(options)
+    response = _fetch_with_retry(url, headers, to_json_string(body, indent=None), options.logger)
+    if not response.ok:
+        _raise_invoke_failed(response, "MCP notification failed")
+
+
+def _mcp_invoke_params(options: McpInvokeOptions, body: dict[str, Any]) -> dict[str, Any]:
+    params: dict[str, Any] = {
+        "agentRuntimeArn": options.runtime_arn,
+        "payload": to_json_string(body, indent=None).encode("utf-8"),
+        "contentType": "application/json",
+        "accept": _ACCEPT_JSON_OR_SSE,
+        "mcpProtocolVersion": MCP_PROTOCOL_VERSION,
+        "runtimeUserId": options.user_id if options.user_id is not None else DEFAULT_RUNTIME_USER_ID,
+    }
+    if options.mcp_session_id is not None:
+        params["mcpSessionId"] = options.mcp_session_id
+    return params
+
+
+def _mcp_rpc_call(options: McpInvokeOptions, body: dict[str, Any]) -> _McpRpcResult:
+    """Send a JSON-RPC request through InvokeAgentRuntime and parse the response."""
+    if options.bearer_token:
+        return _mcp_rpc_call_with_bearer(options, body)
+
+    client = _create_agentcore_client(options.region, options.headers)
+    _log(options.logger, f"MCP request: {to_json_string(body, indent=None)}")
+    response = client.invoke_agent_runtime(**_mcp_invoke_params(options, body))
+    stream = response.get("response")
+    if stream is None:
+        raise RuntimeError("No response from AgentCore Runtime")
+    text = stream.read().decode("utf-8", errors="replace")
+    _log(options.logger, f"MCP response: {text}")
+    return _rpc_result(parse_json_rpc_response(text), response.get("mcpSessionId"))
+
+
+def _mcp_rpc_call_strict(options: McpInvokeOptions, body: dict[str, Any]) -> _McpRpcResult:
+    """:func:`_mcp_rpc_call`, raising on JSON-RPC errors."""
+    result = _mcp_rpc_call(options, body)
+    if result.error:
+        message = result.error.get("message")
+        raise RuntimeError(message if message is not None else f"MCP error (code {result.error.get('code')})")
+    return result
+
+
+def _mcp_rpc_notify(options: McpInvokeOptions, body: dict[str, Any]) -> None:
+    """Send a JSON-RPC notification (no id, response ignored)."""
+    if options.bearer_token:
+        _mcp_rpc_notify_with_bearer(options, body)
+        return
+    client = _create_agentcore_client(options.region, options.headers)
+    response = client.invoke_agent_runtime(**_mcp_invoke_params(options, body))
+    stream = response.get("response")
+    if stream is not None:
+        stream.close()
+
+
+def _initialize_request() -> dict[str, Any]:
+    return {
+        "jsonrpc": "2.0",
+        "id": next(_mcp_request_ids),
+        "method": "initialize",
+        "params": {
+            "protocolVersion": MCP_PROTOCOL_VERSION,
+            "capabilities": {},
+            "clientInfo": {"name": "agentcore-cli", "version": "1.0.0"},
+        },
+    }
+
+
+def _is_cold_start(err: BaseException) -> bool:
+    msg = str(err)
+    return "initialization time exceeded" in msg or "initialization" in msg
+
+
+def mcp_list_tools(options: McpInvokeOptions) -> McpListToolsResult:
+    """Initialize an MCP session and list tools; retries cold-start initialization timeouts."""
+    for attempt in range(_COLD_START_RETRIES):
+        try:
+            return _mcp_list_tools_once(options)
+        except Exception as exc:
+            if _is_cold_start(exc) and attempt < _COLD_START_RETRIES - 1:
+                _log(options.logger, f"MCP cold start (attempt {attempt + 1}/{_COLD_START_RETRIES}), retrying...")
+                _sleep(_COLD_START_DELAY_S)
+                continue
+            raise
+    raise RuntimeError("Failed to list MCP tools after retries")
+
+
+def _mcp_list_tools_once(options: McpInvokeOptions) -> McpListToolsResult:
+    # Stateless servers may reject initialize but still return a session ID, so errors are tolerated.
+    init_result = _mcp_rpc_call(options, _initialize_request())
+    if init_result.error:
+        _log(
+            options.logger,
+            "MCP initialize returned error (expected for stateless servers): "
+            f"{js_template(init_result.error.get('message'))}",
+        )
+
+    session_id = init_result.mcp_session_id
+    with_session = replace(options, mcp_session_id=session_id)
+    _mcp_rpc_notify(with_session, {"jsonrpc": "2.0", "method": "notifications/initialized"})
+
+    list_result = _mcp_rpc_call_strict(
+        with_session, {"jsonrpc": "2.0", "id": next(_mcp_request_ids), "method": "tools/list", "params": {}}
+    )
+    tools = list_result.result.get("tools") or []
+    mapped: list[McpToolDef] = [
+        drop_none({"name": t.get("name"), "description": t.get("description"), "inputSchema": t.get("inputSchema")})  # type: ignore[misc]
+        for t in tools
+    ]
+    return McpListToolsResult(tools=mapped, mcp_session_id=session_id)
+
+
+def mcp_init_session(options: McpInvokeOptions) -> str | None:
+    """Initialize an MCP session (without listing tools); returns its session ID."""
+    init_result = _mcp_rpc_call(options, _initialize_request())
+    session_id = init_result.mcp_session_id
+    _mcp_rpc_notify(
+        replace(options, mcp_session_id=session_id), {"jsonrpc": "2.0", "method": "notifications/initialized"}
+    )
+    return session_id
+
+
+def mcp_call_tool(options: McpInvokeOptions, tool_name: str, args: dict[str, Any]) -> str:
+    """Call an MCP tool; returns its concatenated text content (or pretty JSON of the result)."""
+    for attempt in range(_COLD_START_RETRIES):
+        try:
+            rpc = _mcp_rpc_call_strict(
+                options,
+                {
+                    "jsonrpc": "2.0",
+                    "id": next(_mcp_request_ids),
+                    "method": "tools/call",
+                    "params": {"name": tool_name, "arguments": args},
+                },
+            )
+            content = rpc.result.get("content")
+            if content:
+                texts = [item["text"] for item in content if isinstance(item, dict) and item.get("text") is not None]
+                if texts:
+                    return "".join(texts)
+            return to_json_string(rpc.result)
+        except Exception as exc:
+            if _is_cold_start(exc) and attempt < _COLD_START_RETRIES - 1:
+                _log(options.logger, f"MCP cold start (attempt {attempt + 1}/{_COLD_START_RETRIES}), retrying...")
+                _sleep(_COLD_START_DELAY_S)
+                continue
+            raise
+    raise RuntimeError("Failed to call MCP tool after retries")
+
+
+# ---------------------------------------------------------------------------
+# A2A: JSON-RPC message/send over InvokeAgentRuntime
+# ---------------------------------------------------------------------------
+
+
+@dataclass(slots=True)
+class A2AInvokeOptions:
+    region: str
+    runtime_arn: str
+    user_id: str | None = None
+    session_id: str | None = None
+    logger: SSELogger | None = None
+    headers: dict[str, str] | None = None
+    bearer_token: str | None = None
+
+
+def _single_value_stream(value: str) -> Iterator[str]:
+    yield value
+
+
+def invoke_a2a_runtime(options: A2AInvokeOptions, message: str) -> StreamingInvokeResult:
+    """Send an A2A ``message/send`` request and return the agent's text as a one-item stream."""
+    body = {
+        "jsonrpc": "2.0",
+        "id": next(_a2a_request_ids),
+        "method": "message/send",
+        "params": {
+            "message": {
+                "role": "user",
+                "parts": [{"kind": "text", "text": message}],
+                "messageId": f"msg-{int(time.time() * 1000)}",
+            }
+        },
+    }
+    serialized = to_json_string(body, indent=None)
+    _log(options.logger, f"A2A request: {serialized}")
+
+    if options.bearer_token:
+        url = _build_invoke_url(options.region, options.runtime_arn)
+        headers = build_bearer_invoke_headers(options, _ACCEPT_JSON_OR_SSE)
+        response = http_request("POST", url, headers=headers, body=serialized)
+        if not response.ok:
+            _raise_invoke_failed(response)
+        text = response.text()
+        _log(options.logger, f"A2A response: {text}")
+        return StreamingInvokeResult(
+            stream=_single_value_stream(parse_a2a_response(text)), session_id=response.header(SESSION_ID_HEADER)
+        )
+
+    client = _create_agentcore_client(options.region, options.headers)
+    params: dict[str, Any] = {
+        "agentRuntimeArn": options.runtime_arn,
+        "payload": serialized.encode("utf-8"),
+        "contentType": "application/json",
+        "accept": _ACCEPT_JSON_OR_SSE,
+        "runtimeUserId": options.user_id if options.user_id is not None else DEFAULT_RUNTIME_USER_ID,
+    }
+    if options.session_id:
+        params["runtimeSessionId"] = options.session_id
+    response = client.invoke_agent_runtime(**params)
+    stream = response.get("response")
+    if stream is None:
+        raise RuntimeError("No response from AgentCore Runtime")
+    text = stream.read().decode("utf-8", errors="replace")
+    _log(options.logger, f"A2A response: {text}")
+    # Upstream does not surface the session ID on the SigV4 A2A path.
+    return StreamingInvokeResult(stream=_single_value_stream(parse_a2a_response(text)), session_id=None)
+
+
+def _text_parts(parts: Any) -> list[str]:
+    if not isinstance(parts, list):
+        return []
+    return [
+        js_string(p["text"])
+        for p in parts
+        if isinstance(p, dict) and (p.get("kind") == "text" or p.get("type") == "text") and p.get("text") is not None
+    ]
+
+
+def parse_a2a_response(text: str) -> str:
+    """Text of an A2A JSON-RPC response (``kind: 'text'`` and ``type: 'text'`` parts)."""
+    try:
+        parsed = json.loads(text)
+        if not isinstance(parsed, dict):
+            return text
+
+        error = parsed.get("error")
+        if isinstance(error, (dict, list)):
+            message = error.get("message") if isinstance(error, dict) else None
+            return f"Error: {js_string(message) if message is not None else to_json_string(error, indent=None)}"
+
+        result = parsed.get("result")
+        if not result:
+            return text
+
+        artifacts = result.get("artifacts") if isinstance(result, dict) else None
+        if artifacts:
+            texts = [
+                t for artifact in artifacts if isinstance(artifact, dict) for t in _text_parts(artifact.get("parts"))
+            ]
+            if texts:
+                return "".join(texts)
+
+        # Fallback: last agent message in the history.
+        history = result.get("history") if isinstance(result, dict) else None
+        if history:
+            for msg in reversed(history):
+                if isinstance(msg, dict) and msg.get("role") == "agent" and msg.get("parts"):
+                    agent_texts = _text_parts(msg.get("parts"))
+                    if agent_texts:
+                        return "".join(agent_texts)
+
+        return to_json_string(result)
+    except Exception:  # noqa: BLE001 - any malformed shape falls back to the raw text
+        return text
+
+
+# ---------------------------------------------------------------------------
+# AGUI: structured event streaming over InvokeAgentRuntime
+# ---------------------------------------------------------------------------
+
+
+@dataclass(slots=True)
+class AguiInvokeOptions:
+    region: str
+    runtime_arn: str
+    session_id: str | None = None
+    user_id: str | None = None
+    logger: SSELogger | None = None
+    headers: dict[str, str] | None = None
+    bearer_token: str | None = None
+    """Not yet supported for AG-UI; raises if provided."""
+
+
+@dataclass(slots=True)
+class AguiStreamingInvokeResult:
+    stream: Iterator[AguiEvent]
+    """All AG-UI events, for rich rendering."""
+    text_stream: Iterator[str]
+    """Only TEXT_MESSAGE_CONTENT deltas (and RUN_ERROR messages)."""
+    session_id: str | None
+
+
+def invoke_agui_runtime(options: AguiInvokeOptions, input: AguiRunInput) -> AguiStreamingInvokeResult:
+    """Invoke an AG-UI runtime and stream its structured events."""
+    if options.bearer_token:
+        raise RuntimeError("Bearer token auth is not yet supported for AGUI. Use SigV4 credentials.")
+
+    client = _create_agentcore_client(options.region, options.headers)
+    params: dict[str, Any] = {
+        "agentRuntimeArn": options.runtime_arn,
+        "payload": to_json_string(input, indent=None).encode("utf-8"),
+        "contentType": "application/json",
+        "accept": "text/event-stream",
+        "runtimeUserId": options.user_id if options.user_id is not None else DEFAULT_RUNTIME_USER_ID,
+    }
+    if options.session_id is not None:
+        params["runtimeSessionId"] = options.session_id
+    response = client.invoke_agent_runtime(**params)
+    body = response.get("response")
+    if body is None:
+        raise RuntimeError("No response from AgentCore Runtime")
+
+    streams = parse_agui_sse_stream(iter_stream_chunks(body), logger=options.logger)
+    if streams.text_stream is None:
+        raise RuntimeError("AGUI parser created in single-consumer mode — textStream unavailable")
+    return AguiStreamingInvokeResult(
+        stream=streams.event_stream, text_stream=streams.text_stream, session_id=response.get("runtimeSessionId")
+    )
+
+
+def stop_runtime_session(*, region: str, runtime_arn: str, session_id: str) -> StopRuntimeSessionResult:
+    """Stop a runtime session."""
+    client = create_aws_client("bedrock-agentcore", region)
+    response = client.stop_runtime_session(agentRuntimeArn=runtime_arn, runtimeSessionId=session_id)
+    return {"sessionId": response.get("runtimeSessionId"), "statusCode": response.get("statusCode")}
+
+
+# ---------------------------------------------------------------------------
+# Delete capacity provider session (data plane)
+# ---------------------------------------------------------------------------
+
+
+class DeleteCapacityProviderSessionResult(TypedDict, total=False):
+    capacityProviderArn: str
+    sessionId: str
+    status: str
+    """Session status after the asynchronous delete, typically ``Deprovisioning``."""
+
+
+def delete_capacity_provider_session(
+    *, region: str, capacity_provider_id: str, session_id: str
+) -> DeleteCapacityProviderSessionResult:
+    """Deprovision one capacity provider session (idempotent, asynchronous; deletes attached EBS volumes).
+
+    ``capacity_provider_id`` is the ``{cpName}-{suffix}`` id, not the ARN.
+    """
+    client = _create_agentcore_client(region)
+    response = client.delete_capacity_provider_session(capacityProviderId=capacity_provider_id, sessionId=session_id)
+    return drop_none(  # type: ignore[return-value]
+        {
+            "capacityProviderArn": response.get("capacityProviderArn"),
+            "sessionId": response.get("sessionId"),
+            "status": response.get("status"),
+        }
+    )
+
+
+# ---------------------------------------------------------------------------
+# Execute bash: run shell commands in runtime containers
+# ---------------------------------------------------------------------------
+
+
+@dataclass(slots=True)
+class ExecuteBashOptions:
+    region: str
+    runtime_arn: str
+    command: str
+    session_id: str | None = None
+    timeout: int | None = None
+    headers: dict[str, str] | None = None
+    bearer_token: str | None = None
+    """Not yet supported for exec; raises if provided."""
+
+
+class ExecuteBashStreamEvent(TypedDict):
+    type: Literal["start", "stdout", "stderr", "stop"]
+    data: NotRequired[str]
+    exitCode: NotRequired[int | None]
+    status: NotRequired[str | None]
+
+
+@dataclass(slots=True)
+class ExecuteBashResult:
+    stream: Iterator[ExecuteBashStreamEvent]
+    session_id: str | None = field(default=None)
+
+
+def execute_bash_command(options: ExecuteBashOptions) -> ExecuteBashResult:
+    """Run a shell command in a running runtime container, streaming stdout/stderr and the exit code."""
+    if options.bearer_token:
+        raise RuntimeError("Bearer token auth for exec is not yet supported. Use SigV4 credentials.")
+
+    client = _create_agentcore_client(options.region, options.headers)
+    body: dict[str, Any] = {"command": options.command}
+    if options.timeout is not None:
+        body["timeout"] = options.timeout
+    params: dict[str, Any] = {"agentRuntimeArn": options.runtime_arn, "body": body}
+    if options.session_id is not None:
+        params["runtimeSessionId"] = options.session_id
+    response = client.invoke_agent_runtime_command(**params)
+
+    def stream_events() -> Iterator[ExecuteBashStreamEvent]:
+        event_stream = response.get("stream")
+        if event_stream is None:
+            raise RuntimeError("No stream in response from AgentCore Runtime")
+        for event in event_stream:
+            chunk = event.get("chunk") if isinstance(event, dict) else None
+            if not isinstance(chunk, dict):
+                continue
+            if chunk.get("contentStart") is not None:
+                yield {"type": "start"}
+            delta = chunk.get("contentDelta")
+            if delta:
+                if delta.get("stdout"):
+                    yield {"type": "stdout", "data": delta["stdout"]}
+                if delta.get("stderr"):
+                    yield {"type": "stderr", "data": delta["stderr"]}
+            stop = chunk.get("contentStop")
+            if stop is not None:
+                yield {"type": "stop", "exitCode": stop.get("exitCode"), "status": stop.get("status")}
+
+    return ExecuteBashResult(stream=stream_events(), session_id=response.get("runtimeSessionId"))