deploy.py

#!/usr/bin/env python3
"""Deploy MCP tools as an AWS Lambda and register it as an AgentCore gateway target.

A boto3 CLI that handles the full lifecycle: per-function IAM role, container or
zip packaging, Lambda create/update, environment variables, and gateway target
registration. It runs each step in order and uses waiters to gate on readiness.

Optional per-project convention files (both support ${VAR} expansion from the
current environment, so dynamic values like ARNs resolved by direnv/.envrc flow
straight through):

  env.json     Lambda environment variables in AWS format:
                 {"Variables": {"EXAMPLE_STATIC": "value",
                                "EXAMPLE_FROM_ENV": "${SOME_SHELL_VAR}"}}
  policy.json  An IAM policy document attached inline to the function's role:
                 {"Version": "2012-10-17", "Statement": [ ... ]}

Commands:
  schema   Generate tools.json from the @tool-decorated functions
  create   Create the Lambda + role (+ env/policy) and register the gateway target
  deploy   Update code (+ env/policy) and update the gateway target
  delete   Delete the Lambda, its per-function role, and the gateway target
  invoke   Invoke a tool locally against the deployed Lambda
"""

from __future__ import annotations

import argparse
import base64
import json
import os
import subprocess
import sys
import time
from pathlib import Path

import boto3
from botocore.exceptions import ClientError

HERE = Path(__file__).resolve().parent
REGION = os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION") or "us-east-1"


# --------------------------------------------------------------------------- #
# helpers
# --------------------------------------------------------------------------- #
def log(msg: str) -> None:
    print(f"[deploy] {msg}", flush=True)


def die(msg: str) -> "None":
    print(f"[deploy] ERROR: {msg}", file=sys.stderr, flush=True)
    sys.exit(1)


def run(cmd: list[str], **kwargs) -> subprocess.CompletedProcess:
    log("$ " + " ".join(cmd))
    return subprocess.run(cmd, check=True, **kwargs)


def client(name: str):
    return boto3.client(name, region_name=REGION)


def account_id() -> str:
    return client("sts").get_caller_identity()["Account"]


def python_version() -> str:
    """Read the python major.minor from .tool-versions (fallback 3.13)."""
    tv = HERE / ".tool-versions"
    if tv.exists():
        for line in tv.read_text().splitlines():
            parts = line.split()
            if len(parts) == 2 and parts[0] == "python":
                return ".".join(parts[1].split(".")[:2])
    return "3.13"


def _expand(obj):
    """Recursively expand ${VAR} / $VAR in all string values from the environment."""
    if isinstance(obj, str):
        return os.path.expandvars(obj)
    if isinstance(obj, list):
        return [_expand(v) for v in obj]
    if isinstance(obj, dict):
        return {k: _expand(v) for k, v in obj.items()}
    return obj


def load_config(filename: str):
    """Load an optional JSON config file with ${VAR} expansion. None if absent."""
    path = HERE / filename
    if not path.exists():
        return None
    data = json.loads(path.read_text())
    expanded = _expand(data)
    # Fail loudly on unresolved placeholders so we never ship a literal "${VAR}".
    leftovers = [s for s in _iter_strings(expanded) if "${" in s]
    if leftovers:
        die(f"{filename} has unresolved variables: {leftovers}. "
            f"Set them in your environment (see .envrc) before deploying.")
    return expanded


def _iter_strings(obj):
    if isinstance(obj, str):
        yield obj
    elif isinstance(obj, list):
        for v in obj:
            yield from _iter_strings(v)
    elif isinstance(obj, dict):
        for v in obj.values():
            yield from _iter_strings(v)


# --------------------------------------------------------------------------- #
# schema
# --------------------------------------------------------------------------- #
def generate_schema() -> list[dict]:
    """Import the tools module and export the MCP tool schema to tools.json."""
    sys.path.insert(0, str(HERE))
    import tools  # noqa: F401  -- registers the tool handlers
    from registry import auto_register, export_schema

    auto_register(tools)
    schema = export_schema()
    (HERE / "tools.json").write_text(json.dumps(schema, indent=2))
    log(f"generated tools.json with {len(schema)} tool(s): "
        + ", ".join(t["name"] for t in schema))
    return schema


# --------------------------------------------------------------------------- #
# IAM role (one per function so per-project policies never collide)
# --------------------------------------------------------------------------- #
BASIC_EXECUTION = "arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole"
INLINE_POLICY_NAME = "tool-permissions"
ASSUME_ROLE_DOC = json.dumps({
    "Version": "2012-10-17",
    "Statement": [{
        "Effect": "Allow",
        "Principal": {"Service": "lambda.amazonaws.com"},
        "Action": "sts:AssumeRole",
    }],
})


