agentcore-cli 0.31.1__py3-none-any.whl

agentcore/operations/traces/get_trace.py

@@ -0,0 +1,202 @@
+"""Port of src/cli/operations/traces/get-trace.ts."""
+
+import json
+import os
+import re
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any
+
+from agentcore.aws.cloudwatch import runtime_log_group
+from agentcore.lib.errors.types import ResourceNotFoundError, ValidationError
+from agentcore.lib.result import Err, Ok
+from agentcore.operations.traces import insights_query
+from agentcore.operations.traces.types import (
+    CloudWatchSpanRecord,
+    CloudWatchTraceRecord,
+    FetchTraceRecordsValue,
+    GetTraceValue,
+)
+
+SPANS_LOG_GROUP = "aws/spans"
+TRACE_ID_PATTERN = re.compile(r"^[a-fA-F0-9-]+$")
+_INVALID_TRACE_ID = "Invalid trace ID format. Expected a hex string (e.g., abc123def456)."
+
+
+def _invalid_trace_id(trace_id: str) -> bool:
+    return not TRACE_ID_PATTERN.match(trace_id)
+
+
+def _number(value: str | None) -> float | int | None:
+    if not value:
+        return None
+    try:
+        number = float(value)
+    except ValueError:
+        return float("nan")
+    return int(number) if number.is_integer() else number
+
+
+def _fetch_spans(
+    region: str, runtime_id: str, trace_id: str, start_time: int | None, end_time: int | None
+) -> Ok[dict[str, list[CloudWatchSpanRecord]]] | Err[BaseException]:
+    if _invalid_trace_id(trace_id):
+        return Err(ValidationError(_INVALID_TRACE_ID))
+
+    query_string = f"""fields traceId, spanId, parentSpanId, name, kind,
+  startTimeUnixNano, endTimeUnixNano, durationNano,
+  status.code as statusCode,
+  resource.attributes.service.name as serviceName,
+  attributes.gen_ai.usage.input_tokens as inputTokens,
+  attributes.gen_ai.usage.output_tokens as outputTokens,
+  attributes.gen_ai.usage.total_tokens as totalTokens,
+  attributes.http.status_code as httpStatusCode,
+  attributes.session.id as sessionId
+| filter ispresent(traceId) and ispresent(resource.attributes.service.name)
+| filter resource.attributes.aws.service.type = "gen_ai_agent"
+| filter ispresent(kind)
+| filter traceId = '{trace_id}'
+| sort startTimeUnixNano asc"""
+
+    def run(log_group: str):
+        return insights_query.run_insights_query(
+            region=region, log_group_name=log_group, query_string=query_string, start_time=start_time, end_time=end_time
+        )
+
+    with ThreadPoolExecutor(max_workers=2) as pool:
+        results = list(pool.map(run, [SPANS_LOG_GROUP, runtime_log_group(runtime_id)]))
+
+    all_rows: list[dict[str, str]] = []
+    for result in results:
+        if result.success:
+            all_rows.extend(result.value["rows"])
+        elif isinstance(result.error, ResourceNotFoundError):
+            continue
+        else:
+            return Err(result.error)
+
+    spans: list[CloudWatchSpanRecord] = []
+    for row in all_rows:
+        if not (row.get("traceId") and row.get("spanId")):
+            continue
+        span: dict[str, Any] = {"traceId": row["traceId"], "spanId": row["spanId"]}
+        for key in (
+            "parentSpanId",
+            "name",
+            "kind",
+            "startTimeUnixNano",
+            "endTimeUnixNano",
+            "durationNano",
+            "statusCode",
+            "serviceName",
+        ):
+            if row.get(key) is not None:
+                span[key] = row[key]
+        for key in ("inputTokens", "outputTokens", "totalTokens", "httpStatusCode"):
+            number = _number(row.get(key))
+            if number is not None:
+                span[key] = number
+        if row.get("sessionId") is not None:
+            span["sessionId"] = row["sessionId"]
+        spans.append(span)  # type: ignore[arg-type]
+    return Ok({"spans": spans})
+
+
+def fetch_trace_records(
+    *,
+    region: str,
+    runtime_id: str,
+    trace_id: str,
+    start_time: int | None = None,
+    end_time: int | None = None,
+    include_spans: bool = False,
+) -> Ok[FetchTraceRecordsValue] | Err[BaseException]:
+    """Typed trace records (and optionally spans) for the web UI API."""
+    if _invalid_trace_id(trace_id):
+        return Err(ValidationError(_INVALID_TRACE_ID))
+
+    def records_query():
+        return insights_query.run_insights_query(
+            region=region,
+            log_group_name=runtime_log_group(runtime_id),
+            start_time=start_time,
+            end_time=end_time,
+            query_string=(
+                f"fields @timestamp, @message, @ptr\n| filter traceId = '{trace_id}'\n"
+                "| sort @timestamp asc\n| limit 10000"
+            ),
+        )
+
+    with ThreadPoolExecutor(max_workers=2) as pool:
+        records_future = pool.submit(records_query)
+        spans_future = (
+            pool.submit(_fetch_spans, region, runtime_id, trace_id, start_time, end_time) if include_spans else None
+        )
+        records_result = records_future.result()
+        spans_result = spans_future.result() if spans_future else None
+
+    if not records_result.success:
+        return Err(records_result.error)
+
+    trace_data = records_result.value["rows"]
+    if not trace_data and (spans_result is None or not spans_result.success or not spans_result.value["spans"]):
+        return Err(ResourceNotFoundError(f"No trace data found for trace ID: {trace_id}"))
+
+    records: list[CloudWatchTraceRecord] = []
+    for entry in trace_data:
+        raw = entry.get("@message", "{}")
+        try:
+            message: Any = json.loads(raw)
+        except ValueError:
+            message = raw
+        record: dict[str, Any] = {"@timestamp": entry.get("@timestamp", ""), "@message": message}
+        if entry.get("@ptr"):
+            record["@ptr"] = entry["@ptr"]
+        records.append(record)  # type: ignore[arg-type]
+
+    if spans_result is not None and spans_result.success:
+        return Ok({"records": records, "spans": spans_result.value["spans"]})
+    return Ok({"records": records})
+
+
+def get_trace(
+    *,
+    region: str,
+    runtime_id: str,
+    agent_name: str,
+    trace_id: str,
+    output_path: str | None = None,
+    start_time: int | None = None,
+    end_time: int | None = None,
+) -> Ok[GetTraceValue] | Err[BaseException]:
+    """Download a full trace and write the raw Insights rows (``@message`` parsed) to a JSON file."""
+    if _invalid_trace_id(trace_id):
+        return Err(ValidationError(_INVALID_TRACE_ID))
+
+    result = insights_query.run_insights_query(
+        region=region,
+        log_group_name=runtime_log_group(runtime_id),
+        start_time=start_time,
+        end_time=end_time,
+        query_string=f"fields @timestamp, @message\n| filter traceId = '{trace_id}'\n| sort @timestamp asc\n| limit 10000",
+    )
+    if not result.success:
+        return Err(result.error)
+
+    trace_data = result.value["rows"]
+    if not trace_data:
+        return Err(ResourceNotFoundError(f"No trace data found for trace ID: {trace_id}"))
+
+    parsed_trace = []
+    for entry in trace_data:
+        try:
+            parsed_trace.append({**entry, "@message": json.loads(entry.get("@message", "{}"))})
+        except ValueError:
+            parsed_trace.append(entry)
+
+    file_path = output_path or os.path.join("agentcore", ".cli", "traces", f"{agent_name}-{trace_id}.json")
+    directory = os.path.dirname(file_path)
+    if directory:
+        os.makedirs(directory, exist_ok=True)
+    with open(file_path, "w", encoding="utf-8") as f:
+        f.write(json.dumps(parsed_trace, indent=2, ensure_ascii=False))
+    return Ok({"filePath": os.path.abspath(file_path)})

