Bedrock AgentCore Runtimeのマネージドセッションストレージでエージェントの作業状態を停止・再開をまたいで永続化する

import json
import sys
 
def handle_invoke(event):
    user_input = event.get("input", {}).get("text", "")
    return {
        "output": {
            "text": f"Received: {user_input}. This is a minimal test agent."
        }
    }
 
def main():
    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue
        try:
            request = json.loads(line)
            response = handle_invoke(request)
            print(json.dumps(response), flush=True)
        except json.JSONDecodeError:
            print(json.dumps({"error": "Invalid JSON"}), flush=True)
 
if __name__ == "__main__":
    main()

ACCOUNT_ID=$(aws sts get-caller-identity --query Account --output text)
BUCKET_NAME="agentcore-session-storage-test-${ACCOUNT_ID}"
REGION="us-west-2"
 
# コードをZIP化してS3にアップロード
zip agent.zip main.py
aws s3 mb "s3://${BUCKET_NAME}" --region "$REGION"
aws s3 cp agent.zip "s3://${BUCKET_NAME}/agent.zip"
 
# IAMロール作成
aws iam create-role \
  --role-name AgentCoreSessionStorageTestRole \
  --assume-role-policy-document '{
    "Version": "2012-10-17",
    "Statement": [{
      "Effect": "Allow",
      "Principal": {"Service": "bedrock-agentcore.amazonaws.com"},
      "Action": "sts:AssumeRole"
    }]
  }'
 
aws iam attach-role-policy \
  --role-name AgentCoreSessionStorageTestRole \
  --policy-arn arn:aws:iam::aws:policy/AmazonBedrockFullAccess
 
aws iam put-role-policy \
  --role-name AgentCoreSessionStorageTestRole \
  --policy-name S3Access \
  --policy-document "{
    \"Version\": \"2012-10-17\",
    \"Statement\": [{
      \"Effect\": \"Allow\",
      \"Action\": [\"s3:GetObject\", \"s3:ListBucket\"],
      \"Resource\": [
        \"arn:aws:s3:::${BUCKET_NAME}\",
        \"arn:aws:s3:::${BUCKET_NAME}/*\"
      ]
    }]
  }"

aws bedrock-agentcore-control create-agent-runtime \
  --region "$REGION" \
  --agent-runtime-name session_storage_test_agent \
  --role-arn "arn:aws:iam::${ACCOUNT_ID}:role/AgentCoreSessionStorageTestRole" \
  --agent-runtime-artifact "{
    \"codeConfiguration\": {
      \"code\": {\"s3\": {\"bucket\": \"${BUCKET_NAME}\", \"prefix\": \"agent.zip\"}},
      \"runtime\": \"PYTHON_3_13\",
      \"entryPoint\": [\"main.py\"]
    }
  }" \
  --network-configuration '{"networkMode": "PUBLIC"}' \
  --filesystem-configurations '[{
    "sessionStorage": {
      "mountPath": "/mnt/workspace"
    }
  }]'
# → レスポンスの agentRuntimeId を控える

RUNTIME_ID="session_storage_test_agent-XXXXXXXXXX"  # 実際のIDに置き換え
 
aws bedrock-agentcore-control create-agent-runtime-endpoint \
  --region us-west-2 \
  --agent-runtime-id "$RUNTIME_ID" \
  --name session_storage_test_endpoint
 
# READYまでポーリング
while true; do
  STATUS=$(aws bedrock-agentcore-control get-agent-runtime-endpoint \
    --region us-west-2 \
    --agent-runtime-id "$RUNTIME_ID" \
    --endpoint-name session_storage_test_endpoint \
    --query 'status' --output text)
  echo "Endpoint status: $STATUS"
  [ "$STATUS" = "READY" ] && break
  sleep 10
done

import boto3, uuid, json, time
 
REGION = "us-west-2"
RUNTIME_ARN = "arn:aws:bedrock-agentcore:us-west-2:111122223333:runtime/RUNTIME_ID"
 
client = boto3.client("bedrock-agentcore", region_name=REGION)
 
def make_session_id():
    """セッションIDは33文字以上が必要なため、UUIDを結合して生成する。"""
    return str(uuid.uuid4()) + "-" + str(uuid.uuid4())[:8]
 
def run_command(command, timeout=60, session_id=None):
    """コマンドを /bin/bash -c で実行し、EventStreamを処理して結果を返す。"""
    response = client.invoke_agent_runtime_command(
        agentRuntimeArn=RUNTIME_ARN,
        runtimeSessionId=session_id,
        qualifier="DEFAULT",
        contentType="application/json",
        accept="application/vnd.amazon.eventstream",
        body={"command": f'/bin/bash -c {json.dumps(command)}', "timeout": timeout},
    )
    stdout, stderr, exit_code, status = [], [], None, None
    for event in response.get("stream", []):
        if "chunk" in event:
            chunk = event["chunk"]
            if "contentDelta" in chunk:
                d = chunk["contentDelta"]
                if d.get("stdout"): stdout.append(d["stdout"])
                if d.get("stderr"): stderr.append(d["stderr"])
            if "contentStop" in chunk:
                exit_code = chunk["contentStop"].get("exitCode")
                status = chunk["contentStop"].get("status")
    return {"stdout": "".join(stdout), "stderr": "".join(stderr),
            "exit_code": exit_code, "status": status}
 
