a2a_adapter_example.py

"""
A2AAdapter Example
==================

Demonstrates A2AAdapter — the **server-side** (receiver) half of the
Agent2Agent (A2A) protocol.  Use A2AAdapter in any agent service you deploy so
that other services (A2AClient, orchestrators) can call it over HTTP using the
standard A2A JSON-RPC 2.0 wire format.

Companion: a2a_client_example.py shows the **caller** side.

What A2AAdapter does
--------------------
- Constructs and serves the A2A Agent Card (``GET /.well-known/agent.json``)
- Parses incoming ``POST /rpc`` JSON-RPC 2.0 ``tasks/send`` requests
- Extracts the task text and context metadata from the request body
- Calls YOUR execution callback with ``(task_text, context)``
- Formats the ``AgentResult`` you return into a valid A2A JSON-RPC response
- Formats protocol-level errors (unknown method, bad payload, etc.)

What A2AAdapter does NOT do
---------------------------
- It does not know about tools, checkpoints, routing, or any other
  orchestration concern.  Those belong to the agent behind it.
- It does not start an HTTP server — you bind it to whatever HTTP framework
  you use (aiohttp, FastAPI, stdlib http.server, Django, etc.).

Wire protocol summary
---------------------

    Discovery::

        GET /.well-known/agent.json
        → {"name": "...", "description": "...", "url": "...",
           "version": "...", "skills": [...]}

    Invocation::

        POST /rpc
        Content-Type: application/json

        {"jsonrpc": "2.0", "id": "<uuid>", "method": "tasks/send",
         "params": {"id": "<task-uuid>",
                    "message": {"role": "user",
                                "parts": [{"type": "text", "text": "<task>"}],
                                "metadata": {"obo_token": "...", "locale": "en"}}}}

        Response:
        {"jsonrpc": "2.0", "id": "<uuid>",
         "result": {"id": "<task-uuid>",
                    "status": {"state": "completed"},
                    "artifacts": [{"parts": [{"type": "text", "text": "<output>"}]}],
                    "metadata": {"steps": [...]}}}

Examples in this file
---------------------
1. A2AAdapter in isolation — no HTTP server needed (test the protocol logic directly)
2. Building and serialising the Agent Card
3. Parsing ``tasks/send`` requests and building responses manually
4. Error handling — malformed requests, unknown methods, execution failures
5. Minimal stdlib HTTP server (no extra dependencies) using A2AAdapter
6. Integration test helper — simulate A2AClient ↔ A2AAdapter round-trip in memory

Prerequisites::

    pip install -e packages/orchestration-layer
    pip install -e packages/shared-core

Usage::

    cd packages/orchestration-layer
    python examples/a2a_adapter_example.py
"""

import asyncio
import json
import os
import threading
from http.server import BaseHTTPRequestHandler, HTTPServer
from pathlib import Path
from typing import Any, Dict
from uuid import uuid4

from dotenv import load_dotenv

from gmf_forge_ai_shared_core.observability import BasicLogger, BasicMetricsCollector

from gmf_forge_ai_orchestration.agents.base import AgentResult, AgentStep
from gmf_forge_ai_orchestration.protocols.a2a.a2a_adapter import A2AAdapter, A2AAdapterError

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

load_dotenv(Path(__file__).parent / ".env")

WORKSPACE_ROOT = Path(__file__).parent.parent.parent.parent
logger = BasicLogger("a2a_adapter_example")
metrics = BasicMetricsCollector()

# Service identity — what the adapter advertises to callers
AGENT_SERVICE_URL = "http://localhost:8080"
AGENT_VERSION = "1.0.0"


# ---------------------------------------------------------------------------
# Shared execution callback
# ---------------------------------------------------------------------------

async def my_agent_execute(task: str, context: Dict[str, Any]) -> AgentResult:
    """
    The business-logic callback that A2AAdapter calls after parsing a request.

    In production this would call a ReActAgent, PlanExecuteAgent, or any other
    agent.  Here we simulate a simple echo + context-aware response.

    Args:
        task: The plain-text task extracted from the A2A request body.
        context: The ``message.metadata`` dict forwarded by the caller.
                 May contain obo_token, user_assertion, locale, session_id, etc.
    """
    logger.info(
        "Agent executing task",
        task_preview=task[:80],
        context_keys=list(context.keys()),
    )

    # Simulate work — in production: agent.execute(task, context)
    output = (
        f"Task received: '{task}'. "
        f"Processing with context keys: {list(context.keys())}. "
        "Result: analysis complete."
    )

    # Build a realistic AgentResult with steps
    steps = [
        AgentStep(
            thought="Analysing the incoming task and available context.",
            action="analyse_task",
            action_input={"task": task, "context_keys": list(context.keys())},
            observation="Task analysis complete. No tool calls required for this request.",
        ),
        AgentStep(
            thought="Composing the final answer.",
            action="Final Answer",
            action_input={"answer": output},
            observation=output,
        ),
    ]

    return AgentResult(
        output=output,
        steps=steps,
        success=True,
        metadata={
            "agent_id": "my-example-agent",
            "task_length": len(task),
        },
    )