agentcore/operations/traces/insights_query.py

@@ -0,0 +1,74 @@
+"""Port of src/cli/operations/traces/insights-query.ts (CloudWatch Logs Insights query + polling)."""
+
+import time
+from typing import Any
+
+from agentcore.aws._aws_errors import aws_error_name
+from agentcore.aws.account import create_aws_client
+from agentcore.lib.errors.types import ResourceNotFoundError, TimeoutError, to_error
+from agentcore.lib.result import Err, Ok
+from agentcore.operations.traces.types import InsightsQueryValue
+
+DEFAULT_LOOKBACK_MS = 12 * 60 * 60 * 1000
+
+InsightsQueryResult = Ok[InsightsQueryValue] | Err[BaseException]
+
+# Indirection so tests can skip the 1s poll interval.
+_sleep = time.sleep
+
+
+def _now_ms() -> int:
+    return int(time.time() * 1000)
+
+
+def _poll_query_results(client: Any, query_id: str) -> InsightsQueryResult:
+    for _ in range(60):
+        _sleep(1)
+        query_results = client.get_query_results(queryId=query_id)
+        status = query_results.get("status") or "Unknown"
+        if status in ("Complete", "Failed", "Cancelled"):
+            if status != "Complete":
+                return Err(Exception(f"Query {status.lower()}"))
+            rows: list[dict[str, str]] = []
+            for row in query_results.get("results") or []:
+                fields: dict[str, str] = {}
+                for item in row:
+                    if item.get("field") and item.get("value"):
+                        fields[item["field"]] = item["value"]
+                rows.append(fields)
+            return Ok({"rows": rows})
+    return Err(TimeoutError("Query timed out after 60 seconds"))
+
+
+def run_insights_query(
+    *,
+    region: str,
+    log_group_name: str,
+    query_string: str,
+    start_time: int | None = None,
+    end_time: int | None = None,
+) -> InsightsQueryResult:
+    """Run a Logs Insights query (times in epoch ms; default window: the 12h before ``end_time``)."""
+    end = end_time if end_time is not None else _now_ms()
+    start = start_time if start_time is not None else end - DEFAULT_LOOKBACK_MS
+    try:
+        client = create_aws_client("logs", region)
+        start_query = client.start_query(
+            logGroupName=log_group_name,
+            startTime=start // 1000,
+            endTime=end // 1000,
+            queryString=query_string,
+        )
+        query_id = start_query.get("queryId")
+        if not query_id:
+            return Err(Exception("Failed to start CloudWatch Logs Insights query"))
+        return _poll_query_results(client, query_id)
+    except Exception as exc:  # noqa: BLE001 - converted to a Result, as upstream
+        if aws_error_name(exc) == "ResourceNotFoundException":
+            return Err(
+                ResourceNotFoundError(
+                    f"Log group '{log_group_name}' not found. The agent may not have been invoked yet, "
+                    "or traces may not be enabled."
+                )
+            )
+        return Err(to_error(exc))

