From 3a436c95383e6ff9f83af23b3a65db39773fa611 Mon Sep 17 00:00:00 2001
From: Joshua Medvinsky <joshuam33@Joshuas-MacBook-Pro.local>
Date: Sat, 23 May 2026 14:38:11 -0700
Subject: [PATCH] fix: add optional API key authentication to all mutating HTTP
 endpoints

Introduces an api_key parameter to create_flask_app(). When set, every
POST route (/a2a, /stream, /tasks/send, /tasks/cancel, /) requires an
Authorization: Bearer <token> header and returns 401 otherwise.

Read-only GET routes (agent card, health, metadata) are intentionally
left unauthenticated so service-discovery clients can operate without
credentials, consistent with the A2A spec.
---
 python_a2a/server/a2a_server.py | 14 +++++++++++++-
 python_a2a/server/http.py       | 33 +++++++++++++++++++++++++++------
 2 files changed, 40 insertions(+), 7 deletions(-)

diff --git a/python_a2a/server/a2a_server.py b/python_a2a/server/a2a_server.py
index 5bf772b..a3571b5 100644
--- a/python_a2a/server/a2a_server.py
+++ b/python_a2a/server/a2a_server.py
@@ -341,6 +341,9 @@ def handle_task(self, task):
     
     def setup_routes(self, app):
         """Setup Flask routes for A2A endpoints"""
+        # Retrieve the auth checker injected by create_flask_app, or a no-op fallback
+        _check_auth = getattr(app, "check_api_key", lambda: None)
+
         # Root endpoint for both GET and POST
         @app.route("/", methods=["GET"])
         def a2a_root_get():
@@ -352,10 +355,13 @@ def a2a_root_get():
                 "protocol": "a2a",
                 "capabilities": self.agent_card.capabilities
             })
-            
+
         @app.route("/", methods=["POST"])
         def a2a_root_post():
             """Root endpoint for A2A (POST) - handle message in appropriate format"""
+            auth_error = _check_auth()
+            if auth_error:
+                return auth_error
             try:
                 data = request.json
                 
@@ -442,6 +448,9 @@ def well_known_agent_card():
         @app.route("/a2a/tasks/send", methods=["POST"])
         def a2a_tasks_send():
             """Handle POST request to create or update a task"""
+            auth_error = _check_auth()
+            if auth_error:
+                return auth_error
             try:
                 # Parse JSON data
                 request_data = request.json
@@ -596,6 +605,9 @@ def tasks_get():
         @app.route("/a2a/tasks/cancel", methods=["POST"])
         def a2a_tasks_cancel():
             """Handle POST request to cancel a task"""
+            auth_error = _check_auth()
+            if auth_error:
+                return auth_error
             try:
                 # Parse JSON data
                 request_data = request.json
diff --git a/python_a2a/server/http.py b/python_a2a/server/http.py
index 24644df..92577f8 100644
--- a/python_a2a/server/http.py
+++ b/python_a2a/server/http.py
@@ -23,16 +23,19 @@
 from .ui_templates import AGENT_INDEX_HTML, JSON_HTML_TEMPLATE
 
 
-def create_flask_app(agent: BaseA2AServer) -> Flask:
+def create_flask_app(agent: BaseA2AServer, api_key: Optional[str] = None) -> Flask:
     """
     Create a Flask application that serves an A2A agent
-    
+
     Args:
         agent: The A2A agent server
-        
+        api_key: Optional bearer token. When set, all POST/mutating routes require
+                 an ``Authorization: Bearer <api_key>`` header; requests without a
+                 valid token receive a 401 response.
+
     Returns:
         A Flask application
-        
+
     Raises:
         A2AImportError: If Flask is not installed
     """
@@ -41,8 +44,17 @@ def create_flask_app(agent: BaseA2AServer) -> Flask:
             "Flask is not installed. "
             "Install it with 'pip install flask'"
         )
-    
+
     app = Flask(__name__)
+
+    def _check_api_key():
+        """Return a 401 Response if api_key is configured and not matched, else None."""
+        if not api_key:
+            return None
+        auth_header = request.headers.get("Authorization", "")
+        if auth_header == f"Bearer {api_key}":
+            return None
+        return jsonify({"error": "Unauthorized", "message": "Valid Authorization: Bearer <token> header required"}), 401
     
     # Allow CORS for all routes
     @app.after_request
@@ -187,10 +199,13 @@ def enhanced_wellknown_agent_json():
     def handle_streaming_request():
         """
         Handle streaming requests.
-        
+
         This endpoint enables Server-Sent Events (SSE) streaming from the agent.
         It uses the agent's stream_response method if it implements it.
         """
+        auth_error = _check_api_key()
+        if auth_error:
+            return auth_error
         try:
             # CORS for streaming - important for browser compatibility
             if request.method == 'OPTIONS':
@@ -382,6 +397,9 @@ async def process_stream():
             # Return error response for any other exception
             return jsonify({"error": str(e)}), 500
     
+    # Store the auth checker on the app so setup_routes implementations can use it
+    app.check_api_key = _check_api_key  # type: ignore[attr-defined]
+
     # Only AFTER registering our enhanced routes, set up the agent's routes
     if hasattr(agent, 'setup_routes'):
         agent.setup_routes(app)
@@ -390,6 +408,9 @@ async def process_stream():
     @app.route("/a2a", methods=["POST"])
     def handle_a2a_request() -> Union[Response, tuple]:
         """Handle A2A protocol requests"""
+        auth_error = _check_api_key()
+        if auth_error:
+            return auth_error
         try:
             data = request.json