agentenv-framework 0.9.1254__py3-none-any.whl

agent_env/task_step/task_steps/a2a_agent_validator/verify_a2a_litellm_attribution_runtime.py

@@ -0,0 +1,173 @@
+"""Verify A2A agent actually emits LiteLLM attribution on outbound LLM calls.
+
+Optional behavioral counterpart to the static
+`VerifyA2ALitellmAttributionStep`. When a runtime advertises its private probe,
+this step asks whether it actually forwards project_id/task_id to the LLM call.
+The probe is not part of the general AgentEnv SDK contract.
+
+How it works:
+  1. The validator plants unique probe values in `initial_context.metadata`
+     before any prompt fires.
+  2. The standard `prompt_agent` flow POSTs those values via
+     `/ext/agent-config` and sends a prompt — the agent processes it,
+     makes an outbound LLM call, and records the attribution it claims
+     to have emitted in module-level state.
+  3. This step queries the agent's `/ext/attribution-probe` extension to
+     read back that recorded state.
+  4. Passes when the recorded values match the planted probe values.
+
+Catches:
+  - Bridge ignores project_id/task_id from agent-config
+  - Bridge sets the wrong env var / header names
+  - Agent advertises support but never wires it through
+
+Does NOT catch:
+  - CLI eats the env var and never emits the actual HTTP header (probe
+    trusts the bridge's claim about what it sent)
+  - Proxy strips the header on receive
+
+For full end-to-end coverage of the latter two, a spend-log probe would
+be needed (requires proxy admin access).
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import ClassVar, Optional
+
+import httpx
+
+from agent_env.task_step.context import TaskStepContext
+from agent_env.task_step.task_step import TaskStep, TaskStepDependency
+
+logger = logging.getLogger(__name__)
+
+EXT_ATTRIBUTION_PROBE_URI = "urn:agentenv:attribution-probe/v1"
+EXT_ATTRIBUTION_PROBE_PATH = "/ext/attribution-probe"
+
+
+class VerifyA2ALitellmAttributionRuntimeStep(TaskStep):
+    type: ClassVar[str] = "verify_a2a_litellm_attribution_runtime"
+
+    def __init__(
+        self,
+        id: str,
+        version: Optional[int],
+        a2a_agent_id: str,
+        expected_project_id: str,
+        expected_task_id: str,
+        a2a_agent_version: Optional[int] = None,
+        depends_on: Optional[list[TaskStepDependency]] = None,
+        fail_task_on_error: bool = True,
+    ):
+        super().__init__(id, version, depends_on=depends_on, fail_task_on_error=fail_task_on_error)
+        self.a2a_agent_id = a2a_agent_id
+        self.a2a_agent_version = a2a_agent_version
+        self.expected_project_id = expected_project_id
+        self.expected_task_id = expected_task_id
+
+    def to_dict(self) -> dict:
+        base = super().to_dict()
+        base["a2a_agent_id"] = self.a2a_agent_id
+        base["a2a_agent_version"] = self.a2a_agent_version
+        base["expected_project_id"] = self.expected_project_id
+        base["expected_task_id"] = self.expected_task_id
+        return base
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "VerifyA2ALitellmAttributionRuntimeStep":
+        return cls(
+            **cls._base_from_dict(data),
+            a2a_agent_id=data["a2a_agent_id"],
+            a2a_agent_version=data.get("a2a_agent_version"),
+            expected_project_id=data["expected_project_id"],
+            expected_task_id=data["expected_task_id"],
+        )
+
+    async def execute(self, context: TaskStepContext) -> TaskStepContext:
+        from agent_env.a2a_agent import A2AAgent
+
+        deployed = next((a for a in context.deployed_agents), None)
+        if deployed is None:
+            raise RuntimeError("No deployed agent found in context")
+
+        agent = A2AAgent.get(self.a2a_agent_id, self.a2a_agent_version)
+
+        # This is a runtime-specific diagnostic, not a general SDK capability.
+        # Its absence is therefore not a conformance failure.
+        probe_ext = A2AAgent.find_extension(deployed.a2a_card or {}, EXT_ATTRIBUTION_PROBE_URI)
+        if probe_ext is None:
+            logger.info(
+                "Agent '%s' does not advertise the attribution-probe extension "
+                "(%s) — skipping optional runtime attribution verification.",
+                self.a2a_agent_id, EXT_ATTRIBUTION_PROBE_URI,
+            )
+            validated = {
+                "passed": True,
+                "skipped": True,
+                "extension_present": False,
+                "reason": "optional runtime attribution probe not advertised",
+                "expected_project_id": self.expected_project_id,
+                "expected_task_id": self.expected_task_id,
+            }
+            agent.update_metadata({**agent.metadata, "validated_litellm_attribution_runtime": validated})
+            context.metadata.setdefault("verifications", {})["a2a_litellm_attribution_runtime"] = validated
+            return context
+
+        endpoint = (probe_ext.get("params") or {}).get("endpoint", EXT_ATTRIBUTION_PROBE_PATH)
+        url = (deployed.a2a_url or deployed.api_url) + endpoint
+        logger.info("Querying attribution probe at %s", url)
+
+        try:
+            async with httpx.AsyncClient(timeout=10.0) as client:
+                resp = await client.post(url, json={})
+                resp.raise_for_status()
+                probe_data = resp.json()
+        except Exception as e:
+            logger.warning("attribution-probe request failed: %s", e)
+            validated = {
+                "passed": False,
+                "extension_present": True,
+                "reason": f"probe request failed: {e}",
+                "expected_project_id": self.expected_project_id,
+                "expected_task_id": self.expected_task_id,
+            }
+            agent.update_metadata({**agent.metadata, "validated_litellm_attribution_runtime": validated})
+            context.metadata.setdefault("verifications", {})["a2a_litellm_attribution_runtime"] = validated
+            return context
+
+        last_seen = probe_data.get("last_seen_attribution") or {}
+        observed_project_id = last_seen.get("project_id")
+        observed_task_id = last_seen.get("task_id")
+
+        project_matches = observed_project_id == self.expected_project_id
+        task_matches = observed_task_id == self.expected_task_id
+        passed = project_matches and task_matches
+
+        if passed:
+            logger.info(
+                "Runtime attribution verified for '%s': observed projectId=%s taskId=%s",
+                self.a2a_agent_id, observed_project_id, observed_task_id,
+            )
+        else:
+            logger.warning(
+                "Runtime attribution mismatch for '%s': expected projectId=%s taskId=%s, "
+                "observed projectId=%s taskId=%s",
+                self.a2a_agent_id, self.expected_project_id, self.expected_task_id,
+                observed_project_id, observed_task_id,
+            )
+
+        validated = {
+            "passed": passed,
+            "extension_present": True,
+            "expected_project_id": self.expected_project_id,
+            "expected_task_id": self.expected_task_id,
+            "observed_project_id": observed_project_id,
+            "observed_task_id": observed_task_id,
+            "project_matches": project_matches,
+            "task_matches": task_matches,
+            "last_seen_at_utc": probe_data.get("last_seen_at_utc"),
+        }
+        agent.update_metadata({**agent.metadata, "validated_litellm_attribution_runtime": validated})
+        context.metadata.setdefault("verifications", {})["a2a_litellm_attribution_runtime"] = validated
+        return context

agent_env/task_step/task_steps/a2a_agent_validator/verify_a2a_modalities.py

@@ -0,0 +1,270 @@
+"""Verify which input modalities an A2A agent functionally supports.
+
+Covers `text`, `image/{png,jpeg,gif}`, `audio/{wav,mpeg,mp4,ogg}`,
+`application/pdf`, and `video/mp4`.
+The companion `PromptAgentTaskStep` probes (configured in
+`A2AAgent.validate()`) actually send the messages via A2A; this step purely
+grades the resulting `PromptResponse`s.
+
+Grading is a deterministic case-insensitive substring match against a known
+ground-truth string per probe. The outcome for each modality is one of:
+- `passed`: response contains the expected string.
+- `no_ingestion_evidence`: response succeeded but didn't reference the asset.
+- `protocol_reject`: response was appended with `error_type` (wrapper accepted
+  the message but the agent task ended in `failed`).
+- `probe_step_did_not_run`: no matching response — the upstream
+  `PromptAgentTaskStep` crashed before `context.prompt_responses.append(...)`.
+
+The gemini wrapper currently passes all 12 modality probes (text, image
+{png,jpeg,gif}, audio {wav,mpeg,mp4,ogg}, PDF, video/mp4, plus image/png
+delivered via `FileWithUri` over `s3://` and presigned `https://`). New
+A2A wrappers (claude-code, codex, grok) inherit this same probe set and
+will report which subset they actually ingest.
+"""
+
+from __future__ import annotations
+
+import base64
+import functools
+import logging
+from pathlib import Path
+from typing import ClassVar, Optional
+
+from agent_env.task_step.context import TaskStepContext
+from agent_env.entity_refs import EntityRef
+from agent_env.task_step.task_step import TaskStep, TaskStepDependency
+
+logger = logging.getLogger(__name__)
+
+
+_FIXTURES_DIR = Path(__file__).parent / "fixtures"
+
+
+@functools.lru_cache(maxsize=None)
+def _b64(name: str) -> str:
+    """Load a binary fixture and return its base64-encoded form.
+
+    Fixtures live as native binary files in `fixtures/` so they can be
+    inspected (e.g. `afplay clip.wav`, `qlmanage -p document.pdf`) and
+    replaced in-place without re-chunking base64 into Python source.
+    """
+    return base64.b64encode((_FIXTURES_DIR / name).read_bytes()).decode("ascii")
+
+
+# 8×8 solid red squares in three formats — unambiguously red so the probe's
+# deterministic "red" substring match holds. See `fixtures/README.md`.
+IMAGE_PROBE_PNG_B64 = _b64("red.png")
+IMAGE_PROBE_JPEG_B64 = _b64("red.jpg")
+IMAGE_PROBE_GIF_B64 = _b64("red.gif")
+
+# Per-format spoken "Say the word: <X>" clips with distinct target words —
+# WAV/banana, MP3/telephone, M4A/bicycle, OGG/rainbow. Distinct targets force
+# the model to actually decode each clip. See `fixtures/README.md` for the
+# encoding pipeline, voice/volume/padding rationale, and regen commands.
+AUDIO_PROBE_WAV_B64 = _b64("clip.wav")
+AUDIO_PROBE_MP3_B64 = _b64("clip.mp3")
+AUDIO_PROBE_M4A_B64 = _b64("clip.m4a")
+AUDIO_PROBE_OGG_B64 = _b64("clip.ogg")
+
+TEXT_PROBE_PROMPT = "Reply with exactly the word PINEAPPLE and nothing else."
+TEXT_PROBE_EXPECTED = "PINEAPPLE"
+
+IMAGE_PROBE_PROMPT = (
+    "What is the dominant color of the image attached? Answer in one word."
+)
+IMAGE_PROBE_EXPECTED = "red"
+
+AUDIO_PROBE_PROMPT = (
+    "An audio clip is attached. Follow the spoken instruction in it and "
+    "respond with only the requested word."
+)
+AUDIO_WAV_EXPECTED = "banana"
+AUDIO_MP3_EXPECTED = "telephone"
+AUDIO_M4A_EXPECTED = "bicycle"
+AUDIO_OGG_EXPECTED = "rainbow"
+
+# 1.1KB single-page PDF with "ELEPHANT" inside an instructional paragraph
+# (bare word alone trips letter-truncation edge cases). See `fixtures/README.md`.
+PDF_PROBE_PDF_B64 = _b64("document.pdf")
+PDF_PROBE_PROMPT = (
+    "A PDF document is attached. Read it and reply with only the target "
+    "word it specifies."
+)
+PDF_PROBE_EXPECTED = "elephant"
+
+# 8.7KB h.264 mp4 (3s @ 480×360, 15fps, no audio) showing "GIRAFFE" on a
+# static frame with on-screen instructional context. Delivered via
+# FileWithUri (s3://) at validator runtime. See `fixtures/README.md`.
+VIDEO_PROBE_MP4_B64 = _b64("clip.mp4")
+VIDEO_PROBE_PROMPT = (
+    "A short video clip is attached. Read the target word displayed in "
+    "the video frames and reply with only that word."
+)
+VIDEO_PROBE_EXPECTED = "giraffe"
+
+TEXT_PROBE_PARTS: list[dict] = [{"kind": "text", "text": TEXT_PROBE_PROMPT}]
+IMAGE_PROBE_PARTS: list[dict] = [
+    {"kind": "text", "text": IMAGE_PROBE_PROMPT},
+    {"kind": "file", "file": {
+        "bytes": IMAGE_PROBE_PNG_B64,
+        "mimeType": "image/png",
+        "name": "red.png",
+    }},
+]
+IMAGE_JPEG_PROBE_PARTS: list[dict] = [
+    {"kind": "text", "text": IMAGE_PROBE_PROMPT},
+    {"kind": "file", "file": {
+        "bytes": IMAGE_PROBE_JPEG_B64,
+        "mimeType": "image/jpeg",
+        "name": "red.jpg",
+    }},
+]
+IMAGE_GIF_PROBE_PARTS: list[dict] = [
+    {"kind": "text", "text": IMAGE_PROBE_PROMPT},
+    {"kind": "file", "file": {
+        "bytes": IMAGE_PROBE_GIF_B64,
+        "mimeType": "image/gif",
+        "name": "red.gif",
+    }},
+]
+AUDIO_WAV_PROBE_PARTS: list[dict] = [
+    {"kind": "text", "text": AUDIO_PROBE_PROMPT},
+    {"kind": "file", "file": {
+        "bytes": AUDIO_PROBE_WAV_B64,
+        "mimeType": "audio/wav",
+        "name": "clip.wav",
+    }},
+]
+AUDIO_MP3_PROBE_PARTS: list[dict] = [
+    {"kind": "text", "text": AUDIO_PROBE_PROMPT},
+    {"kind": "file", "file": {
+        "bytes": AUDIO_PROBE_MP3_B64,
+        "mimeType": "audio/mpeg",
+        "name": "clip.mp3",
+    }},
+]
+AUDIO_M4A_PROBE_PARTS: list[dict] = [
+    {"kind": "text", "text": AUDIO_PROBE_PROMPT},
+    {"kind": "file", "file": {
+        "bytes": AUDIO_PROBE_M4A_B64,
+        "mimeType": "audio/mp4",
+        "name": "clip.m4a",
+    }},
+]
+AUDIO_OGG_PROBE_PARTS: list[dict] = [
+    {"kind": "text", "text": AUDIO_PROBE_PROMPT},
+    {"kind": "file", "file": {
+        "bytes": AUDIO_PROBE_OGG_B64,
+        "mimeType": "audio/ogg",
+        "name": "clip.ogg",
+    }},
+]
+PDF_PROBE_PARTS: list[dict] = [
+    {"kind": "text", "text": PDF_PROBE_PROMPT},
+    {"kind": "file", "file": {
+        "bytes": PDF_PROBE_PDF_B64,
+        "mimeType": "application/pdf",
+        "name": "document.pdf",
+    }},
+]
+
+
+class VerifyA2AModalitiesStep(TaskStep):
+    """Grade modality probes by reading `context.prompt_responses`.
+
+    `probes` is a list of dicts, each `{"modality": str, "prompt_id": str,
+    "expected": str}`. Each `prompt_id` should match an upstream
+    `PromptAgentTaskStep`'s `prompt_id`.
+    """
+
+    type: ClassVar[str] = "verify_a2a_modalities"
+    entity_refs = (EntityRef.agent("a2a_agent_id"),)
+
+    def __init__(
+        self,
+        id: str,
+        version: Optional[int],
+        a2a_agent_id: str,
+        probes: list[dict],
+        depends_on: Optional[list[TaskStepDependency]] = None,
+        fail_task_on_error: bool = True,
+    ):
+        super().__init__(id, version, depends_on=depends_on, fail_task_on_error=fail_task_on_error)
+        self.a2a_agent_id = a2a_agent_id
+        self.probes = probes
+
+    def to_dict(self) -> dict:
+        base = super().to_dict()
+        base["a2a_agent_id"] = self.a2a_agent_id
+        base["probes"] = self.probes
+        return base
+
+    @classmethod
+    def from_dict(cls, data: dict) -> VerifyA2AModalitiesStep:
+        return cls(
+            **cls._base_from_dict(data),
+            a2a_agent_id=data["a2a_agent_id"],
+            probes=data["probes"],
+        )
+
+    async def execute(self, context: TaskStepContext) -> TaskStepContext:
+        from agent_env.a2a_agent import A2AAgent
+
+        results: dict[str, dict] = {}
+        for probe in self.probes:
+            modality = probe["modality"]
+            prompt_id = probe["prompt_id"]
+            expected = probe["expected"]
+            response = next(
+                (r for r in context.prompt_responses if r.prompt_id == prompt_id),
+                None,
+            )
+            if response is None:
+                # Upstream PromptAgentTaskStep crashed before appending its response
+                # (e.g., wrapper returned HTTP 4xx on message/send).
+                results[modality] = {
+                    "supported": False,
+                    "reason": "probe_step_did_not_run",
+                }
+            elif response.error_type:
+                results[modality] = {
+                    "supported": False,
+                    "reason": "protocol_reject",
+                    "error_type": response.error_type,
+                    "probe_response": (response.response or "")[:200],
+                }
+            elif expected.lower() in (response.response or "").lower():
+                results[modality] = {
+                    "supported": True,
+                    "reason": "passed",
+                    "probe_response": response.response[:200],
+                }
+            else:
+                results[modality] = {
+                    "supported": False,
+                    "reason": "no_ingestion_evidence",
+                    "probe_response": (response.response or "")[:200],
+                }
+            logger.info(
+                f"Modality '{modality}' graded: supported={results[modality]['supported']} "
+                f"reason={results[modality]['reason']}"
+            )
+
+        # Echo what the AgentCard declared, for cross-reference with the
+        # actual probe results.
+        agent = A2AAgent.get(self.a2a_agent_id)
+        card = agent.metadata.get("agent_card") or {}
+        declared = {
+            "default_input_modes": card.get("defaultInputModes"),
+            "default_output_modes": card.get("defaultOutputModes"),
+        }
+
+        agent.update_metadata({
+            **agent.metadata,
+            "validated_modalities": {"input": results, "declared": declared},
+        })
+        context.metadata.setdefault("verifications", {})["a2a_modalities"] = {
+            "input": results,
+            "declared": declared,
+        }
+        return context

agent_env/task_step/task_steps/a2a_agent_validator/verify_a2a_peer_agents.py

@@ -0,0 +1,102 @@
+"""Verify the A2A peer-agents extension end-to-end."""
+from __future__ import annotations
+
+import logging
+from typing import ClassVar, Optional
+
+import httpx
+
+from agent_env.a2a_agent import A2AAgent
+from agent_env.task_step.context import TaskStepContext
+from agent_env.entity_refs import EntityRef
+from agent_env.task_step.task_step import TaskStep, TaskStepDependency
+
+logger = logging.getLogger(__name__)
+
+
+class VerifyA2APeerAgentsStep(TaskStep):
+    """Consolidate four checks of the peer-agents extension and stamp results on agent.metadata.
+
+    Assumes upstream steps have deployed the agent under test plus a peer-target
+    agent, peered them via PeerAgentsTaskStep, and prompted the agent under test
+    to use peer_send_message to relay an `expected_token` from the peer."""
+
+    type: ClassVar[str] = "verify_a2a_peer_agents"
+    entity_refs = (EntityRef.agent("a2a_agent_id", version_field="a2a_agent_version"),)
+
+    def __init__(self, id: str, version: Optional[int], a2a_agent_id: str, a2a_agent_version: Optional[int], peer_prompt_id: str, expected_token: str, depends_on: Optional[list[TaskStepDependency]] = None, fail_task_on_error: bool = True):
+        super().__init__(id, version, depends_on=depends_on, fail_task_on_error=fail_task_on_error)
+        self.a2a_agent_id = a2a_agent_id
+        self.a2a_agent_version = a2a_agent_version
+        self.peer_prompt_id = peer_prompt_id
+        self.expected_token = expected_token
+
+    def to_dict(self) -> dict:
+        return {**super().to_dict(), "a2a_agent_id": self.a2a_agent_id, "a2a_agent_version": self.a2a_agent_version, "peer_prompt_id": self.peer_prompt_id, "expected_token": self.expected_token}
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "VerifyA2APeerAgentsStep":
+        return cls(**cls._base_from_dict(data), a2a_agent_id=data["a2a_agent_id"], a2a_agent_version=data.get("a2a_agent_version"), peer_prompt_id=data["peer_prompt_id"], expected_token=data["expected_token"])
+
+    async def execute(self, context: TaskStepContext) -> TaskStepContext:
+        deployed = next((a for a in context.deployed_agents if a.agent_name == TaskStep.DEFAULT_AGENT_NAME), None)
+        if deployed is None:
+            raise RuntimeError(f"VerifyA2APeerAgentsStep: no deployed agent named '{TaskStep.DEFAULT_AGENT_NAME}' in context")
+        a2a_url = deployed.a2a_url or deployed.api_url
+        if not a2a_url:
+            raise RuntimeError(f"VerifyA2APeerAgentsStep: deployed agent has no a2a_url")
+
+        extension = A2AAgent.find_extension(deployed.a2a_card or {}, A2AAgent.EXT_PEER_AGENTS)
+        extension_advertised = extension is not None
+        endpoint = (extension.get("params") or {}).get("endpoint", "/ext/peer-agents") if extension else "/ext/peer-agents"
+        post_ok, list_ok = await self._check_round_trip(a2a_url, endpoint)
+        sidecar_registered = await self._check_sidecar_registered(a2a_url)
+        peer_call_ok, response_excerpt = self._check_peer_call(context)
+
+        peer_entry = {
+            "supported": extension_advertised and post_ok and list_ok and sidecar_registered and peer_call_ok,
+            "extension_advertised": extension_advertised,
+            "methods": {"set": {"supported": post_ok}, "list": {"supported": list_ok}},
+            "sidecar_self_registered": sidecar_registered,
+            "peer_call_succeeded": peer_call_ok,
+            "peer_call_response_excerpt": response_excerpt,
+        }
+
+        agent = A2AAgent.get(self.a2a_agent_id, self.a2a_agent_version)
+        validated_ext = agent.metadata.get("validated_a2a_extensions", {})
+        validated_ext[A2AAgent.EXT_PEER_AGENTS] = peer_entry
+        agent.update_metadata({**agent.metadata, "validated_a2a_extensions": validated_ext})
+
+        context.metadata.setdefault("verifications", {})["a2a_peer_agents"] = peer_entry
+        logger.info(f"peer-agents validation: {peer_entry}")
+        return context
+
+    @staticmethod
+    async def _check_round_trip(a2a_url: str, endpoint: str) -> tuple[bool, bool]:
+        probe = {"peers": [{"name": "probe", "url": "http://example.invalid"}]}
+        post_ok = list_ok = False
+        async with httpx.AsyncClient() as client:
+            try:
+                post_resp = await client.post(f"{a2a_url}{endpoint}", json=probe, timeout=30)
+                post_ok = post_resp.status_code < 400
+                get_resp = await client.get(f"{a2a_url}{endpoint}", timeout=30)
+                list_ok = get_resp.status_code < 400 and any(p.get("name") == "probe" for p in (get_resp.json().get("peers") or []))
+            finally:
+                await client.post(f"{a2a_url}{endpoint}", json={"peers": []}, timeout=30)
+        return post_ok, list_ok
+
+    @staticmethod
+    async def _check_sidecar_registered(a2a_url: str) -> bool:
+        async with httpx.AsyncClient() as client:
+            resp = await client.get(f"{a2a_url}/ext/mcp-config", timeout=30)
+        if resp.status_code >= 400:
+            return False
+        servers = (resp.json() or {}).get("mcp_servers") or {}
+        return any("127.0.0.1" in (info.get("url") or "") and (info.get("url") or "").endswith("/mcp") for info in servers.values())
+
+    def _check_peer_call(self, context: TaskStepContext) -> tuple[bool, str]:
+        prompt_resp = next((p for p in context.prompt_responses if p.prompt_id == self.peer_prompt_id), None)
+        if prompt_resp is None:
+            return False, ""
+        response = prompt_resp.response or ""
+        return self.expected_token in response, response[:200]