def ensure_role(role_name: str) -> str:
    """Create (idempotently) the execution role and attach basic-exec + policy.json."""
    iam = client("iam")
    created = False
    try:
        iam.create_role(RoleName=role_name, AssumeRolePolicyDocument=ASSUME_ROLE_DOC)
        created = True
        log(f"created role {role_name}")
    except iam.exceptions.EntityAlreadyExistsException:
        log(f"role {role_name} already exists")

    iam.attach_role_policy(RoleName=role_name, PolicyArn=BASIC_EXECUTION)

    policy = load_config("policy.json")
    if policy:
        iam.put_role_policy(
            RoleName=role_name,
            PolicyName=INLINE_POLICY_NAME,
            PolicyDocument=json.dumps(policy),
        )
        log(f"attached inline policy {INLINE_POLICY_NAME} to {role_name}")
    else:
        # Remove a stale inline policy if policy.json was deleted.
        try:
            iam.delete_role_policy(RoleName=role_name, PolicyName=INLINE_POLICY_NAME)
            log(f"removed inline policy {INLINE_POLICY_NAME} (no policy.json present)")
        except iam.exceptions.NoSuchEntityException:
            pass

    role_arn = iam.get_role(RoleName=role_name)["Role"]["Arn"]
    if created:
        log("waiting for role to propagate...")
        time.sleep(10)
    return role_arn


# --------------------------------------------------------------------------- #
# container image (ECR)
# --------------------------------------------------------------------------- #
def platform_for(arch: str) -> str:
    return "linux/arm64" if arch == "arm64" else "linux/amd64"


def build_and_push_image(function: str, arch: str) -> str:
    """Ensure the ECR repo exists, then build and push the image. The repo is
    created before the image URI is resolved, so the tag is always valid."""
    ecr = client("ecr")
    acct = account_id()
    registry = f"{acct}.dkr.ecr.{REGION}.amazonaws.com"

    try:
        ecr.describe_repositories(repositoryNames=[function])
    except ecr.exceptions.RepositoryNotFoundException:
        ecr.create_repository(repositoryName=function)
        log(f"created ECR repository {function}")

    uri = ecr.describe_repositories(
        repositoryNames=[function])["repositories"][0]["repositoryUri"]
    image = f"{uri}:latest"

    # docker login using an ECR auth token
    token = ecr.get_authorization_token()["authorizationData"][0]["authorizationToken"]
    user, password = base64.b64decode(token).decode().split(":", 1)
    run(["docker", "login", "--username", user, "--password-stdin", registry],
        input=password.encode())

    run([
        "docker", "buildx", "build", "--push", "--provenance=false",
        "--platform", platform_for(arch),
        "--build-arg", f"PYTHON_VERSION={python_version()}",
        "-t", image, ".",
    ], cwd=HERE)
    log(f"pushed image {image}")
    return image


# --------------------------------------------------------------------------- #
# zip package
# --------------------------------------------------------------------------- #
def build_zip() -> Path:
    """Build lambda.zip from the installed venv site-packages plus the *.py files."""
    import zipfile

    pyver = python_version()
    site = HERE / ".venv" / "lib" / f"python{pyver}" / "site-packages"
    if not site.exists():
        die(f"{site} not found. Run 'make install' first (zip packaging).")
    zip_path = HERE / "lambda.zip"
    if zip_path.exists():
        zip_path.unlink()
    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
        for f in site.rglob("*"):
            if f.is_file():
                zf.write(f, f.relative_to(site))
        for py in HERE.glob("*.py"):
            zf.write(py, py.name)
    log(f"built {zip_path.name}")
    return zip_path


# --------------------------------------------------------------------------- #
# Lambda
# --------------------------------------------------------------------------- #
def function_exists(lam, function: str) -> bool:
    try:
        lam.get_function(FunctionName=function)
        return True
    except lam.exceptions.ResourceNotFoundException:
        return False


def wait_active(lam, function: str) -> None:
    lam.get_waiter("function_active_v2").wait(FunctionName=function)


def wait_updated(lam, function: str) -> None:
    lam.get_waiter("function_updated_v2").wait(FunctionName=function)


def apply_env(lam, function: str) -> None:
    env = load_config("env.json")
    if not env:
        return
    variables = env.get("Variables", env)  # accept either {"Variables":{...}} or {...}
    wait_updated(lam, function)
    lam.update_function_configuration(
        FunctionName=function, Environment={"Variables": variables})
    wait_updated(lam, function)
    log(f"applied {len(variables)} environment variable(s) from env.json")


