I tried out "Temporal Policies" added to Amazon Bedrock AgentCore Policy

"""Create the verification environment for the article and save the created resource IDs to a file."""

import io, json, time, zipfile
from pathlib import Path
import boto3

# Creation destination is Tokyo region. Save progress and load already-created information on re-execution.
ROOT = Path(__file__).resolve().parents[1]
STATE = ROOT / ".gateway-state.json"
s = (
    json.loads(STATE.read_text())
    if STATE.exists()
    else {"region": "ap-northeast-1", "name": f"dogwood-shopping-{int(time.time())}"}
)
if s.get("cleanedUp"):
    s = {"region": s["region"], "name": f"dogwood-shopping-{int(time.time())}"}
session = boto3.Session(region_name=s["region"])
c = session.client("bedrock-agentcore-control")
iam = session.client("iam")
lam = session.client("lambda")
s["account"] = session.client("sts").get_caller_identity()["Account"]
base = f"arn:aws:bedrock-agentcore:{s['region']}:{s['account']}"

def save():
    STATE.write_text(json.dumps(s, indent=2, default=str))

def wait(get, key, good=("READY", "ACTIVE"), **args):
    for _ in range(120):
        r = get(**args)
        status = r[key]
        if status in good:
            return r
        if "FAIL" in status:
            raise RuntimeError(json.dumps(r, default=str))
        time.sleep(5)
    raise TimeoutError(args)

def role(key, service):
    if key not in s:
        statement = {
            "Effect": "Allow",
            "Principal": {"Service": service},
            "Action": "sts:AssumeRole",
        }
        if service == "bedrock-agentcore.amazonaws.com":
            statement["Condition"] = {
                "StringEquals": {"aws:SourceAccount": s["account"]},
                "ArnLike": {"aws:SourceArn": base + ":*"},
            }
        r = iam.create_role(
            RoleName=s["name"] + "-" + key,
            AssumeRolePolicyDocument=json.dumps(
                {"Version": "2012-10-17", "Statement": [statement]}
            ),
        )
        s[key] = r["Role"]["Arn"]
        save()
    return s[key]

# 1. Create a Policy Engine to register shopping rules.
if "engineId" not in s:
    r = c.create_policy_engine(name=s["name"].replace("-", "_"))
    s.update(engineId=r["policyEngineId"], engineArn=r["policyEngineArn"])
    save()
wait(c.get_policy_engine, "status", policyEngineId=s["engineId"])
# 2. Create the Lambda for shopping tools and an execution role for writing logs.
role("lambdaRole", "lambda.amazonaws.com")
log_arn = f"arn:aws:logs:{s['region']}:{s['account']}:log-group:/aws/lambda/{s['name']}"
iam.put_role_policy(
    RoleName=s["lambdaRole"].split("/")[-1],
    PolicyName="ArticleLogs",
    PolicyDocument=json.dumps(
        {
            "Version": "2012-10-17",
            "Statement": [
                {
                    "Effect": "Allow",
                    "Action": [
                        "logs:CreateLogGroup",
                        "logs:CreateLogStream",
                        "logs:PutLogEvents",
                    ],
                    "Resource": [log_arn, log_arn + ":*"],
                }
            ],
        }
    ),
)
if "lambdaArn" not in s:
    code = io.BytesIO()
    with zipfile.ZipFile(code, "w") as z:
        z.write(ROOT / "gateway/lambda_function.py", "lambda_function.py")
    for attempt in range(12):
        try:
            r = lam.create_function(
                FunctionName=s["name"],
                Runtime="python3.14",
                Role=s["lambdaRole"],
                Handler="lambda_function.lambda_handler",
                Code={"ZipFile": code.getvalue()},
                Timeout=10,
                MemorySize=128,
            )
            break
        except lam.exceptions.InvalidParameterValueException:
            if attempt == 11:
                raise
            time.sleep(5)
    s["lambdaArn"] = r["FunctionArn"]
    save()