# ---------------------------------------------------------------------------
# Example 1: A2AAdapter in isolation
#             — call handle_rpc directly, no HTTP server needed
# ---------------------------------------------------------------------------

async def example_adapter_in_isolation() -> None:
    """
    Test the A2AAdapter protocol logic without any HTTP server.

    This is the most useful pattern for unit tests — construct the JSON-RPC
    body yourself, call handle_rpc, and assert on the JSON-RPC response.
    No port binding, no network calls.
    """
    logger.info("=" * 60)
    logger.info("Example 1: A2AAdapter in isolation (no HTTP server)")
    logger.info("=" * 60)

    adapter = A2AAdapter(
        agent_id="example-agent",
        description="Example A2A agent for demonstration purposes.",
        url=AGENT_SERVICE_URL,
        version=AGENT_VERSION,
        skills=[
            {"id": "analysis", "name": "Document Analysis",
             "description": "Analyse documents and answer questions about them."},
            {"id": "summary", "name": "Summarisation",
             "description": "Produce concise summaries of long content."},
        ],
        logger=logger,
    )

    # Construct a valid tasks/send JSON-RPC request body
    request_body = {
        "jsonrpc": "2.0",
        "id": str(uuid4()),
        "method": "tasks/send",
        "params": {
            "id": str(uuid4()),
            "message": {
                "role": "user",
                "parts": [{"type": "text", "text": "Analyse the Q1 financial report."}],
                "metadata": {
                    "locale": "en-US",
                    "session_id": "sess-001",
                    "obo_token": "eyJhbGci...(truncated)",
                },
            },
        },
    }

    # handle_rpc parses the request, calls our execution callback, and
    # returns a complete JSON-RPC response dict ready to serialise.
    response = await adapter.handle_rpc(request_body, my_agent_execute)

    # The response is a plain dict — serialise it as JSON for HTTP
    response_json = json.dumps(response, indent=2)
    logger.info("Protocol response received", method="tasks/send")
    print("\n--- JSON-RPC response (Example 1) ---")
    print(response_json[:800])
    print("--- End response ---\n")

    # Validate the response structure
    assert response.get("jsonrpc") == "2.0", "jsonrpc version mismatch"
    assert "result" in response, "Expected 'result' key in successful response"
    task_result = response["result"]
    assert task_result["status"]["state"] in ("completed", "failed"), "Unexpected task state"
    logger.info("Response validation passed", state=task_result["status"]["state"])


# ---------------------------------------------------------------------------
# Example 2: Agent Card
#             — advertise capabilities to callers via discovery
# ---------------------------------------------------------------------------

async def example_agent_card() -> None:
    """
    Agent Card construction and serialisation.

    Callers (A2AClient, orchestrators, service registries) send
    GET /.well-known/agent.json to discover the agent's identity and skills.

    The AgentCard is a Pydantic model — call .model_dump() to get a plain dict
    ready for JSON serialisation or HTTP response bodies.
    """
    logger.info("=" * 60)
    logger.info("Example 2: Agent Card")
    logger.info("=" * 60)

    adapter = A2AAdapter(
        agent_id="policy-review-agent",
        description="Reviews corporate policies for compliance with EU AI Act and GDPR.",
        url="https://policy-review.internal.company.com",
        version="2.1.0",
        skills=[
            {
                "id": "eu_ai_act",
                "name": "EU AI Act compliance check",
                "description": "Assess policies against EU AI Act requirements.",
            },
            {
                "id": "gdpr_review",
                "name": "GDPR review",
                "description": "Identify GDPR obligations and data subject rights.",
            },
            {
                "id": "risk_classification",
                "name": "AI risk classification",
                "description": "Classify AI system risk level (minimal, limited, high, unacceptable).",
            },
        ],
        logger=logger,
    )

    # agent_card() returns a fully populated AgentCard Pydantic model.
    card = adapter.agent_card()
    card_dict = card.model_dump()

    logger.info("Agent Card built", agent_id=card.name, skills=len(card.skills))
    print("\n--- Agent Card (GET /.well-known/agent.json) ---")
    print(json.dumps(card_dict, indent=2))
    print("--- End agent card ---\n")

    # In an HTTP handler:
    #   response.body = json.dumps(card_dict)
    #   response.headers["Content-Type"] = "application/json"
    assert card.name == "policy-review-agent"
    assert len(card.skills) == 3
    logger.info("Agent card assertions passed")