def create_function(lam, function: str, role_arn: str, arch: str, packaging: str) -> None:
    if packaging == "container":
        image = build_and_push_image(function, arch)
        lam.create_function(
            FunctionName=function, PackageType="Image",
            Code={"ImageUri": image}, Architectures=[arch], Role=role_arn)
    else:
        zip_path = build_zip()
        lam.create_function(
            FunctionName=function, Runtime=f"python{python_version()}",
            Handler="lambda_function.lambda_handler", Role=role_arn,
            Architectures=[arch], Code={"ZipFile": zip_path.read_bytes()})
    log(f"created function {function}")
    wait_active(lam, function)


def update_function(lam, function: str, role_arn: str, arch: str, packaging: str) -> None:
    if packaging == "container":
        image = build_and_push_image(function, arch)
        lam.update_function_code(FunctionName=function, ImageUri=image)
    else:
        zip_path = build_zip()
        lam.update_function_code(FunctionName=function, ZipFile=zip_path.read_bytes())
    wait_updated(lam, function)
    # keep the role in sync (e.g. if it changed) without failing on identical values
    lam.update_function_configuration(FunctionName=function, Role=role_arn)
    wait_updated(lam, function)
    log(f"updated function code for {function}")


# --------------------------------------------------------------------------- #
# gateway target
# --------------------------------------------------------------------------- #
def ensure_bucket(function: str) -> str:
    s3 = client("s3")
    bucket = f"agentcore-tools-{account_id()}-{REGION}"
    try:
        s3.head_bucket(Bucket=bucket)
    except ClientError:
        if REGION == "us-east-1":
            s3.create_bucket(Bucket=bucket)
        else:
            s3.create_bucket(
                Bucket=bucket,
                CreateBucketConfiguration={"LocationConstraint": REGION})
        log(f"created bucket {bucket}")
    return bucket


def upload_schema(function: str) -> str:
    bucket = ensure_bucket(function)
    key = f"{function}/tools.json"
    client("s3").upload_file(str(HERE / "tools.json"), bucket, key)
    log(f"uploaded schema to s3://{bucket}/{key}")
    return f"s3://{bucket}/{key}"


def target_config(function: str, s3_uri: str) -> dict:
    return {"mcp": {"lambda": {
        "lambdaArn": f"arn:aws:lambda:{REGION}:{account_id()}:function:{function}",
        "toolSchema": {"s3": {"uri": s3_uri, "bucketOwnerAccountId": account_id()}},
    }}}


def find_target(agc, gateway_id: str, name: str) -> str | None:
    targets = agc.list_gateway_targets(gatewayIdentifier=gateway_id).get("items", [])
    for t in targets:
        if t.get("name") == name:
            return t.get("targetId")
    return None


def allow_gateway_invoke(lam, function: str, gateway_id: str) -> None:
    acct = account_id()
    for stmt_id, principal, extra in [
        ("allow-agentcore-gateway", "bedrock-agentcore.amazonaws.com",
         {"SourceAccount": acct}),
        ("allow-gateway-role",
         client("bedrock-agentcore-control").get_gateway(
             gatewayIdentifier=gateway_id)["roleArn"], {}),
    ]:
        try:
            lam.add_permission(
                FunctionName=function, StatementId=stmt_id,
                Action="lambda:InvokeFunction", Principal=principal, **extra)
        except lam.exceptions.ResourceConflictException:
            pass
    log("granted gateway invoke permissions")


def register_target(gateway_id: str, name: str, function: str) -> None:
    generate_schema()
    s3_uri = upload_schema(function)
    agc = client("bedrock-agentcore-control")
    cfg = target_config(function, s3_uri)
    cred = [{"credentialProviderType": "GATEWAY_IAM_ROLE"}]
    existing = find_target(agc, gateway_id, name)

    # The Lambda may take a beat to become invokable by the gateway; retry briefly.
    for attempt in range(1, 7):
        try:
            if existing:
                agc.update_gateway_target(
                    gatewayIdentifier=gateway_id, targetId=existing, name=name,
                    targetConfiguration=cfg, credentialProviderConfigurations=cred)
                log(f"updated gateway target {name} ({existing})")
            else:
                resp = agc.create_gateway_target(
                    gatewayIdentifier=gateway_id, name=name,
                    targetConfiguration=cfg, credentialProviderConfigurations=cred)
                log(f"registered gateway target {name} ({resp.get('targetId')})")
            return
        except ClientError as e:
            msg = str(e)
            if "not ready" in msg or "resource conflict" in msg.lower():
                log(f"lambda not ready yet, retrying ({attempt}/6)...")
                time.sleep(10)
                continue
            raise
    die("gateway target registration timed out waiting for the Lambda to be ready")


# --------------------------------------------------------------------------- #
# commands
# --------------------------------------------------------------------------- #
def cmd_schema(args) -> None:
    generate_schema()


