1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
# echo_agent.py
"""
Simple Echo Agent - Demonstrates AgentWeave basics.

This agent has a single capability: echo messages back to the caller.
All security (identity, mTLS, authorization) is handled by the SDK.
"""

import asyncio
from typing import Dict, Any

from agentweave import SecureAgent, capability
from agentweave.types import TaskResult, Message, TextPart


class EchoAgent(SecureAgent):
    """
    An agent that echoes messages back.

    This demonstrates:
    - Minimal agent implementation
    - @capability decorator for exposing functionality
    - Type-safe message handling
    - Automatic security enforcement
    """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._echo_count = 0  # Track number of echoes

    @capability("echo")
    async def echo(self, message: str, metadata: Dict[str, Any] = None) -> TaskResult:
        """
        Echo a message back to the caller.

        Security automatically enforced by SDK:
        - Caller's SPIFFE ID verified
        - OPA policy checked before this runs
        - Request/response logged for audit

        Args:
            message: The message to echo
            metadata: Optional metadata to include in response

        Returns:
            TaskResult with the echoed message
        """
        self._echo_count += 1

        # Get caller information from request context
        # This is automatically populated by the SDK
        caller_id = self.context.caller_spiffe_id

        # Log for observability
        self.logger.info(
            "Echo request received",
            extra={
                "caller": caller_id,
                "message_length": len(message),
                "echo_count": self._echo_count
            }
        )

        # Build response
        response_text = f"[Echo #{self._echo_count}] {message}"

        if metadata:
            response_text += f"\nMetadata: {metadata}"

        # Return as A2A TaskResult
        return TaskResult(
            status="completed",
            messages=[
                Message(
                    role="assistant",
                    parts=[TextPart(text=response_text)]
                )
            ],
            artifacts=[
                {
                    "type": "echo_stats",
                    "data": {
                        "original_message": message,
                        "echo_count": self._echo_count,
                        "caller": caller_id
                    }
                }
            ]
        )

    @capability("stats")
    async def get_stats(self) -> TaskResult:
        """
        Get agent statistics.

        Returns:
            TaskResult with agent stats
        """
        return TaskResult(
            status="completed",
            messages=[
                Message(
                    role="assistant",
                    parts=[
                        TextPart(
                            text=f"Echo Agent Statistics\n"
                                 f"Total echoes: {self._echo_count}\n"
                                 f"Agent ID: {self.spiffe_id}"
                        )
                    ]
                )
            ]
        )


async def main():
    """Run the echo agent."""
    # Load configuration from file
    # This includes identity, authorization, transport settings
    agent = EchoAgent.from_config("config/agent.yaml")

    # Start the agent server
    # This will:
    # 1. Connect to SPIRE to get SVID
    # 2. Start HTTPS server with mTLS
    # 3. Publish Agent Card at /.well-known/agent.json
    # 4. Listen for incoming tasks
    await agent.run()


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