agentcore/operations/traces/list_traces.py

@@ -0,0 +1,46 @@
+"""Port of src/cli/operations/traces/list-traces.ts."""
+
+from agentcore.aws.cloudwatch import runtime_log_group
+from agentcore.lib.result import Err, Ok
+from agentcore.operations.traces import insights_query
+from agentcore.operations.traces.types import ListTracesValue, TraceEntry
+
+
+def list_traces(
+    *,
+    region: str,
+    runtime_id: str,
+    agent_name: str,
+    limit: int = 20,
+    start_time: int | None = None,
+    end_time: int | None = None,
+) -> Ok[ListTracesValue] | Err[BaseException]:
+    """Recent traces for a deployed runtime, via Logs Insights on its runtime log group."""
+    result = insights_query.run_insights_query(
+        region=region,
+        log_group_name=runtime_log_group(runtime_id),
+        start_time=start_time,
+        end_time=end_time,
+        query_string=(
+            "stats earliest(@timestamp) as firstSeen, latest(@timestamp) as lastSeen, count(*) as spanCount, "
+            "earliest(attributes.session.id) as sessionId by traceId\n"
+            "| sort lastSeen desc\n"
+            f"| limit {limit}"
+        ),
+    )
+    if not result.success:
+        return Err(result.error)
+
+    traces: list[TraceEntry] = []
+    for row in result.value["rows"]:
+        if row.get("traceId"):
+            entry: TraceEntry = {
+                "traceId": row["traceId"],
+                "timestamp": row.get("lastSeen") or row.get("firstSeen") or "unknown",
+            }
+            if row.get("sessionId") is not None:
+                entry["sessionId"] = row["sessionId"]
+            if row.get("spanCount") is not None:
+                entry["spanCount"] = row["spanCount"]
+            traces.append(entry)
+    return Ok({"traces": traces})