lam.get_waiter("function_active_v2").wait(FunctionName=s["name"])
# 3. Grant the Gateway execution role permissions for Lambda invocation, Policy evaluation, and token retrieval.
role("gatewayRole", "bedrock-agentcore.amazonaws.com")
gw_resource = s.get("gatewayArn", base + ":gateway/" + s["name"] + "-*")
permissions = {
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Action": ["lambda:InvokeFunction"],
            "Resource": [s["lambdaArn"]],
        },
        {
            "Effect": "Allow",
            "Action": [
                "bedrock-agentcore:GetPolicyEngine",
                "bedrock-agentcore:AuthorizeAction",
                "bedrock-agentcore:PartiallyAuthorizeActions",
            ],
            "Resource": [s["engineArn"], gw_resource],
        },
        {
            "Effect": "Allow",
            "Action": ["bedrock-agentcore:GetWorkloadAccessToken"],
            "Resource": [
                base + ":workload-identity-directory/default",
                base
                + ":workload-identity-directory/default/workload-identity/"
                + s.get("gatewayId", s["name"])
                + "*",
            ],
        },
    ],
}
iam.put_role_policy(
    RoleName=s["gatewayRole"].split("/")[-1],
    PolicyName="ArticleGateway",
    PolicyDocument=json.dumps(permissions),
)
# 4. Create a Gateway with the Policy Engine linked in ENFORCE mode.
if "gatewayId" not in s:
    time.sleep(10)
    for attempt in range(12):
        try:
            r = c.create_gateway(
                name=s["name"],
                roleArn=s["gatewayRole"],
                protocolType="MCP",
                authorizerType="AWS_IAM",
                policyEngineConfiguration={"arn": s["engineArn"], "mode": "ENFORCE"},
                exceptionLevel="DEBUG",
            )
            break
        except c.exceptions.ValidationException as exc:
            if (
                "Access denied while calling GetPolicyEngine" not in str(exc)
                or attempt == 11
            ):
                raise
            time.sleep(5)  # Retry after waiting for IAM permissions to propagate.
    s.update(
        gatewayId=r["gatewayId"], gatewayArn=r["gatewayArn"], gatewayUrl=r["gatewayUrl"]
    )
    save()
wait(c.get_gateway, "status", gatewayIdentifier=s["gatewayId"])

def obj(props):
    return {
        "type": "object",
        "properties": {k: {"type": v} for k, v in props.items()},
        "required": list(props),
    }

# 5. Define the input/output for the 2 tools and register Lambda as a Gateway target.
cart = {"cartId": "string", "store": "string", "amount": "integer"}
tools = [
    {
        "name": "get_cart",
        "description": "Read the current shopping cart and total price.",
        "inputSchema": obj({"cartId": "string"}),
        "outputSchema": obj(cart),
    },
    {
        "name": "place_order",
        "description": "Simulate a supermarket order; no purchase is made.",
        "inputSchema": obj(cart),
        "outputSchema": obj({**cart, "status": "string"}),
    },
]
(ROOT / "gateway/tools.json").write_text(json.dumps(tools, indent=2))
if "targetId" not in s:
    r = c.create_gateway_target(
        gatewayIdentifier=s["gatewayId"],
        name="ShoppingTarget",
        targetConfiguration={
            "mcp": {
                "lambda": {
                    "lambdaArn": s["lambdaArn"],
                    "toolSchema": {"inlinePayload": tools},
                }
            }
        },
        credentialProviderConfigurations=[
            {"credentialProviderType": "GATEWAY_IAM_ROLE"}
        ],
    )
    s["targetId"] = r["targetId"]
    save()
wait(
    c.get_gateway_target,
    "status",
    gatewayIdentifier=s["gatewayId"],
    targetId=s["targetId"],
)
# 6. Register the cart confirmation and order policies. The meaning of the conditions will be explained in the next section of the article.
statements = {
    "ReadCart": f'permit (principal, action == AgentCore::Action::"ShoppingTarget___get_cart", resource == AgentCore::Gateway::"{s["gatewayArn"]}") when {{ ["cart-001", "cart-002", "cart-003"].contains(context.input.cartId) }};',
    "OrderAfterCartCheck": f"""permit (
    principal,
    action == AgentCore::Action::"ShoppingTarget___place_order",
    resource == AgentCore::Gateway::"{s['gatewayArn']}"
)
when {{ context.input.amount > 0 && context.input.amount <= 3000 }}
when temporal {{
    formerly within 1m AgentCore::Action::"ShoppingTarget___get_cart"::response{{
        eventResource: resource,
        output.cartId: context.input.cartId,
        output.store: context.input.store,
        output.amount: context.input.amount
    }}
}};""",
}
s.setdefault("policies", {})
save()
for name, statement in statements.items():
    (ROOT / "gateway" / (name + ".dw")).write_text(
        statement.replace(s["gatewayArn"], "<GATEWAY_ARN>") + "\n"
    )
    if name not in s["policies"]:
        r = c.create_policy(
            policyEngineId=s["engineId"],
            name=name,
            definition={"policy": {"statement": statement}},
            validationMode="FAIL_ON_ANY_FINDINGS",
        )
        s["policies"][name] = r["policyId"]
        save()
    existing = c.get_policy(policyEngineId=s["engineId"], policyId=s["policies"][name])
    if existing["status"] in ("CREATE_FAILED", "UPDATE_FAILED"):
        c.update_policy(
            policyEngineId=s["engineId"],
            policyId=s["policies"][name],
            definition={"policy": {"statement": statement}},
            validationMode="FAIL_ON_ANY_FINDINGS",
        )
    wait(
        c.get_policy,
        "status",
        policyEngineId=s["engineId"],
        policyId=s["policies"][name],
    )
