Useful A2A extensions

from urllib.parse import urljoin
from uuid import uuid4

import httpx

from agentenv_protocol.a2a_agent import PEER_AGENTS_V1, PeerAgentsSetRequest, extension


class ClaudiusRLMAgent(AgentEnvAgent):
    def __init__(self) -> None:
        self.peers = {}

    @extension(PEER_AGENTS_V1.set)
    async def set_peers(self, request: PeerAgentsSetRequest):
        self.peers = {peer.name: peer for peer in request.peers}
        return {"status": "updated"}

    @extension(PEER_AGENTS_V1.list)
    async def list_peers(self):
        return {"peers": [peer.model_dump() for peer in self.peers.values()]}

    async def ask_peer(self, name: str, text: str) -> str:
        """Send a peer one A2A message and return its reply."""
        peer = self.peers[name]
        message = {"kind": "message", "messageId": uuid4().hex, "role": "user", "parts": [{"kind": "text", "text": text}]}
        async with httpx.AsyncClient(timeout=600) as client:
            response = await client.post(
                urljoin(peer.url, peer.card.get("url", "/a2a")),
                json={"jsonrpc": "2.0", "id": 1, "method": "message/send", "params": {"message": message}},
            )
        task = response.json()["result"]
        return "".join(part["text"] for part in task["status"]["message"]["parts"] if part["kind"] == "text")

import io
import json
import tarfile
import tempfile
from pathlib import Path
from uuid import uuid4

from agentenv_protocol.a2a_agent import (
    SNAPSHOT_V1,
    ObjectSnapshotLoadRequest,
    ObjectSnapshotSaveRequest,
    ObjectSnapshotSaveResponse,
    SnapshotUploadedObjects,
    download,
    extension,
    upload,
)

WORKSPACE = Path("/app/workspace")


def archive(directory: Path) -> bytes:
    buffer = io.BytesIO()
    with tarfile.open(fileobj=buffer, mode="w") as tar:
        tar.add(directory, arcname=".")
    return buffer.getvalue()


class ClaudiusRLMAgent(AgentEnvAgent):
    def __init__(self) -> None:
        self.conversations = {}
        WORKSPACE.mkdir(parents=True, exist_ok=True)

    @extension(SNAPSHOT_V1.save)
    async def save_snapshot(self, request: ObjectSnapshotSaveRequest):
        conversation = json.dumps(self.conversations.get(request.context_id, [])).encode()
        objects = {"trajectory": await upload(request.objects.trajectory, conversation)}
        if request.objects.workspace is not None:
            objects["workspace"] = await upload(request.objects.workspace, archive(WORKSPACE))
        return ObjectSnapshotSaveResponse(
            context_id=request.context_id, objects=SnapshotUploadedObjects(**objects)
        )

    @extension(SNAPSHOT_V1.load)
    async def load_snapshot(self, request: ObjectSnapshotLoadRequest):
        with tempfile.TemporaryDirectory() as tmp:
            await download(request.objects.trajectory, Path(tmp, "trajectory"))
            conversation = json.loads(Path(tmp, "trajectory").read_bytes())
            if request.objects.workspace is not None:
                await download(request.objects.workspace, Path(tmp, "workspace"))
                with tarfile.open(Path(tmp, "workspace")) as tar:
                    tar.extractall(WORKSPACE, filter="data")
        context_id = request.target_context_id or uuid4().hex
        self.conversations[context_id] = conversation
        return {"context_id": context_id}

{"id": "rules", "type": "register_agent_triggers", "agent_name": "user-sim",
 "triggers": [
   {"id": "confirm", "when": {"type": "conversational", "where": {"message": {"regex": "Send it\\?"}}},
    "actions": [{"type": "say", "text": "Yes, send it."}]},
   {"id": "sent", "when": {"type": "conversational", "where": {"message": {"regex": "Sent"}}},
    "actions": [{"type": "say", "text": "Thanks."}, {"type": "end"}]}]}

from agentenv_protocol.a2a_agent import INSTALL_V1, enable

INSTALL = enable(
    INSTALL_V1,
    install_commands=[
        "docker exec {container} mkdir -p /opt/claudius",
        "docker cp {agent_ctx_tar} {container}:/opt/claudius/claudius.tar.gz",
        "docker exec -w /opt/claudius {container} sh -c 'tar -xzf claudius.tar.gz && python3 -m venv .venv"
        " && .venv/bin/pip install \"agentenv-framework-protocol[agent]\" \"mcp>=1.25,<2\" anthropic'",
        "docker exec -d -w /opt/claudius -e A2A_PORT={a2a_port} \\\n"
        "  -e LITELLM_BASE_URL={litellm_base_url} -e LITELLM_API_KEY={litellm_api_key} {container} .venv/bin/python agent.py",
    ],
    required_params=["container", "agent_ctx_tar", "a2a_port", "litellm_base_url", "litellm_api_key"],
    a2a_port=8000,
)


@a2a_agent(identity=..., config=ClaudiusRLMConfig, extensions=(MCP_CONFIG_V1, TRAJECTORY_V1, INSTALL))
class ClaudiusRLMAgent(AgentEnvAgent):
    ...