agentcore/operations/traces/trace_url.py

@@ -0,0 +1,25 @@
+"""Port of src/cli/operations/traces/trace-url.ts."""
+
+from urllib.parse import quote
+
+from agentcore.aws.partition import arn_prefix, console_domain
+
+# From upstream src/cli/constants.ts.
+DEFAULT_ENDPOINT_NAME = "DEFAULT"
+
+
+def _encode_uri_component(value: str) -> str:
+    return quote(value, safe="-_.!~*'()")
+
+
+def build_trace_console_url(*, region: str, account_id: str, runtime_id: str, agent_name: str) -> str:
+    """CloudWatch GenAI-observability console URL for an agent's traces."""
+    ep = DEFAULT_ENDPOINT_NAME
+    resource_id = _encode_uri_component(
+        f"{arn_prefix(region)}:bedrock-agentcore:{region}:{account_id}:runtime/{runtime_id}/runtime-endpoint/{ep}:{ep}"
+    )
+    return (
+        f"https://{region}.{console_domain(region)}/cloudwatch/home?region={region}"
+        f"#/gen-ai-observability/agent-core/agent-alias/{runtime_id}/endpoint/{ep}/agent/{agent_name}"
+        f"?start=-43200000&resourceId={resource_id}&serviceName={agent_name}.{ep}&tabId=traces"
+    )

agentcore/operations/traces/types.py

@@ -0,0 +1,55 @@
+"""Port of src/cli/operations/traces/types.ts.
+
+Records are TypedDicts keyed exactly like upstream's JSON (camelCase); optional fields that are
+``undefined`` upstream are omitted. Results use ``agentcore.lib.result`` ``Ok``/``Err``: the ``Ok``
+value is the dict upstream spreads onto its success branch.
+"""
+
+from typing import Any, NotRequired, TypedDict
+
+CloudWatchTraceRecord = TypedDict(
+    "CloudWatchTraceRecord",
+    {"@timestamp": str, "@message": Any, "@ptr": NotRequired[str]},
+)
+
+
+class CloudWatchSpanRecord(TypedDict):
+    traceId: str
+    spanId: str
+    parentSpanId: NotRequired[str]
+    name: NotRequired[str]
+    kind: NotRequired[str]
+    startTimeUnixNano: NotRequired[str]
+    endTimeUnixNano: NotRequired[str]
+    durationNano: NotRequired[str]
+    statusCode: NotRequired[str]
+    serviceName: NotRequired[str]
+    inputTokens: NotRequired[float]
+    outputTokens: NotRequired[float]
+    totalTokens: NotRequired[float]
+    httpStatusCode: NotRequired[float]
+    sessionId: NotRequired[str]
+
+
+class TraceEntry(TypedDict):
+    traceId: str
+    timestamp: str
+    sessionId: NotRequired[str]
+    spanCount: NotRequired[str]
+
+
+class ListTracesValue(TypedDict):
+    traces: list[TraceEntry]
+
+
+class FetchTraceRecordsValue(TypedDict):
+    records: list[CloudWatchTraceRecord]
+    spans: NotRequired[list[CloudWatchSpanRecord]]
+
+
+class GetTraceValue(TypedDict):
+    filePath: str
+
+
+class InsightsQueryValue(TypedDict):
+    rows: list[dict[str, str]]