print(
    json.dumps(
        {"gateway": s["gatewayId"], "target": s["targetId"], "policies": s["policies"]},
        indent=2,
    )
)

"""Run a real Strands/Bedrock agent against the shopping Gateway."""
import argparse
import json
import time
import uuid
from pathlib import Path

from mcp_proxy_for_aws.client import aws_iam_streamablehttp_client
from strands import Agent
from strands.models import BedrockModel
from strands.tools.mcp import MCPClient

ROOT = Path(__file__).resolve().parents[1]
MODEL_ID = 'jp.anthropic.claude-haiku-4-5-20251001-v1:0'
PROMPTS = {
    'normal': 'Check the cart for cart-001, and place an order with the same store and amount returned.',
    'skip': 'Without checking the cart, directly place an order for cart-001 at store favorite-supermarket for 3000 yen. If denied, stop there and do not confirm or retry.',
    'over_budget': 'Check the cart for cart-003, and place an order with the same store and amount returned. Even if over budget, call the order tool once to check the result, and stop if denied.',
}

def run(case):
    state = json.loads((ROOT / '.gateway-state.json').read_text())
    if state.get('cleanedUp'):
        raise RuntimeError('Please run gateway/setup.py first.')
    session_id = str(uuid.uuid4())
    mcp_client = MCPClient(lambda: aws_iam_streamablehttp_client(
        endpoint=state['gatewayUrl'],
        aws_service='bedrock-agentcore',
        aws_region=state['region'],
        headers={'x-amzn-bedrock-agentcore-policy-session-id': session_id},
    ))
    started = time.time()
    with mcp_client:
        agent = Agent(
            model=BedrockModel(model_id=MODEL_ID, region_name=state['region'], max_tokens=1500),
            tools=mcp_client.list_tools_sync(),
            system_prompt=(
                'You are an agent that helps with supermarket shopping.'
                'Please use tools in the order specified by the user.'
                'The Gateway determines whether orders are allowed. Do not retry denied operations; briefly report the result in Japanese.'
            ),
        )
        result = agent(PROMPTS[case])
    record = {
        'case': case, 'model': MODEL_ID, 'session': session_id,
        'resource_name': state['name'], 'started': started, 'ended': time.time(),
        'prompt': PROMPTS[case], 'messages': agent.messages, 'answer': str(result),
    }
    path = ROOT / 'verification' / f'strands-{case}.json'
    path.write_text(json.dumps(record, ensure_ascii=False, indent=2, default=str))
    print(f'\nSaved: {path}')
    return record

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('case', choices=[*PROMPTS, 'all'], default='normal', nargs='?')
    args = parser.parse_args()
    for case in PROMPTS if args.case == 'all' else [args.case]:
        run(case)

"""Add a cumulative amount limit for the last 5 minutes to the existing order policy."""
import json
import time
from pathlib import Path
import boto3

ROOT = Path(__file__).resolve().parents[1]
path = ROOT / '.gateway-state.json'
state = json.loads(path.read_text())
client = boto3.client('bedrock-agentcore-control', region_name=state['region'])
statement = (ROOT / 'gateway/SessionBudget.dw').read_text().replace(
    '<GATEWAY_ARN>', state['gatewayArn']
)
result = client.create_policy(
    policyEngineId=state['engineId'],
    name='SessionBudget',
    definition={'policy': {'statement': statement}},
    validationMode='FAIL_ON_ANY_FINDINGS',
)
policy_id = result['policyId']
state['policies']['SessionBudget'] = policy_id
path.write_text(json.dumps(state, indent=2))
for _ in range(120):
    result = client.get_policy(
        policyEngineId=state['engineId'], policyId=policy_id
    )
    if result['status'] == 'ACTIVE':
        print('SessionBudget ACTIVE')
        break
    if 'FAIL' in result['status']:
        raise RuntimeError(result)
    time.sleep(5)
else:
    raise TimeoutError('SessionBudget')