# ---------------------------------------------------------------------------
# Example 3: Manual parse + build response
#             — fine-grained control over each step
# ---------------------------------------------------------------------------

async def example_manual_parse_and_build() -> None:
    """
    Manual request parsing and response construction.

    Use parse_tasks_send() + build_task_response() directly when you need
    to inspect the parsed values (e.g. log the task_id, check the context,
    apply request-level auth) before calling your agent.

    Equivalent to calling handle_rpc() in two steps.
    """
    logger.info("=" * 60)
    logger.info("Example 3: Manual parse and build")
    logger.info("=" * 60)

    adapter = A2AAdapter(
        agent_id="manual-example-agent",
        description="Demo adapter for manual parse/build example.",
        url=AGENT_SERVICE_URL,
        version=AGENT_VERSION,
        logger=logger,
    )

    request_body = {
        "jsonrpc": "2.0",
        "id": "req-123",
        "method": "tasks/send",
        "params": {
            "id": "task-456",
            "message": {
                "role": "user",
                "parts": [{"type": "text", "text": "Summarise the AML policy document."}],
                "metadata": {
                    "user_assertion": "bearer-token-xyz",
                    "locale": "en-GB",
                },
            },
        },
    }

    # Step 1 — Parse the JSON-RPC body.
    #   Returns: (rpc_id, task_id, task_text, context)
    rpc_id, task_id, task_text, context = adapter.parse_tasks_send(request_body)
    logger.info(
        "Parsed tasks/send",
        rpc_id=rpc_id,
        task_id=task_id,
        task_preview=task_text[:60],
        context_keys=list(context.keys()),
    )

    # Step 2 — Run the agent (or apply any pre-execution logic here).
    #   e.g. validate the user_assertion from context before executing:
    user_assertion = context.get("user_assertion", "")
    if not user_assertion:
        logger.warning("No user_assertion in context — proceeding anyway for demo")

    result = await my_agent_execute(task_text, context)

    # Step 3 — Build the JSON-RPC response dict.
    response = adapter.build_task_response(rpc_id, task_id, result)
    logger.info(
        "Response built",
        state=response["result"]["status"]["state"],
        artifact_count=len(response["result"].get("artifacts", [])),
    )

    assert response["result"]["id"] == task_id, "Task ID mismatch in response"
    logger.info("Manual parse/build assertions passed")


# ---------------------------------------------------------------------------
# Example 4: Error handling
#             — malformed requests, unknown methods, execution failures
# ---------------------------------------------------------------------------

async def example_error_handling() -> None:
    """
    A2AAdapter error handling.

    A2AAdapter converts protocol-level errors into well-formed JSON-RPC error
    responses rather than raising exceptions to the HTTP handler.

    Error codes follow JSON-RPC 2.0 conventions:
    - -32601 : Method not found
    - -32602 : Invalid params (bad payload shape)
    - -32603 : Internal error (agent execution failure)
    """
    logger.info("=" * 60)
    logger.info("Example 4: Error handling")
    logger.info("=" * 60)

    adapter = A2AAdapter(
        agent_id="error-demo-agent",
        description="Agent for error handling examples.",
        url=AGENT_SERVICE_URL,
        version=AGENT_VERSION,
        logger=logger,
    )

    # --- 4a: Unsupported method ---
    unsupported_method = {
        "jsonrpc": "2.0",
        "id": "req-unsupported",
        "method": "tasks/cancel",   # A2AAdapter only handles tasks/send
        "params": {},
    }
    response = await adapter.handle_rpc(unsupported_method, my_agent_execute)
    logger.info("4a — Unsupported method", error=response.get("error"))
    assert "error" in response, "Expected error response for unsupported method"
    assert response["error"]["code"] == -32601

    # --- 4b: Malformed request (params is not a dict/Mapping) ---
    # Missing params defaults to {}, which is valid. To trigger the params
    # validation branch, pass params as a non-Mapping value (e.g. a string).
    malformed = {
        "jsonrpc": "2.0",
        "id": "req-malformed",
        "method": "tasks/send",
        "params": "this-is-not-a-dict",  # must be a Mapping
    }
    response = await adapter.handle_rpc(malformed, my_agent_execute)
    logger.info("4b — Malformed request", error=response.get("error"))
    assert "error" in response, "Expected error response for malformed request"

    # --- 4c: Agent execution failure ---
    async def failing_execute(task: str, context: Dict[str, Any]) -> AgentResult:
        raise RuntimeError("Agent internal error: database connection lost")

    valid_request = {
        "jsonrpc": "2.0",
        "id": "req-exec-fail",
        "method": "tasks/send",
        "params": {
            "id": str(uuid4()),
            "message": {
                "role": "user",
                "parts": [{"type": "text", "text": "Any task"}],
                "metadata": {},
            },
        },
    }
    response = await adapter.handle_rpc(valid_request, failing_execute)
    logger.info("4c — Execution failure", error=response.get("error"))
    assert "error" in response, "Expected error response when agent raises"
    assert response["error"]["code"] == -32603

    # --- 4d: Manually build a JSON-RPC error response ---
    error_response = adapter.build_jsonrpc_error(
        rpc_id="req-custom-error",
        code=-32700,
        message="Parse error: request body was not valid JSON",
        data={"received_content_type": "text/plain"},
    )
    logger.info("4d — Custom error response built", error=error_response.get("error"))

    logger.info("All error handling assertions passed")


