Creating your own agent

import os
from contextlib import AsyncExitStack

from anthropic import AsyncAnthropic
from mcp import ClientSession
from mcp.client.streamable_http import streamablehttp_client

from agentenv_protocol.a2a_agent import (
    MCP_CONFIG_V1,
    TRAJECTORY_V1,
    AgentConfig,
    AgentEnvAgent,
    AgentIdentity,
    TaskRequest,
    TaskResult,
    TextPart,
    Usage,
    a2a_agent,
)

RECURSE = {
    "name": "recurse",
    "description": (
        "Hand a self-contained sub-task to a fresh call of yourself, with its own context and the same "
        "tools. It returns only its answer, so use it for work whose details you do not need to keep."
    ),
    "input_schema": {
        "type": "object",
        "properties": {"task": {"type": "string", "description": "The sub-task, with everything it needs."}},
        "required": ["task"],
    },
}


class ClaudiusRLMConfig(AgentConfig):
    model: str = "claude-opus-5-5"
    system_prompt: str = "You work an email inbox. Use the tools to do what you are asked."
    effort: str = "high"
    max_turns: int = 10
    max_depth: int = 2


class Stopped(Exception):
    """A call that ended without an answer."""


@a2a_agent(
    identity=AgentIdentity(
        name="claudius",
        description="Claude that works an inbox through the tools of the envs it is given, and calls itself on sub-tasks.",
        version="1.0.0",
    ),
    config=ClaudiusRLMConfig,
    extensions=(MCP_CONFIG_V1, TRAJECTORY_V1),
)
class ClaudiusRLMAgent(AgentEnvAgent):
    async def run(self, request: TaskRequest[ClaudiusRLMConfig]) -> TaskResult:
        config = request.config
        claude = AsyncAnthropic(
            base_url=os.environ["LITELLM_BASE_URL"].removesuffix("/v1"), api_key=os.environ["LITELLM_API_KEY"]
        )
        calls = []

        async def solve(task: str, depth: int) -> str:
            tools = env_tools + ([RECURSE] if depth < config.max_depth else [])
            messages = [{"role": "user", "content": task}]
            for _ in range(config.max_turns):
                response = await claude.messages.create(
                    model=config.model,
                    max_tokens=16000,
                    system=config.system_prompt,
                    tools=tools,
                    messages=messages,
                    output_config={"effort": config.effort},
                )
                if response.stop_reason == "refusal":
                    raise Stopped("refusal", "Claude declined the task")
                if response.stop_reason != "tool_use":
                    return "".join(block.text for block in response.content if block.type == "text")
                messages.append({"role": "assistant", "content": response.content})
                results = []
                for block in response.content:
                    if block.type != "tool_use":
                        continue
                    if block.name == RECURSE["name"]:
                        output = await solve(block.input["task"], depth + 1)
                    else:
                        result = await sessions[block.name].call_tool(block.name, block.input)
                        output = "\n".join(c.text for c in result.content if c.type == "text")
                    calls.append({"depth": depth, "tool": block.name, "input": block.input, "output": output})
                    results.append({"type": "tool_result", "tool_use_id": block.id, "content": output})
                messages.append({"role": "user", "content": results})
            raise Stopped("max_turns", f"Still calling tools after {config.max_turns} turns at depth {depth}")

        async with AsyncExitStack() as stack:
            # One MCP session per env the framework handed over, and every tool they offer.
            sessions, env_tools = {}, []
            for server in request.mcp_servers.values():
                read, write, _ = await stack.enter_async_context(
                    streamablehttp_client(server["url"], headers=server.get("headers"))
                )
                session = await stack.enter_async_context(ClientSession(read, write))
                await session.initialize()
                for tool in (await session.list_tools()).tools:
                    sessions[tool.name] = session
                    env_tools.append(
                        {"name": tool.name, "description": tool.description or "", "input_schema": tool.inputSchema}
                    )

            prompt = "\n".join(part.text for part in request.parts if isinstance(part, TextPart))
            try:
                answer = await solve(prompt, depth=0)
            except Stopped as stop:
                return TaskResult.failure(*stop.args)

        return (
            TaskResult.builder()
            .succeeded()
            .add_text(answer)
            .usage(Usage(tool_call_count=len(calls)))
            .native_trajectory(format="claudius/v1", payload=calls)
            .build()
        )


if __name__ == "__main__":
    ClaudiusRLMAgent().serve()