def cmd_create(args) -> None:
    lam = client("lambda")
    role_arn = ensure_role(args.role or f"{args.function}-role")
    if function_exists(lam, args.function):
        log(f"function {args.function} exists; updating instead")
        update_function(lam, args.function, role_arn, args.arch, args.packaging)
    else:
        create_function(lam, args.function, role_arn, args.arch, args.packaging)
    apply_env(lam, args.function)
    allow_gateway_invoke(lam, args.function, args.gateway_id)
    register_target(args.gateway_id, args.target or args.function, args.function)


def cmd_deploy(args) -> None:
    lam = client("lambda")
    role_arn = ensure_role(args.role or f"{args.function}-role")
    if not function_exists(lam, args.function):
        die(f"function {args.function} does not exist; run 'create' first")
    update_function(lam, args.function, role_arn, args.arch, args.packaging)
    apply_env(lam, args.function)
    register_target(args.gateway_id, args.target or args.function, args.function)


def cmd_delete(args) -> None:
    lam = client("lambda")
    agc = client("bedrock-agentcore-control")
    name = args.target or args.function
    if args.gateway_id:
        tid = find_target(agc, args.gateway_id, name)
        if tid:
            agc.delete_gateway_target(gatewayIdentifier=args.gateway_id, targetId=tid)
            log(f"deleted gateway target {name} ({tid})")
    try:
        lam.delete_function(FunctionName=args.function)
        log(f"deleted function {args.function}")
    except lam.exceptions.ResourceNotFoundException:
        log(f"function {args.function} not found, skipping")

    role_name = args.role or f"{args.function}-role"
    iam = client("iam")
    try:
        for p in iam.list_role_policies(RoleName=role_name).get("PolicyNames", []):
            iam.delete_role_policy(RoleName=role_name, PolicyName=p)
        for p in iam.list_attached_role_policies(RoleName=role_name).get(
                "AttachedPolicies", []):
            iam.detach_role_policy(RoleName=role_name, PolicyArn=p["PolicyArn"])
        iam.delete_role(RoleName=role_name)
        log(f"deleted role {role_name}")
    except iam.exceptions.NoSuchEntityException:
        pass


def cmd_invoke(args) -> None:
    lam = client("lambda")
    name = args.target or args.function
    payload = Path(args.payload).read_bytes() if Path(args.payload).exists() else b"{}"
    ctx = base64.b64encode(json.dumps({"custom": {
        "bedrockAgentCoreToolName": f"{name}___{args.tool}",
        "bedrockAgentCoreGatewayId": "invoke-test",
        "bedrockAgentCoreTargetId": "invoke-test",
        "bedrockAgentCoreAwsRequestId": "invoke-test",
        "bedrockAgentCoreMcpMessageId": "invoke-test",
        "bedrockAgentCoreMessageVersion": "1.0",
    }}).encode()).decode()
    resp = lam.invoke(FunctionName=args.function, Payload=payload, ClientContext=ctx)
    body = resp["Payload"].read().decode()
    if resp.get("FunctionError"):
        print(f"FunctionError: {resp['FunctionError']}")
    try:
        print(json.dumps(json.loads(body), indent=2))
    except json.JSONDecodeError:
        print(body)


# --------------------------------------------------------------------------- #
# CLI
# --------------------------------------------------------------------------- #
def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(description=__doc__,
                                formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = p.add_subparsers(dest="command", required=True)

    def add_common(sp, gateway=True):
        sp.add_argument("--function", required=True, help="Lambda function name")
        sp.add_argument("--target", help="Gateway target name (default: function name)")
        sp.add_argument("--role", help="Execution role name (default: <function>-role)")
        sp.add_argument("--arch", default="arm64", choices=["arm64", "x86_64"])
        sp.add_argument("--packaging", default="container",
                        choices=["container", "zip"])
        if gateway:
            sp.add_argument("--gateway-id", required=True, help="AgentCore gateway id")

    sub.add_parser("schema", help="Generate tools.json")

    add_common(sub.add_parser("create", help="Create Lambda + register gateway target"))
    add_common(sub.add_parser("deploy", help="Update code + update gateway target"))

    d = sub.add_parser("delete", help="Delete Lambda, role, and gateway target")
    d.add_argument("--function", required=True)
    d.add_argument("--target")
    d.add_argument("--role")
    d.add_argument("--gateway-id")

    i = sub.add_parser("invoke", help="Invoke a tool on the deployed Lambda")
    i.add_argument("--function", required=True)
    i.add_argument("--target")
    i.add_argument("--tool", required=True)
    i.add_argument("--payload", default="request.json")

    return p


def main() -> None:
    args = build_parser().parse_args()
    {
        "schema": cmd_schema,
        "create": cmd_create,
        "deploy": cmd_deploy,
        "delete": cmd_delete,
        "invoke": cmd_invoke,
    }[args.command](args)


if __name__ == "__main__":
    main()