agentenv-framework 0.9.1254__py3-none-any.whl

agent_env/a2a_agent/protocol.py

@@ -0,0 +1,238 @@
+"""Helper Functions for interacting with A2A agents (like sending a task to an A2A Agent)."""
+from __future__ import annotations
+
+import asyncio
+import logging
+import re
+import time
+from collections.abc import Mapping
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import httpx
+from a2a.types import TaskState
+
+logger = logging.getLogger(__name__)
+
+_TERMINAL_TASK_STATES = frozenset({
+    TaskState.completed, TaskState.failed, TaskState.canceled, TaskState.rejected,
+})
+
+
+_ERROR_CODE = re.compile(r"[a-z][a-z0-9_]{0,63}")
+
+
+def raise_for_extension_status(
+    response: httpx.Response, *, operation: str, include_body: bool = False
+) -> None:
+    """Raise ``httpx.HTTPStatusError`` naming the SDK error code, but no other part of a body
+    that may echo a grant. ``include_body`` keeps the body, for a request that sent no grant."""
+    try:
+        response.raise_for_status()
+    except httpx.HTTPStatusError as exc:
+        if include_body:
+            detail = response.text[:500]
+        else:
+            try:
+                body = response.json()
+            except ValueError:
+                body = None
+            error = body.get("error") if isinstance(body, Mapping) else None
+            code = error.get("code") if isinstance(error, Mapping) else None
+            detail = code if isinstance(code, str) and _ERROR_CODE.fullmatch(code) else ""
+        message = f"{operation} failed with HTTP {response.status_code}"
+        if detail:
+            message += f": {detail}"
+        raise httpx.HTTPStatusError(message, request=exc.request, response=response) from exc
+
+
+async def post_agent_config(url: str, payload: dict, timeout_seconds: int = 60) -> None:
+    """POST to an agent's /ext/agent-config endpoint with retry on transient 404/connect errors.
+
+    `url` should be the full endpoint URL (a2a_url + endpoint path from the agent card).
+    """
+    max_attempts = 10
+    delay = 2.0
+    last_exc: Exception | None = None
+    for attempt in range(1, max_attempts + 1):
+        try:
+            async with httpx.AsyncClient() as client:
+                resp = await client.post(url, json=payload, timeout=timeout_seconds)
+            if resp.status_code == 404:
+                raise httpx.HTTPStatusError(
+                    f"agent-config endpoint not available (404) at {url} after {max_attempts} "
+                    f"attempts — the agent gateway isn't serving it (either boot up/initialization "
+                    f"failed, or the agent gateway was torn down by TTL reasons).",
+                    request=resp.request, response=resp,
+                )
+            resp.raise_for_status()
+            if attempt > 1:
+                logger.info(f"agent-config POST succeeded on attempt {attempt}")
+            return
+        except (httpx.ConnectError, httpx.ReadTimeout, httpx.RemoteProtocolError) as e:
+            last_exc = e
+            logger.warning(f"agent-config POST failed (attempt {attempt}/{max_attempts}, will retry): {e}")
+        except httpx.HTTPStatusError as e:
+            if e.response.status_code != 404:
+                raise
+            last_exc = e
+            logger.warning(f"agent-config POST got 404 (attempt {attempt}/{max_attempts}, will retry): {url}")
+        if attempt < max_attempts:
+            await asyncio.sleep(delay)
+            delay = min(delay * 2, 15.0)
+    assert last_exc is not None
+    raise last_exc
+
+
+async def send_a2a_message(
+    a2a_url: str,
+    parts: list[dict],
+    message_id: str,
+    context_id: Optional[str],
+    timeout_seconds: int,
+) -> tuple[str, Optional[str]]:
+    """POST /a2a message/send. Returns (task_id, resolved_context_id). Retries on connect errors."""
+    message: dict[str, Any] = {
+        "messageId": message_id, "role": "user",
+        "parts": parts,
+    }
+    if context_id is not None:
+        message["contextId"] = context_id
+    payload = {
+        "jsonrpc": "2.0", "id": "1", "method": "message/send",
+        "params": {"message": message, "configuration": {"blocking": False}},
+    }
+    max_attempts = 5
+    delay = 2.0
+    last_exc: Exception | None = None
+    for attempt in range(1, max_attempts + 1):
+        try:
+            async with httpx.AsyncClient() as client:
+                resp = await client.post(f"{a2a_url}/a2a", json=payload, timeout=timeout_seconds)
+            resp.raise_for_status()
+            send_body = resp.json()
+            if "error" in send_body:
+                raise RuntimeError(f"A2A message/send failed: {send_body['error']}")
+            send_result = send_body.get("result") or {}
+            task_id = send_result.get("id")
+            if not task_id:
+                raise RuntimeError(f"A2A message/send succeeded but returned no task id: {send_result}")
+            resolved_context_id = context_id or send_result.get("contextId")
+            return task_id, resolved_context_id
+        except (httpx.ConnectError, httpx.ConnectTimeout) as e:
+            last_exc = e
+            logger.warning(f"message/send connect error (attempt {attempt}/{max_attempts}, will retry): {e}")
+        if attempt < max_attempts:
+            await asyncio.sleep(delay)
+            delay = min(delay * 2, 15.0)
+    assert last_exc is not None
+    raise last_exc
+
+
+async def poll_a2a_task(
+    a2a_url: str,
+    task_id: str,
+    timeout_seconds: int,
+    poll_interval_seconds: int = 10,
+) -> dict:
+    """POST /a2a tasks/get until status.state is 'completed' or 'failed'. Returns the result dict."""
+    deadline = time.monotonic() + timeout_seconds
+    consecutive_failures = 0
+    while time.monotonic() < deadline:
+        backoff = min(poll_interval_seconds + consecutive_failures * 5, 60)
+        await asyncio.sleep(backoff)
+        try:
+            async with httpx.AsyncClient() as client:
+                resp = await client.post(f"{a2a_url}/a2a", json={
+                    "jsonrpc": "2.0", "id": "poll", "method": "tasks/get",
+                    "params": {"id": task_id},
+                }, timeout=30)
+            resp.raise_for_status()
+            data = resp.json()
+            if "error" in data:
+                consecutive_failures += 1
+                logger.warning(f"A2A poll returned error (consec={consecutive_failures}, will retry): {data['error']}")
+                continue
+            result = data["result"]
+        except httpx.HTTPStatusError as e:
+            if 400 <= e.response.status_code < 500:
+                raise
+            consecutive_failures += 1
+            logger.warning(f"A2A poll got {e.response.status_code} (consec={consecutive_failures}, will retry)")
+            continue
+        except httpx.HTTPError as e:
+            consecutive_failures += 1
+            logger.warning(f"A2A poll failed (consec={consecutive_failures}, will retry): {type(e).__name__}: {e}")
+            continue
+
+        consecutive_failures = 0
+        state = result["status"]["state"]
+        if state in _TERMINAL_TASK_STATES:
+            return result
+
+    raise TimeoutError(f"A2A task {task_id} did not complete within {timeout_seconds}s")
+
+
+@dataclass(frozen=True)
+class TerminalResponse:
+    """What's pulled from an A2A terminal status message: the text reply, the typed
+    ``structured_output`` (when ``output_format`` was set), and tool/error telemetry."""
+
+    response_text: str
+    tool_call_count: Optional[int] = None
+    error_type: Optional[str] = None
+    error_code: Optional[str] = None
+    error_class: Optional[str] = None
+    error_message: Optional[str] = None
+    structured_output: Optional[Any] = None
+
+    @classmethod
+    def from_message(cls, status_message: dict) -> "TerminalResponse":
+        """Parse an A2A terminal status message. Per A2A spec 0.3 the final agent message
+        (TaskUpdater.complete/failed) lives in result.status.message, not result.history."""
+        response_text = ""
+        tool_call_count: Optional[int] = None
+        error_type: Optional[str] = None
+        error_code: Optional[str] = None
+        error_class: Optional[str] = None
+        error_message: Optional[str] = None
+        structured_output: Optional[Any] = None
+        for part in (status_message or {}).get("parts", []):
+            kind = part.get("kind")
+            if kind == "text":
+                response_text = part["text"]
+            elif kind == "data":
+                data = part.get("data", {})
+                if "tool_call_count" in data:
+                    tool_call_count = data["tool_call_count"]
+                usage = data.get("usage")
+                if isinstance(usage, dict) and "tool_call_count" in usage:
+                    tool_call_count = usage["tool_call_count"]
+                if "error_type" in data:
+                    error_type = data["error_type"]
+                if "error_code" in data:
+                    error_code = data["error_code"]
+                if "error_class" in data:
+                    error_class = data["error_class"]
+                if "error_message" in data:
+                    error_message = data["error_message"]
+                if "structured_output" in data:
+                    structured_output = data["structured_output"]
+        return cls(
+            response_text=response_text,
+            tool_call_count=tool_call_count,
+            error_type=error_type,
+            error_code=error_code,
+            error_class=error_class,
+            error_message=error_message,
+            structured_output=structured_output,
+        )
+
+
+def extract_terminal_response(
+    status_message: dict,
+) -> tuple[str, Optional[int], Optional[str], Optional[str], Optional[str]]:
+    """Back-compat shim over ``TerminalResponse.from_message``; external callers
+    (the hub's agents.py) still unpack the 5-tuple. Drop once they migrate."""
+    tr = TerminalResponse.from_message(status_message)
+    return tr.response_text, tr.tool_call_count, tr.error_type, tr.error_class, tr.error_message

