agentcore-cli 0.31.1__py3-none-any.whl

agentcore/operations/jobs/ab_test/handler.py

@@ -0,0 +1,412 @@
+"""AB-test job handler (port of operations/jobs/ab-test/handler.ts).
+
+- create():  resolve region + gateway ARN, build variants + eval config, create (or reuse) the
+             execution role, ONE CreateABTest (AccessDenied retry while IAM propagates).
+- refresh(): GET latest state (404 -> NOT_FOUND).
+- stop/pause/resume(): UpdateABTest executionStatus = STOPPED / PAUSED / RUNNING.
+- promote(): dry-run validation, wait until RUNNING, stop, apply the winning variant.
+- archive(): stop -> poll STOPPED -> DeleteABTest -> delete the role if the CLI created it.
+"""
+
+import time
+from collections.abc import Callable
+from typing import Any
+
+from agentcore.aws import agentcore_ab_tests as api
+from agentcore.aws import agentcore_control as control
+from agentcore.aws._aws_errors import aws_error_name, aws_http_status
+from agentcore.aws.account import create_aws_client
+from agentcore.cli.errors import get_error_message
+from agentcore.lib import ConfigIO, JobNotFoundError, ResourceNotFoundError
+from agentcore.lib.errors.types import to_error
+from agentcore.lib.result import Err, Ok
+from agentcore.lib.schemas.io.config_io import to_jsonable
+from agentcore.lib.utils.json_format import to_json_string
+from agentcore.operations.jobs.ab_test import promote as _promote
+from agentcore.operations.jobs.ab_test import resolve as _resolve
+from agentcore.operations.jobs.ab_test.build_options import build_ab_test_request
+from agentcore.operations.jobs.batch_evaluation.handler import new_logger
+from agentcore.operations.jobs.shared.constants import NOT_FOUND_STATUS
+from agentcore.operations.jobs.shared.format import iso_now
+from agentcore.operations.jobs.shared.region import region_of, resolve_job_region
+from agentcore.operations.jobs.shared.resolve_agent_state import deployed_dict
+from agentcore.operations.jobs.shared.storage import drop_undefined
+
+MAX_CREATE_RETRIES = 5
+BASE_RETRY_DELAY_MS = 5_000
+
+# Injectable for tests.
+sleep: Callable[[float], None] = time.sleep
+
+
+def merge_deployed_resources(deployed_state: dict[str, Any]) -> dict[str, Any]:
+    """Merge per-target deployed resources (AB tests resolve names across all targets)."""
+    merged: dict[str, Any] = {}
+    for target in (deployed_state.get("targets") or {}).values():
+        r = (target or {}).get("resources")
+        if not r:
+            continue
+        merged["mcp"] = {**(merged.get("mcp") or {}), **(r.get("mcp") or {})}
+        for key in ("gateways", "configBundles", "onlineEvalConfigs"):
+            merged[key] = {**(merged.get(key) or {}), **(r.get(key) or {})}
+    return merged
+
+
+def _poll_until_stopped(region: str, ab_test_id: str, attempts: int = 20, delay_ms: int = 3_000) -> bool:
+    for _ in range(attempts):
+        try:
+            if api.get_ab_test(region=region, ab_test_id=ab_test_id).get("executionStatus") == "STOPPED":
+                return True
+        except JobNotFoundError:
+            return True
+        except Exception:  # noqa: BLE001, S110 - transient: keep polling
+            pass
+        sleep(delay_ms / 1000)
+    return False
+
+
+def _wait_for_running_then_stop(region: str, ab_test_id: str, attempts: int = 12, delay_ms: int = 10_000) -> None:
+    status = None
+    for _ in range(attempts):
+        status = api.get_ab_test(region=region, ab_test_id=ab_test_id).get("executionStatus")
+        if status == "RUNNING":
+            break
+        if status == "STOPPED":
+            return
+        sleep(delay_ms / 1000)
+    if status != "RUNNING":
+        raise Exception(f'A/B test "{ab_test_id}" did not reach RUNNING (current: {status}); cannot promote.')
+    api.update_ab_test(region=region, ab_test_id=ab_test_id, execution_status="STOPPED")
+
+
+def _is_retryable_create_error(exc: BaseException) -> bool:
+    message = str(exc)
+    return (
+        aws_error_name(exc) == "AccessDeniedException"
+        or aws_http_status(exc) == 403
+        or "Access denied" in message
+        or "Gateway validation error" in message
+    )
+
+
+def create(opts: dict[str, Any], config_io: ConfigIO) -> Ok | Err:
+    logger = new_logger("ab-test")
+    region = ""
+    role_arn: str | None = None
+    role_created_by_cli = False
+    on_progress = opts.get("onProgress") or (lambda _s, _m: None)
+    try:
+        if logger:
+            logger.start_step("Load project config")
+        project = to_jsonable(config_io.read_project_spec())
+        deployed_state = deployed_dict(config_io.read_deployed_state())
+        aws_targets = config_io.resolve_aws_deployment_targets()
+        region = resolve_job_region(opts.get("region"), aws_targets)
+        deployed = merge_deployed_resources(deployed_state)
+        if logger:
+            logger.end_step("success")
+            logger.start_step("Resolve gateway")
+        gateway_arn = _resolve.resolve_gateway_arn(opts["gateway"], deployed)
+        if not gateway_arn or not gateway_arn.startswith("arn:") or len(gateway_arn.split(":")) < 6:
+            err = ResourceNotFoundError(
+                f'Gateway "{opts["gateway"]}" is not deployed. Run `agentcore add gateway` and `agentcore deploy` first.'
+            )
+            if logger:
+                logger.end_step("error", err.message)
+                logger.finalize(False)
+            return Err(err)
+        if logger:
+            logger.log(f"Gateway ARN: {gateway_arn}")
+            logger.end_step("success")
+
+        built = build_ab_test_request(opts, deployed)
+
+        if logger:
+            logger.start_step("Resolve execution role")
+        if opts.get("roleArn"):
+            role_arn = opts["roleArn"]
+        else:
+            on_progress("role", "Creating execution role (waiting for IAM propagation)...")
+            role_arn = _resolve.get_or_create_ab_test_role(
+                region=region, project_name=project["name"], test_name=opts["name"], gateway_arn=gateway_arn
+            )
+            role_created_by_cli = True
+        if logger:
+            logger.log(f"Role ARN: {role_arn}")
+            logger.end_step("success")
+            logger.start_step("Create A/B test")
+        on_progress("starting", f'Creating A/B test "{opts["name"]}"...')
+
+        created = None
+        for attempt in range(MAX_CREATE_RETRIES):
+            try:
+                created = api.create_ab_test(
+                    region=region,
+                    name=f"{project['name']}_{opts['name']}",
+                    description=opts.get("description"),
+                    gateway_arn=gateway_arn,
+                    role_arn=role_arn,  # type: ignore[arg-type]
+                    variants=built["variants"],
+                    evaluation_config=built["evaluationConfig"],
+                    gateway_filter=built["gatewayFilter"],
+                    enable_on_create=opts.get("enableOnCreate"),
+                )
+                break
+            except Exception as exc:
+                if _is_retryable_create_error(exc) and attempt < MAX_CREATE_RETRIES - 1:
+                    on_progress("retry", f"Access not yet propagated; retrying (attempt {attempt + 2})...")
+                    sleep(BASE_RETRY_DELAY_MS * (2**attempt) / 1000)
+                    continue
+                raise
+        if not created:
+            raise Exception("A/B test creation failed after retries.")
+        if logger:
+            logger.log(f"Response: {to_json_string(created)}")
+            logger.end_step("success")
+        on_progress("started", f"A/B test created: {created['abTestId']} ({created['executionStatus']})")
+        if logger:
+            logger.finalize(True)
+
+        target_names = (
+            []
+            if opts["mode"] == "target-based"
+            else _resolve.resolve_runtime_target_names(
+                opts["gateway"], opts.get("runtime") or opts.get("agent"), project
+            )
+        )
+        record = {
+            "type": "ab-test",
+            "id": created["abTestId"],
+            "arn": created["abTestArn"],
+            "status": created["status"],
+            "lifecycleStatus": created["executionStatus"],
+            "createdAt": created.get("createdAt") or iso_now(),
+            "agent": opts.get("agent") or opts.get("runtime") or opts["name"],
+            "logFilePath": logger.log_file_path if logger else None,
+            "name": opts["name"],
+            "mode": opts["mode"],
+            "gatewayArn": gateway_arn,
+            "gatewayName": opts["gateway"],
+            "targetName": target_names[0] if len(target_names) == 1 else None,
+            "targetCandidates": target_names if len(target_names) > 1 else None,
+            "roleArn": role_arn,
+            "roleCreatedByCli": role_created_by_cli,
+            "variants": built["variantSummaries"],
+            "evaluationConfig": built["evaluationConfig"],
+            "gatewayFilter": built["gatewayFilter"],
+        }
+        return Ok(drop_undefined(record))
+    except Exception as exc:  # noqa: BLE001
+        if role_created_by_cli and role_arn and region:
+            try:
+                _resolve.delete_ab_test_role(region, role_arn)
+            except Exception:  # noqa: BLE001, S110
+                pass
+        if logger:
+            logger.finalize(False)
+        return Err(to_error(exc))
+
+
+def refresh(record: dict[str, Any]) -> Ok | Err:
+    try:
+        response = api.get_ab_test(region=region_of(record["arn"]), ab_test_id=record["id"])
+    except JobNotFoundError:
+        return Ok({**record, "status": NOT_FOUND_STATUS, "lifecycleStatus": NOT_FOUND_STATUS})
+    except Exception as exc:  # noqa: BLE001
+        return Err(to_error(exc))
+    failure = response.get("failureReason")
+    if failure is None and response.get("errorDetails") is not None:
+        failure = "; ".join(response["errorDetails"])
+    if failure is None:
+        failure = record.get("failureReason")
+    updated = {
+        **record,
+        "status": response.get("status"),
+        "lifecycleStatus": response.get("executionStatus"),
+        "completedAt": response.get("stoppedAt") or record.get("completedAt"),
+        "maxDurationExpiresAt": response.get("maxDurationExpiresAt") or record.get("maxDurationExpiresAt"),
+        "results": response.get("results") or record.get("results"),
+        "failureReason": failure,
+    }
+    return Ok(drop_undefined(updated))
+
+
+def _set_execution(record: dict[str, Any], status: str) -> None:
+    api.update_ab_test(region=region_of(record["arn"]), ab_test_id=record["id"], execution_status=status)  # type: ignore[arg-type]
+
+
+def stop(record: dict[str, Any]) -> Ok | Err:
+    try:
+        _set_execution(record, "STOPPED")
+        return Ok(None)
+    except Exception as exc:  # noqa: BLE001
+        return Err(to_error(exc))
+
+
+def pause(record: dict[str, Any]) -> Ok | Err:
+    try:
+        _set_execution(record, "PAUSED")
+        return Ok({**record, "lifecycleStatus": "PAUSED"})
+    except Exception as exc:  # noqa: BLE001
+        return Err(to_error(exc))
+
+
+def resume(record: dict[str, Any]) -> Ok | Err:
+    try:
+        _set_execution(record, "RUNNING")
+        return Ok({**record, "lifecycleStatus": "RUNNING"})
+    except Exception as exc:  # noqa: BLE001
+        return Err(to_error(exc))
+
+
+def promote(record: dict[str, Any], config_io: ConfigIO | None = None) -> Ok | Err:
+    region = region_of(record["arn"])
+    try:
+        preflight = _promote.promote_ab_test_config(record, True)
+        if not preflight["promoted"]:
+            return Err(
+                Exception(
+                    f'Cannot promote A/B test "{record["id"]}": {preflight["promotionDetail"]} '
+                    "The test was left running (not stopped)."
+                )
+            )
+        _wait_for_running_then_stop(region, record["id"])
+        promotion = _promote.promote_ab_test_config(record)
+        if not promotion["promoted"]:
+            return Err(
+                Exception(
+                    f'A/B test "{record["id"]}" was stopped, but the winning variant could not be applied to '
+                    f"agentcore.json: {promotion['promotionDetail']}"
+                )
+            )
+        return Ok({**record, "lifecycleStatus": "STOPPED"})
+    except Exception as exc:  # noqa: BLE001
+        return Err(to_error(exc))
+
+
+def archive(record: dict[str, Any]) -> Ok | Err:
+    region = region_of(record["arn"])
+    try:
+        try:
+            _set_execution(record, "STOPPED")
+            _poll_until_stopped(region, record["id"])
+        except Exception:  # noqa: BLE001, S110 - already stopped / transient: proceed to delete
+            pass
+        deleted = api.delete_ab_test(region=region, ab_test_id=record["id"])
+        if not deleted.get("success") and "404" not in (deleted.get("error") or ""):
+            return Err(Exception(deleted.get("error") or "Failed to delete A/B test."))
+        if record.get("roleCreatedByCli") and record.get("roleArn"):
+            try:
+                all_tests = api.list_ab_tests(region=region, max_results=100)
+                active = [
+                    t
+                    for t in all_tests["abTests"]
+                    if t.get("abTestId") != record["id"]
+                    and t.get("status") not in ("STOPPED", "CREATE_FAILED", "DELETE_FAILED")
+                ]
+                in_use = any(
+                    api.get_ab_test(region=region, ab_test_id=t["abTestId"]).get("roleArn") == record["roleArn"]
+                    for t in active
+                )
+                if not in_use:
+                    _resolve.delete_ab_test_role(region, record["roleArn"])
+            except Exception:  # noqa: BLE001, S110 - can't verify: keep the role
+                pass
+        return Ok(None)
+    except JobNotFoundError:
+        return Ok(None)
+    except Exception as exc:  # noqa: BLE001
+        return Err(to_error(exc))
+
+
+def debug(record: dict[str, Any]) -> Ok | Err:
+    region = region_of(record["arn"])
+    checks: list[dict[str, str]] = []
+    try:
+        test = api.get_ab_test(region=region, ab_test_id=record["id"])
+        checks.append(
+            {
+                "label": "AB Test Status",
+                "status": "pass"
+                if test.get("status") == "ACTIVE" and test.get("executionStatus") == "RUNNING"
+                else "warn",
+                "detail": f"{test.get('status')} / {test.get('executionStatus')}",
+            }
+        )
+    except Exception as exc:  # noqa: BLE001
+        checks.append({"label": "AB Test Status", "status": "fail", "detail": get_error_message(exc)})
+        return Ok({"checks": checks})
+
+    checks.append(
+        {
+            "label": "AB Test Role",
+            "status": "pass" if test.get("roleArn") else "warn",
+            "detail": test.get("roleArn") or "No role ARN",
+        }
+    )
+
+    evaluation_config = test.get("evaluationConfig") or {}
+    if "perVariantOnlineEvaluationConfig" in evaluation_config:
+        eval_arns = [
+            (v.get("name", ""), v.get("onlineEvaluationConfigArn", ""))
+            for v in evaluation_config["perVariantOnlineEvaluationConfig"]
+        ]
+    else:
+        eval_arns = [("", evaluation_config.get("onlineEvaluationConfigArn", ""))]
+    for variant_name, eval_arn in eval_arns:
+        config_id = eval_arn.split("/")[-1] or eval_arn
+        suffix = f" ({variant_name})" if variant_name else ""
+        try:
+            config = control.get_online_evaluation_config(region=region, config_id=config_id)
+            checks.append(
+                {
+                    "label": f"Online Eval Config{suffix}",
+                    "status": "pass" if config.get("executionStatus") == "ENABLED" else "fail",
+                    "detail": f"{config.get('configName')} — {config.get('executionStatus')}",
+                }
+            )
+        except Exception as exc:  # noqa: BLE001
+            checks.append({"label": f"Online Eval Config{suffix}", "status": "fail", "detail": get_error_message(exc)})
+
+    gateway_id = (test.get("gatewayArn") or "").split("/")[-1]
+    try:
+        gateway = control.get_gateway_detail(region=region, gateway_id=gateway_id)
+        checks.append(
+            {
+                "label": "Gateway Role",
+                "status": "pass" if gateway.get("roleArn") else "warn",
+                "detail": gateway.get("roleArn") or "No role ARN",
+            }
+        )
+    except Exception as exc:  # noqa: BLE001
+        checks.append({"label": "Gateway Role", "status": "fail", "detail": get_error_message(exc)})
+
+    two_hours_ago = int(time.time() * 1000) - 2 * 60 * 60 * 1000
+    try:
+        logs = create_aws_client("logs", region)
+        for v in test.get("variants") or []:
+            name = v.get("name")
+            try:
+                response = logs.filter_log_events(
+                    logGroupName="aws/spans",
+                    startTime=two_hours_ago,
+                    filterPattern=f'"{test.get("abTestArn")}" "{name}"',
+                    limit=5,
+                )
+                count = len(response.get("events") or [])
+                checks.append(
+                    {
+                        "label": f"Experiment Spans — {name} (2h)",
+                        "status": "pass" if count > 0 else "warn",
+                        "detail": f"{count}+ spans with experiment metadata"
+                        if count > 0
+                        else "No spans found — traffic may not be reaching this variant",
+                    }
+                )
+            except Exception as exc:  # noqa: BLE001
+                checks.append(
+                    {"label": f"Experiment Spans — {name}", "status": "warn", "detail": get_error_message(exc)}
+                )
+    except Exception as exc:  # noqa: BLE001
+        checks.append({"label": "Experiment Spans", "status": "warn", "detail": get_error_message(exc)})
+    return Ok({"checks": checks})