def stop_session(session_id):
    """データプレーンAPIでセッションを停止する。"""
    return client.stop_runtime_session(
        agentRuntimeArn=RUNTIME_ARN, runtimeSessionId=session_id
    )

from test_helper import *
 
session = make_session_id()
 
# ファイル・ディレクトリ・シンボリックリンク作成
run_command(
    'mkdir -p /mnt/workspace/project/src && '
    'echo "hello from session storage" > /mnt/workspace/project/README.md && '
    'echo "print(\'hello\')" > /mnt/workspace/project/src/main.py && '
    'chmod 755 /mnt/workspace/project/src/main.py && '
    'ln -s /mnt/workspace/project/README.md /mnt/workspace/project/link-to-readme',
    session_id=session
)
 
# 停止 → 15秒待機 → 再開
stop_session(session)
time.sleep(15)
 
# 復元確認
r = run_command(
    'find /mnt/workspace -type f -o -type l | sort && '
    'echo "---README---" && cat /mnt/workspace/project/README.md && '
    'echo "---SYMLINK---" && readlink /mnt/workspace/project/link-to-readme && '
    'echo "---PERMS---" && stat -c "%a %n" /mnt/workspace/project/src/main.py',
    session_id=session
)
print(r["stdout"])

from test_helper import *
 
session = make_session_id()
 
# pip パッケージをワークスペース内にインストール
run_command(
    'python3 -m pip install --target=/mnt/workspace/pylibs requests',
    session_id=session, timeout=120
)
 
# git リポジトリ作成
run_command(
    'cd /mnt/workspace && git init myproject && cd myproject && '
    'echo "# Session Storage Test" > README.md && '
    'mkdir src && echo "print(\'hello\')" > src/app.py && '
    'git add -A && '
    'git -c user.email="test@example.com" -c user.name="Test" commit -m "initial commit"',
    session_id=session
)
 
# 停止 → 20秒待機 → 再開
stop_session(session)
time.sleep(20)
 
# 復元確認
r = run_command(
    'echo "=== requests import ===" && '
    'PYTHONPATH=/mnt/workspace/pylibs python3 -c "import requests; print(f\'requests {requests.__version__} - OK\')" && '
    'echo "=== git log ===" && cd /mnt/workspace/myproject && git log --oneline && '
    'echo "=== git status ===" && git status && '
    'echo "=== file content ===" && cat src/app.py && '
    'echo "=== workspace du ===" && du -sh /mnt/workspace/*',
    session_id=session
)
print(r["stdout"])

from test_helper import *
 
session_x = make_session_id()
session_y = make_session_id()
 
# Session X にファイル作成
run_command(
    'echo "secret from session X" > /mnt/workspace/secret.txt',
    session_id=session_x
)
 
# Session Y から参照を試みる
r = run_command(
    'ls -la /mnt/workspace/ && '
    'cat /mnt/workspace/secret.txt 2>&1 || echo "FILE NOT FOUND"',
    session_id=session_y
)
print(r["stdout"])

from test_helper import *
 
# --- 小さいファイル（100行テキスト） ---
session_small = make_session_id()
 
run_command(
    'for i in $(seq 1 100); do echo "line $i" >> /mnt/workspace/bigfile.txt; done && '
    'wc -l /mnt/workspace/bigfile.txt',
    session_id=session_small
)
 
stop_session(session_small)
time.sleep(2)
 
r = run_command('wc -l /mnt/workspace/bigfile.txt', session_id=session_small)
print(r["stdout"])
 
# --- 大きいファイル（50MB） ---
session_large = make_session_id()
 
run_command(
    'dd if=/dev/urandom of=/mnt/workspace/random.bin bs=1M count=50 2>&1 && '
    'md5sum /mnt/workspace/random.bin',
    session_id=session_large, timeout=120
)
 
# 即停止 → 待ち時間なしで再開
stop_session(session_large)
 
r = run_command(
    'md5sum /mnt/workspace/random.bin && du -sh /mnt/workspace/random.bin',
    session_id=session_large
)
print(r["stdout"])  # md5sum が一致するか確認

from test_helper import *
 
session = make_session_id()
 
# ファイル作成 → 停止
run_command(
    'echo "version 1 data" > /mnt/workspace/version-test.txt',
    session_id=session
)
stop_session(session)
time.sleep(10)
 
# ランタイムを更新（同じコードだが新バージョンになる）
control = boto3.client("bedrock-agentcore-control", region_name=REGION)
control.update_agent_runtime(
    agentRuntimeId="RUNTIME_ID",
    roleArn="arn:aws:iam::111122223333:role/AgentCoreSessionStorageTestRole",
    agentRuntimeArtifact={...},  # 作成時と同じ設定
    networkConfiguration={"networkMode": "PUBLIC"},
    filesystemConfigurations=[{"sessionStorage": {"mountPath": "/mnt/workspace"}}],
)
 
# READY を待ってから同じセッションIDで再開
time.sleep(15)
r = run_command(
    'cat /mnt/workspace/version-test.txt 2>&1 || echo "FILE NOT FOUND (storage was reset)"',
    session_id=session
)
print(r["stdout"])