Amazon Bedrock AgentCore Policyに追加された「Temporal Policies」を試してみた

"""記事の検証環境を作成し、作成したリソースのIDをファイルへ保存する。"""

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

# 作成先は東京リージョン。途中経過を保存し、再実行時に作成済みの情報を読み込む。
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. 買い物ルールを登録するPolicy Engineを作成する。
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. 買い物ツールのLambdaと、ログを書き込むための実行ロールを作成する。
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. Gatewayの実行ロールに、Lambda呼び出し・Policy評価・トークン取得の権限を付ける。
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. Policy EngineをENFORCEモードで結び付けたGatewayを作成する。
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)  # IAM権限の反映を待って再試行する。
    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. 2つのツールの入出力を定義し、LambdaをGatewayのターゲットとして登録する。
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. カート確認と注文のポリシーを登録する。条件の意味は記事の次の節で説明する。
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': 'cart-001のカートを確認して、返ってきた店舗と金額のまま注文してください。',
    'skip': 'カート確認はしないで、cart-001を店舗favorite-supermarket、金額3000円で直接注文してください。拒否されたらそこで終了し、確認や再試行はしないでください。',
    'over_budget': 'cart-003のカートを確認して、返ってきた店舗と金額のまま注文してください。予算を超えていても注文ツールを1回呼び出して結果を確認し、拒否されたら終了してください。',
}

def run(case):
    state = json.loads((ROOT / '.gateway-state.json').read_text())
    if state.get('cleanedUp'):
        raise RuntimeError('先に gateway/setup.py を実行してください。')
    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=(
                'あなたはスーパーの買い物を手伝うエージェントです。'
                'ユーザーの指定した手順でツールを使ってください。'
                '注文可否はGatewayが判定します。拒否された操作は再試行せず、結果を日本語で短く伝えてください。'
            ),
        )
        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)

"""既存の注文ポリシーに、直近5分間の累計金額の制限を追加する。"""
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')

"""Strands Agentsから同じセッションで2,000円の注文を2回試す。"""
import strands_agent as sample

sample.PROMPTS['cumulative'] = (
    'cart-002を確認し、返された店舗と金額のまま注文してください。'
    '成功したら同じカートを同じ店舗・金額でもう1回注文してください。'
    '累計予算を超えても2回目の注文ツールを1回呼び出してGatewayの結果を確認してください。'
    '拒否されたら再試行せず終了してください。'
)
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
        }
    )
};