agent_env/a2a_agent/store.py

@@ -0,0 +1,201 @@
+"""A2A Agent store for MongoDB persistence."""
+from __future__ import annotations
+
+import dataclasses
+import logging
+import random
+import string
+from datetime import datetime, timedelta, timezone
+from typing import TYPE_CHECKING, Any, Optional, Self
+
+from agent_env.store.base import ConcurrentModificationError, NotFoundError
+from agent_env.config import get_config
+from agent_env.store.document_store import (
+    AbsentOrNull,
+    Filter,
+    LteOrAbsent,
+    UpdateSpec,
+    VersionedEntityStore,
+    VersionedEntityStoreCache,
+)
+from agent_env.store.query import QueryBuilder, to_document_query
+
+if TYPE_CHECKING:
+    from agent_env.a2a_agent.a2a_agent import A2AAgent
+
+logger = logging.getLogger(__name__)
+
+A2A_AGENTS_COLLECTION = "a2a_agents"
+A2A_AGENT_INSTANCES_COLLECTION = "a2a_agent_instances"
+
+
+class A2AAgentQuery(QueryBuilder["A2AAgent"]):
+    def __init__(self, store: Optional[A2AAgentStore] = None) -> None:
+        super().__init__()
+        self._store = store
+
+    def _clone(self) -> Self:
+        clone = A2AAgentQuery(self._store)
+        clone._filters = self._filters.copy()
+        clone._sort_field = self._sort_field
+        clone._sort_desc = self._sort_desc
+        clone._limit_value = self._limit_value
+        clone._offset_value = self._offset_value
+        return clone
+
+    def execute(self) -> list[A2AAgent]:
+        if self._store is None:
+            self._store = get_a2a_agent_store()
+        return self._store.execute_query(self)
+
+    def _execute_count(self) -> int:
+        if self._store is None:
+            self._store = get_a2a_agent_store()
+        return self._store.execute_count(self)
+
+
+class A2AAgentStore:
+    def __init__(self) -> None:
+        self._versioned_cache: VersionedEntityStoreCache[A2AAgent] = VersionedEntityStoreCache(
+            A2A_AGENTS_COLLECTION, self._serialize, self._deserialize, secondary_indexes=[["type"]],
+        )
+
+    @property
+    def _doc_store(self):
+        return get_config().get_document_store()
+
+    @property
+    def _versioned(self) -> VersionedEntityStore[A2AAgent]:
+        return self._versioned_cache.for_store(self._doc_store)
+
+    def _serialize(self, agent: A2AAgent) -> dict:
+        doc = agent.to_dict()
+        doc["created_at_utc"] = datetime.now(timezone.utc)
+        return doc
+
+    def get(self, id: str, version: Optional[int] = None) -> A2AAgent:
+        agent = self._versioned.get(id, version)
+        if agent is None:
+            version_str = f" version={version}" if version is not None else ""
+            raise NotFoundError(f"A2AAgent {id}{version_str} not found")
+        return agent
+
+    def next_version(self, id: str) -> int:
+        return self._versioned.next_version(id)
+
+    def put_document(self, agent: A2AAgent) -> A2AAgent:
+        agent.version = self._versioned.put(agent)
+        return agent
+
+    def update_metadata(self, id: str, version: int, old_metadata: dict[str, Any], new_metadata: dict[str, Any]) -> dict[str, Any]:
+        """Replace metadata using timestamp-based CAS (same pattern as EnvStore)."""
+        now = datetime.now(timezone.utc).isoformat()
+        old_ts = old_metadata.get("updated_at")
+        new_metadata = {**new_metadata, "updated_at": now}
+
+        ts_pred = LteOrAbsent(old_ts) if old_ts is not None else AbsentOrNull()
+        store = self._doc_store
+        matched = store.update(
+            A2A_AGENTS_COLLECTION,
+            Filter.of(id=id, version=version).where("metadata.updated_at", ts_pred),
+            UpdateSpec(set={"metadata": new_metadata}),
+        )
+
+        if matched == 0:
+            if store.find_one(A2A_AGENTS_COLLECTION, Filter.of(id=id, version=version)) is None:
+                raise NotFoundError(f"A2AAgent {id} version={version} not found")
+            raise ConcurrentModificationError(f"A2AAgent {id} version={version} metadata was modified concurrently")
+
+        return new_metadata
+
+    def _deserialize(self, doc: dict) -> A2AAgent:
+        from agent_env.a2a_agent.a2a_agent import A2AAgent
+
+        doc.pop("_id", None)
+        doc.pop("created_at_utc", None)
+        return A2AAgent.from_dict(doc)
+
+    def execute_query(self, query: A2AAgentQuery) -> list[A2AAgent]:
+        filt, sort = to_document_query(query)
+        return self._versioned.query(filt, sort, query._limit_value, query._offset_value)
+
+    def execute_count(self, query: A2AAgentQuery) -> int:
+        filt, _ = to_document_query(query)
+        return self._versioned.count(filt)
+
+
+_a2a_agent_store: Optional[A2AAgentStore] = None
+
+
+def get_a2a_agent_store() -> A2AAgentStore:
+    global _a2a_agent_store
+    if _a2a_agent_store is None:
+        _a2a_agent_store = A2AAgentStore()
+    return _a2a_agent_store
+
+
+def set_a2a_agent_store(store: A2AAgentStore) -> None:
+    global _a2a_agent_store
+    _a2a_agent_store = store
+
+
+def reset_a2a_agent_store() -> None:
+    global _a2a_agent_store
+    _a2a_agent_store = None
+
+
+class A2AAgentInstanceStore:
+    def __init__(self) -> None:
+        self._indexed = None
+
+    @property
+    def _doc_store(self):
+        # Resolved per call: a cached store outlives reset_config(), so a process that
+        # re-pointed would read the new config and write the old backend.
+        store = get_config().get_document_store()
+        if self._indexed is not store:
+            store.ensure_index(A2A_AGENT_INSTANCES_COLLECTION, ["instance_id"], unique=True)
+            self._indexed = store
+        return store
+
+    def create_instance(self, deployed: "DeployedA2AAgent", ttl_seconds: int) -> "DeployedA2AAgent":
+        from agent_env.a2a_agent.a2a_agent import DeployedA2AAgent
+        suffix = "".join(random.choices(string.ascii_lowercase + string.digits, k=8))
+        instance_id = f"{deployed.agent_id}-{suffix}"
+        now = datetime.now(timezone.utc)
+        created_at_utc = now.strftime("%Y-%m-%d %H:%M UTC")
+        expires_at_utc = (now + timedelta(seconds=ttl_seconds)).strftime("%Y-%m-%d %H:%M UTC")
+        doc = dataclasses.asdict(deployed) | {
+            "instance_id": instance_id,
+            "created_at_utc": created_at_utc,
+            "expires_at_utc": expires_at_utc,
+        }
+        self._doc_store.insert(A2A_AGENT_INSTANCES_COLLECTION, doc)
+        return DeployedA2AAgent.from_dict(doc)
+
+    def get(self, instance_id: str) -> "DeployedA2AAgent":
+        from agent_env.a2a_agent.a2a_agent import DeployedA2AAgent
+        doc = self._doc_store.find_one(A2A_AGENT_INSTANCES_COLLECTION, Filter.of(instance_id=instance_id))
+        if not doc:
+            raise NotFoundError(f"A2AAgentInstance '{instance_id}' not found")
+        return DeployedA2AAgent.from_dict(doc)
+
+
+_a2a_agent_instance_store: Optional[A2AAgentInstanceStore] = None
+
+
+def get_a2a_agent_instance_store() -> A2AAgentInstanceStore:
+    global _a2a_agent_instance_store
+    if _a2a_agent_instance_store is None:
+        _a2a_agent_instance_store = A2AAgentInstanceStore()
+    return _a2a_agent_instance_store
+
+
+def set_a2a_agent_instance_store(store: A2AAgentInstanceStore) -> None:
+    global _a2a_agent_instance_store
+    _a2a_agent_instance_store = store
+
+
+def reset_a2a_agent_instance_store() -> None:
+    global _a2a_agent_instance_store
+    _a2a_agent_instance_store = None