Claude Managed Agents

import asyncio, os, re
import anthropic
from boxd import AsyncBoxd, NotFoundError
from fastapi import FastAPI, Request, Response

ENVIRONMENT_ID = os.environ["ANTHROPIC_ENVIRONMENT_ID"]
ENVIRONMENT_KEY = os.environ["ANTHROPIC_ENVIRONMENT_KEY"]

client = anthropic.AsyncAnthropic(auth_token=ENVIRONMENT_KEY)
boxd = AsyncBoxd()
app = FastAPI()
inflight: set[asyncio.Task] = set()

def machine_name(session_id: str) -> str:
    return "agent-" + re.sub(r"[^a-z0-9-]", "-", session_id.lower())

async def ensure_machine(session_id: str) -> str:
    name = machine_name(session_id)
    try:
        machine = await boxd.machines.get(name)
    except NotFoundError:
        machine = await boxd.machines.create(name, from_snapshot="claude-worker-base", isolated=True)
        await boxd.machines.set_auto_hibernate_timeout(machine.id, 0)
    if machine.status == "suspended":      # the SDK's name for standby
        await boxd.machines.resume(machine.id)
    elif machine.status == "hibernated":
        await boxd.machines.wake(machine.id)
    await boxd.machines.wait_until_ready(machine.id)
    return machine.id

async def run_session(work) -> None:
    machine_id = await ensure_machine(work.data.id)
    env = {
        "ANTHROPIC_SESSION_ID": work.data.id,
        "ANTHROPIC_WORK_ID": work.id,
        "ANTHROPIC_ENVIRONMENT_ID": ENVIRONMENT_ID,
        "ANTHROPIC_ENVIRONMENT_KEY": ENVIRONMENT_KEY,
    }
    # Returns when the worker exits, a minute after the session goes idle.
    await boxd.machines.exec(
        machine_id, "ant beta:worker run --workdir /workspace",
        env=env, timeout=24 * 60 * 60,
    )
    await boxd.machines.pause(machine_id)

async def remove_machine(session_id: str) -> None:
    try:
        await boxd.machines.delete(machine_name(session_id))
    except NotFoundError:
        pass

async def drain() -> None:
    async for work in client.beta.environments.work.poller(
        environment_id=ENVIRONMENT_ID, environment_key=ENVIRONMENT_KEY,
        block_ms=None, reclaim_older_than_ms=2000, drain=True, auto_stop=False,
    ):
        if work.data.type == "session":
            task = asyncio.create_task(run_session(work))
            inflight.add(task)
            task.add_done_callback(inflight.discard)

@app.post("/webhook")
async def webhook(request: Request) -> Response:
    try:
        event = client.beta.webhooks.unwrap(
            (await request.body()).decode(), headers=dict(request.headers)
        )
    except Exception:
        return Response("invalid signature", status_code=401)
    if event.data.type == "session.status_run_started":
        task = asyncio.create_task(drain())   # ack now, the session outlives the delivery
    elif event.data.type in ("session.status_terminated", "session.deleted"):
        task = asyncio.create_task(remove_machine(event.data.id))
    else:
        return Response(status_code=204)
    inflight.add(task)
    task.add_done_callback(inflight.discard)
    return Response(status_code=204)