"""Try placing two 2,000 yen orders in the same session from Strands Agents."""
import strands_agent as sample

sample.PROMPTS['cumulative'] = (
    'Check cart-002 and place an order with the returned store and amount.'
    'If successful, place the same cart order again with the same store and amount.'
    'Even if the cumulative budget is exceeded, call the order tool a second time to check the Gateway result.'
    'If denied, stop without retrying.'
)
sample.run('cumulative')

uv run gateway/add_budget.py
uv run gateway/check_budget.py

"""Delete only resources recorded by this article's setup.py."""
import json,time
from pathlib import Path
import boto3
from botocore.exceptions import ClientError
root=Path(__file__).resolve().parents[1]
p=root/'.gateway-state.json'
s=json.loads(p.read_text())
aws=boto3.Session(region_name=s['region']); c=aws.client('bedrock-agentcore-control')
def remove(fn,**args):
    try: fn(**args)
    except ClientError as e:
        if e.response['Error']['Code'] not in ('ResourceNotFoundException','NoSuchEntityException','NoSuchEntity'): raise

def gone(get,**args):
    for _ in range(120):
        try: get(**args)
        except ClientError as e:
            if e.response['Error']['Code']=='ResourceNotFoundException': return
            raise
        time.sleep(3)
    raise TimeoutError(args)
if 'targetId' in s:
    remove(c.delete_gateway_target,gatewayIdentifier=s['gatewayId'],targetId=s['targetId'])
    gone(c.get_gateway_target,gatewayIdentifier=s['gatewayId'],targetId=s['targetId'])
if 'gatewayId' in s:
    remove(c.delete_gateway,gatewayIdentifier=s['gatewayId'])
    gone(c.get_gateway,gatewayIdentifier=s['gatewayId'])
for pid in s.get('policies',{}).values():
    remove(c.delete_policy,policyEngineId=s['engineId'],policyId=pid)
    gone(c.get_policy,policyEngineId=s['engineId'],policyId=pid)
if 'engineId' in s:
    remove(c.delete_policy_engine,policyEngineId=s['engineId'])
    gone(c.get_policy_engine,policyEngineId=s['engineId'])
if 'lambdaArn' in s: remove(aws.client('lambda').delete_function,FunctionName=s['name'])
remove(aws.client('logs').delete_log_group,logGroupName='/aws/lambda/'+s['name'])
iam=aws.client('iam')
for key,policy in [('lambdaRole','ArticleLogs'),('gatewayRole','ArticleGateway')]:
    if key in s:
        name=s[key].split('/')[-1]
        remove(iam.delete_role_policy,RoleName=name,PolicyName=policy)
        remove(iam.delete_role,RoleName=name)
(root/'verification/cleanup.json').write_text(json.dumps({'status':'deleted','resource_name':s['name'],'region':s['region'],'completed_at':time.time()},indent=2))
s['cleanedUp']=True; p.write_text(json.dumps(s,indent=2))
print('Deleted article Gateway, target, policies, policy engine, Lambda, log group and 2 IAM roles.')

forbid (
    principal,
    action == AgentCore::Action::"ShoppingTarget___place_order",
    resource == AgentCore::Gateway::"<GATEWAY_ARN>"
)
when temporal {
    exists (n: Long).
        (count for (t: Timepoint).
            where (formerly within 5m (
                AgentCore::Action::"ShoppingTarget___place_order"::request{
                    eventResource: resource
                } && tp(t)
            ))) == n
        && n > 3
};

forbid (
    principal,
    action == AgentCore::Action::"ShoppingTarget___place_order",
    resource == AgentCore::Gateway::"<GATEWAY_ARN>"
)
when temporal {
    formerly within 1m AgentCore::Action::"ShoppingTarget___place_order"::response{
        eventResource: resource
    }
};

permit (
    principal,
    action == AgentCore::Action::"ShoppingTarget___place_order",
    resource == AgentCore::Gateway::"<GATEWAY_ARN>"
)
when { context.input.amount > 0 && context.input.amount <= 3000 }
when temporal {
    formerly within 1m AgentCore::Action::"ShoppingTarget___get_cart"::response{
        eventResource: resource,
        output.cartId: context.input.cartId,
        output.store: context.input.store,
        output.amount: context.input.amount
    }
    && (
        !AgentCore::Action::"ShoppingTarget___place_order"::response{
            eventResource: resource
        }
        since within 1m AgentCore::Action::"ShoppingTarget___get_cart"::response{
            eventResource: resource
        }
    )
};