More Useful Env Extensions

import random

from agentenv_protocol import AgentEnvEnvironment, environment_card, extension, tool

ERRORS = {"rate_limit": "rate limit exceeded, retry later", "timeout": "request timed out", "unavailable": "service unavailable"}


@environment_card(name="email")
class EmailEnv(AgentEnvEnvironment):
    def __init__(self) -> None:
        self.inbox: list[dict] = []
        self.sent: list[dict] = []
        self.faults: dict[str, dict] = {}

    @extension(uri="urn:agentenv:set-errors/v1", description="Make a tool fail.")
    async def set_errors(self, tool_name: str, every_nth: int = 0, error_rate: float = 0.0, error_type: str = "rate_limit") -> dict:
        self.faults[tool_name] = {"every_nth": every_nth, "error_rate": error_rate, "error_type": error_type, "calls": 0}
        return {"tool_name": tool_name, "every_nth": every_nth, "error_rate": error_rate, "error_type": error_type}

    def _fault(self, tool_name: str) -> None:
        fault = self.faults.get(tool_name)
        if fault is None:
            return
        fault["calls"] += 1
        nth = fault["every_nth"]
        failing = fault["calls"] % nth == 0 if nth else random.random() < fault["error_rate"]
        if failing:
            raise RuntimeError(ERRORS[fault["error_type"]])

    @tool()
    async def send(self, to: str, subject: str, body: str) -> dict:
        """Send an email."""
        self._fault("send")
        self.sent.append({"to": to, "subject": subject, "body": body})
        return {"sent": len(self.sent)}

USERS = {"alex@example.com", "priya@example.com", "jordan@example.com"}


@environment_card(name="email")
class EmailEnv(AgentEnvEnvironment):
    def __init__(self) -> None:
        ...
        self.user_email = "alex@example.com"

    @extension(uri="urn:agentenv:set-acting-user/v1", description="Act as this user.")
    async def set_acting_user(self, user_email: str) -> dict:
        if user_email not in USERS:
            raise ValueError(f"unknown user: {user_email}")
        self.user_email = user_email
        return {"user_email": user_email}

    @tool()
    async def email_search(self, query: str) -> dict:
        """Find the acting user's emails whose subject or body contains the query."""
        q = query.lower()
        mine = [e for e in self.inbox if e["to"] == self.user_email]
        return {"emails": [e for e in mine if q in (e["subject"] + " " + e["body"]).lower()]}

    @tool()
    async def send(self, to: str, subject: str, body: str) -> dict:
        """Send an email as the acting user."""
        self._fault("send")
        self.sent.append({"from": self.user_email, "to": to, "subject": subject, "body": body})
        return {"sent": len(self.sent)}

import random
import time

from agentenv_protocol import AgentEnvEnvironment, environment_card, extension, tool


@environment_card(name="reports")
class ReportsEnv(AgentEnvEnvironment):
    def __init__(self) -> None:
        self.waits: dict[str, tuple[float, float]] = {}
        self.jobs: dict[str, dict] = {}

    @extension(uri="urn:agentenv:set-async-wait/v1", description="Make a tool's jobs take this long.")
    async def set_async_wait(self, tool_name: str, min_seconds: float, max_seconds: float | None = None) -> dict:
        self.waits[tool_name] = (min_seconds, max_seconds if max_seconds is not None else min_seconds)
        return {"tool_name": tool_name, "wait_range_seconds": list(self.waits[tool_name])}

    @tool()
    async def reports_submit(self, query: str) -> dict:
        """Start a report; poll reports_status for it."""
        wait = random.uniform(*self.waits.get("reports_submit", (0, 0)))
        job_id = f"job_{len(self.jobs) + 1}"
        self.jobs[job_id] = {"query": query, "ready_at": time.monotonic() + wait}
        return {"job_id": job_id, "status": "pending"}

    @tool()
    async def reports_status(self, job_id: str) -> dict:
        """Whether a report is ready, and its rows once it is."""
        job = self.jobs[job_id]
        if time.monotonic() < job["ready_at"]:
            return {"job_id": job_id, "status": "pending"}
        return {"job_id": job_id, "status": "ready", "rows": 12}