agentcore/operations/jobs/ab_test/promote.py

@@ -0,0 +1,142 @@
+"""Apply the winning A/B variant to agentcore.json (port of operations/jobs/ab-test/promote.ts)."""
+
+import copy
+from typing import Any
+
+from agentcore.aws import agentcore_config_bundles as bundles_api
+from agentcore.lib import ConfigIO
+from agentcore.lib.schemas.io.config_io import to_jsonable
+from agentcore.operations.jobs.recommendation.build_config import resolve_component_key_for_json_path
+from agentcore.operations.jobs.shared.region import region_from_arn
+from agentcore.operations.jobs.shared.resolve_agent_state import deployed_dict
+
+
+def _bundle_id_from_arn(arn: str) -> str | None:
+    tail = arn.split("/")[-1]
+    return tail or None
+
+
+def restore_placeholder_keys(
+    service_components: dict[str, Any], local_components: dict[str, Any] | None, deployed_state: dict[str, Any]
+) -> dict[str, Any]:
+    """Rewrite service ARN keys back to the bundle's portable ``{{runtime:...}}`` placeholders."""
+    arn_to_placeholder: dict[str, str] = {}
+    for key in local_components or {}:
+        if key.startswith("arn:"):
+            continue
+        arn = resolve_component_key_for_json_path(key, deployed_state)
+        if arn != key:
+            arn_to_placeholder[arn] = key
+    return {arn_to_placeholder.get(k, k): v for k, v in service_components.items()}
+
+
+def _bundle_name_from_arn(deployed_state: dict[str, Any], bundle_arn: str) -> str | None:
+    for target in (deployed_state.get("targets") or {}).values():
+        for name, entry in ((((target or {}).get("resources")) or {}).get("configBundles") or {}).items():
+            if entry.get("bundleArn") == bundle_arn:
+                return name
+    return None
+
+
+def _result(promoted: bool, mode: str | None, detail: str) -> dict[str, Any]:
+    return {"promoted": promoted, "mode": mode, "promotionDetail": detail}
+
+
+def promote_ab_test_config(
+    record: dict[str, Any], dry_run: bool = False, config_io: ConfigIO | None = None
+) -> dict[str, Any]:
+    """{"promoted", "mode", "promotionDetail"}. ``dry_run`` validates without writing agentcore.json."""
+    config_io = config_io or ConfigIO()
+    project = to_jsonable(config_io.read_project_spec())
+    mode = record.get("mode")
+    variants = record.get("variants") or []
+    control = next((v for v in variants if v.get("name") == "C"), None)
+    treatment = next((v for v in variants if v.get("name") == "T1"), None)
+    if not control or not treatment:
+        return _result(False, mode, "A/B test record is missing control (C) or treatment (T1) variant.")
+
+    if mode == "target-based":
+        gateway_name = record.get("gatewayName")
+        if not gateway_name:
+            return _result(False, mode, "A/B test record is missing the gateway name; cannot locate targets.")
+        gateway = next((g for g in project.get("agentCoreGateways") or [] if g.get("name") == gateway_name), None)
+        if not gateway or gateway.get("targets") is None:
+            return _result(False, mode, f'Gateway "{gateway_name}" not found in agentcore.json.')
+        targets = gateway["targets"]
+        control_t = next((t for t in targets if t.get("name") == control.get("targetName")), None)
+        treatment_t = next((t for t in targets if t.get("name") == treatment.get("targetName")), None)
+        c_rt = (control_t or {}).get("httpRuntime") or {}
+        t_rt = (treatment_t or {}).get("httpRuntime") or {}
+        if not c_rt.get("runtime") or not t_rt.get("runtime"):
+            return _result(False, mode, "Could not resolve control/treatment runtime targets for promotion.")
+
+        c_ep, t_ep = c_rt.get("runtimeEndpoint"), t_rt.get("runtimeEndpoint")
+        if c_rt["runtime"] == t_rt["runtime"] and c_ep and t_ep:
+            runtime = next((r for r in project.get("runtimes") or [] if r.get("name") == c_rt["runtime"]), None)
+            endpoints = (runtime or {}).get("endpoints") or {}
+            control_ep, treatment_ep = endpoints.get(c_ep), endpoints.get(t_ep)
+            if control_ep and treatment_ep:
+                if not dry_run:
+                    control_ep["version"] = treatment_ep.get("version")
+                    config_io.write_project_spec(project)
+                return _result(
+                    True,
+                    mode,
+                    f'Control endpoint "{c_ep}" updated to version {treatment_ep.get("version")} '
+                    f'(from treatment "{t_ep}").',
+                )
+
+        if not dry_run:
+            control_t["httpRuntime"] = copy.deepcopy(t_rt)  # type: ignore[index]
+            config_io.write_project_spec(project)
+        treatment_ref = f'{t_rt["runtime"]} (endpoint "{t_ep}")' if t_ep else t_rt["runtime"]
+        return _result(
+            True,
+            mode,
+            f'Control target "{control_t["name"]}" repointed to treatment runtime {treatment_ref}.',  # type: ignore[index]
+        )
+
+    if not control.get("bundleArn") or not treatment.get("bundleArn"):
+        return _result(False, mode, "A/B test record is missing control/treatment bundle ARNs.")
+    if control["bundleArn"] != treatment["bundleArn"]:
+        return _result(
+            False,
+            mode,
+            "Cannot promote: control and treatment reference different config bundles. "
+            "A config-bundle A/B test can only promote between two versions of the SAME bundle.",
+        )
+    if not treatment.get("bundleVersion"):
+        return _result(False, mode, "A/B test record is missing the treatment bundle version; cannot promote.")
+
+    control_name = None
+    deployed_state = None
+    try:
+        deployed_state = deployed_dict(config_io.read_deployed_state())
+        control_name = _bundle_name_from_arn(deployed_state, control["bundleArn"])
+    except Exception:  # noqa: BLE001, S110 - deployed state unavailable
+        pass
+    if not control_name or deployed_state is None:
+        return _result(
+            False, mode, "Could not resolve the config bundle from deployed state (deploy the bundle first)."
+        )
+    control_bundle = next((b for b in project.get("configBundles") or [] if b.get("name") == control_name), None)
+    if not control_bundle:
+        return _result(False, mode, f'Could not find config bundle "{control_name}" in agentcore.json.')
+    bundle_id = _bundle_id_from_arn(treatment["bundleArn"])
+    if not bundle_id:
+        return _result(False, mode, f'Could not parse bundle id from ARN "{treatment["bundleArn"]}".')
+
+    if not dry_run:
+        region = region_from_arn(treatment["bundleArn"]) or region_from_arn(record.get("arn", ""))
+        if not region:
+            return _result(False, mode, "Could not determine region for the config bundle.")
+        winning = bundles_api.get_configuration_bundle_version(
+            region=region, bundle_id=bundle_id, version_id=treatment["bundleVersion"]
+        )
+        control_bundle["components"] = restore_placeholder_keys(
+            winning.get("components") or {}, control_bundle.get("components"), deployed_state
+        )
+        config_io.write_project_spec(project)
+    return _result(
+        True, mode, f'Config bundle "{control_name}" updated to the winning version {treatment["bundleVersion"]}.'
+    )