# ---------------------------------------------------------------------------
# Example 5: Minimal stdlib HTTP server
#             — wire A2AAdapter to Python's built-in http.server
# ---------------------------------------------------------------------------

def run_minimal_http_server(adapter: A2AAdapter, host: str = "localhost", port: int = 8080) -> HTTPServer:
    """
    Create and return a minimal HTTP server that handles the two A2A endpoints:

    - ``GET /.well-known/agent.json`` → serves the Agent Card
    - ``POST /rpc``                   → processes tasks/send JSON-RPC requests

    Returns the HTTPServer instance so the caller can start/stop it.

    In production use a production-grade ASGI framework (FastAPI, aiohttp,
    Starlette) instead of http.server.  This example exists purely to show the
    A2AAdapter integration without any extra dependencies.
    """
    loop = asyncio.new_event_loop()

    class A2ARequestHandler(BaseHTTPRequestHandler):
        def log_message(self, format: str, *args: Any) -> None:  # noqa: A002
            logger.info(f"HTTP {format % args}")

        def do_GET(self) -> None:  # noqa: N802
            if self.path == "/.well-known/agent.json":
                card = adapter.agent_card().model_dump()
                body = json.dumps(card).encode()
                self.send_response(200)
                self.send_header("Content-Type", "application/json")
                self.send_header("Content-Length", str(len(body)))
                self.end_headers()
                self.wfile.write(body)
            else:
                self.send_response(404)
                self.end_headers()

        def do_POST(self) -> None:  # noqa: N802
            if self.path != "/rpc":
                self.send_response(404)
                self.end_headers()
                return

            content_length = int(self.headers.get("Content-Length", 0))
            raw_body = self.rfile.read(content_length)

            try:
                request_body = json.loads(raw_body)
            except json.JSONDecodeError:
                self.send_response(400)
                self.end_headers()
                return

            # Run the async handle_rpc in the event loop created for this server.
            response = loop.run_until_complete(
                adapter.handle_rpc(request_body, my_agent_execute)
            )

            body = json.dumps(response).encode()
            self.send_response(200)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            self.wfile.write(body)

    return HTTPServer((host, port), A2ARequestHandler)


