Persist agent workspace across stop/resume with Bedrock AgentCore Runtime managed session storage

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 agent.zip main.py
aws s3 mb "s3://${BUCKET_NAME}" --region "$REGION"
aws s3 cp agent.zip "s3://${BUCKET_NAME}/agent.zip"
 
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"
    }
  }]'
# → Note the agentRuntimeId from the response

RUNTIME_ID="session_storage_test_agent-XXXXXXXXXX"
 
aws bedrock-agentcore-control create-agent-runtime-endpoint \
  --region us-west-2 \
  --agent-runtime-id "$RUNTIME_ID" \
  --name session_storage_test_endpoint
 
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():
    """Session IDs must be 33+ characters — concatenate UUIDs."""
    return str(uuid.uuid4()) + "-" + str(uuid.uuid4())[:8]
 
def run_command(command, timeout=60, session_id=None):
    """Execute a command via /bin/bash -c and process the 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):
    """Stop a session via the data plane API."""
    return client.stop_runtime_session(
        agentRuntimeArn=RUNTIME_ARN, runtimeSessionId=session_id
    )

from test_helper import *
 
session = make_session_id()
 
# Create files, directories, symlinks
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
)
 
# Stop → wait 15s → resume
stop_session(session)
time.sleep(15)
 
# Verify restoration
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()
 
# Install pip packages to workspace
run_command(
    'python3 -m pip install --target=/mnt/workspace/pylibs requests',
    session_id=session, timeout=120
)
 
# Create git repository
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
)
 
# Stop → wait 20s → resume
stop_session(session)
time.sleep(20)
 
# Verify restoration
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 creates a file
run_command(
    'echo "secret from session X" > /mnt/workspace/secret.txt',
    session_id=session_x
)
 
# Session Y tries to read it
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 *
 
# --- Small file (100-line text) ---
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"])
 
# --- Large file (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 immediately → resume with no wait
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"])  # Check if md5sum matches

from test_helper import *
 
session = make_session_id()
 
# Create file → stop
run_command(
    'echo "version 1 data" > /mnt/workspace/version-test.txt',
    session_id=session
)
stop_session(session)
time.sleep(10)
 
# Update runtime (same code, but creates a new version)
control = boto3.client("bedrock-agentcore-control", region_name=REGION)
control.update_agent_runtime(
    agentRuntimeId="RUNTIME_ID",
    roleArn="arn:aws:iam::111122223333:role/AgentCoreSessionStorageTestRole",
    agentRuntimeArtifact={...},  # Same config as creation
    networkConfiguration={"networkMode": "PUBLIC"},
    filesystemConfigurations=[{"sessionStorage": {"mountPath": "/mnt/workspace"}}],
)
 
# Wait for READY, then resume with same session 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"])