async def example_minimal_http_server() -> None:
    """
    Minimal stdlib HTTP server.

    Starts a temporary server on localhost:18080, fires a few A2A requests
    against it using httpx (same library used by A2AClient), then shuts down.

    The server is started in a background thread so the asyncio event loop
    can continue running the client calls on the main thread.
    """
    logger.info("=" * 60)
    logger.info("Example 5: Minimal stdlib HTTP server")
    logger.info("=" * 60)

    import httpx

    adapter = A2AAdapter(
        agent_id="http-demo-agent",
        description="Agent served via minimal stdlib HTTP server.",
        url="http://localhost:18080",
        version="1.0.0",
        skills=[{"id": "demo", "name": "Demo", "description": "Demo skill."}],
        logger=logger,
    )

    server = run_minimal_http_server(adapter, host="localhost", port=18080)
    server_thread = threading.Thread(target=server.serve_forever, daemon=True)
    server_thread.start()
    logger.info("HTTP server started on http://localhost:18080")

    # Small delay for server readiness
    await asyncio.sleep(0.1)

    async with httpx.AsyncClient() as client:
        # 5a — Agent card discovery
        resp = await client.get("http://localhost:18080/.well-known/agent.json")
        card = resp.json()
        logger.info("Agent card received", name=card.get("name"), url=card.get("url"))
        assert card["name"] == "http-demo-agent"

        # 5b — tasks/send
        payload = {
            "jsonrpc": "2.0",
            "id": str(uuid4()),
            "method": "tasks/send",
            "params": {
                "id": str(uuid4()),
                "message": {
                    "role": "user",
                    "parts": [{"type": "text", "text": "What is the weather today?"}],
                    "metadata": {"locale": "en"},
                },
            },
        }
        resp = await client.post(
            "http://localhost:18080/rpc",
            json=payload,
            headers={"Content-Type": "application/json"},
        )
        result = resp.json()
        logger.info(
            "tasks/send response",
            state=result["result"]["status"]["state"],
            output_preview=result["result"]["artifacts"][0]["parts"][0]["text"][:80],
        )
        assert result["result"]["status"]["state"] == "completed"

    server.shutdown()
    logger.info("HTTP server shut down cleanly")


# ---------------------------------------------------------------------------
# Example 6: In-memory A2AClient ↔ A2AAdapter round-trip
#             — integration test helper
# ---------------------------------------------------------------------------

async def example_in_memory_round_trip() -> None:
    """
    Full A2AClient ↔ A2AAdapter round-trip without a real HTTP server.

    This pattern is useful for integration tests: instead of starting a server,
    intercept the HTTP layer with ``respx`` (already in dev requirements) so you
    can verify that the full JSON-RPC serialisation/deserialisation cycle works
    end-to-end.

    This example shows the pattern using raw A2AAdapter calls to simulate what
    A2AClient would send and receive, proving protocol correctness without network I/O.
    """
    logger.info("=" * 60)
    logger.info("Example 6: In-memory round-trip simulation")
    logger.info("=" * 60)

    adapter = A2AAdapter(
        agent_id="roundtrip-agent",
        description="Agent for round-trip simulation.",
        url=AGENT_SERVICE_URL,
        version="1.0.0",
        logger=logger,
    )

    # Simulate what A2AClient._call_remote would POST as its JSON-RPC payload.
    # (A2AClient constructs this using JsonRpcRequest Pydantic model internally.)
    rpc_request_id = str(uuid4())
    a2a_task_id = str(uuid4())

    simulated_client_payload = {
        "jsonrpc": "2.0",
        "id": rpc_request_id,
        "method": "tasks/send",
        "params": {
            "id": a2a_task_id,
            "message": {
                "role": "user",
                "parts": [{"type": "text", "text": "Generate a risk summary for project Alpha."}],
                "metadata": {
                    "obo_token": "Bearer obo-token-placeholder",
                    "user_assertion": "original-access-token",
                    "locale": "en-US",
                },
            },
        },
    }

    # Pass through A2AAdapter — returns the JSON-RPC response dict
    response = await adapter.handle_rpc(simulated_client_payload, my_agent_execute)

    logger.info(
        "Round-trip complete",
        jsonrpc=response.get("jsonrpc"),
        response_id=response.get("id"),
        task_state=response["result"]["status"]["state"],
    )

    # Verify the response carries back the same rpc_id and task_id
    assert response["id"] == rpc_request_id, "RPC ID must be echoed back"
    assert response["result"]["id"] == a2a_task_id, "Task ID must be preserved"

    # Verify the output text is present in artifacts
    artifacts = response["result"].get("artifacts", [])
    assert len(artifacts) > 0, "Expected at least one artifact"
    assert artifacts[0]["parts"][0]["type"] == "text"
    output_text = artifacts[0]["parts"][0]["text"]
    logger.info("Output from adapter", output_preview=output_text[:80])

    # Verify steps are propagated in metadata
    steps = response["result"]["metadata"].get("steps", [])
    logger.info("Steps propagated", count=len(steps))
    assert len(steps) > 0, "Expected steps in metadata"

    logger.info("All round-trip assertions passed")


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

async def main() -> None:
    await example_adapter_in_isolation()
    await example_agent_card()
    await example_manual_parse_and_build()
    await example_error_handling()
    await example_minimal_http_server()
    await example_in_memory_round_trip()

    logger.info("All A2AAdapter examples completed.")


if __name__ == "__main__":
    asyncio.run(main())