Allow configuring headers and security checks.

Author: gspencergoog

packages/a2a_dart/lib/src/client/http_transport.dart

@@ -25,23 +25,32 @@ class HttpTransport implements Transport {
   /// The logger to use for logging.
   final Logger? log;
 
+  /// Default headers to include in all requests.
+  final Map<String, String> defaultHeaders;
+
   /// Creates an [HttpTransport].
   ///
   /// The [url] is the base URL of the A2A server.
   /// The [client] is an optional HTTP client to use for requests. If not
   /// provided, a new one will be created.
   /// The [log] is an optional logger.
-  HttpTransport({required this.url, http.Client? client, this.log})
-    : client = client ?? http.Client();
+  /// The [defaultHeaders] are optional headers to include in all requests.
+  HttpTransport({
+    required this.url,
+    http.Client? client,
+    this.log,
+    this.defaultHeaders = const {},
+  }) : client = client ?? http.Client();
 
   @override
   Future<Map<String, Object?>> get(
     String path, {
     Map<String, String> headers = const {},
   }) async {
     final uri = Uri.parse('$url$path');
-    log?.fine('Sending GET request to $uri');
-    final response = await client.get(uri, headers: headers);
+    final mergedHeaders = {...defaultHeaders, ...headers};
+    log?.fine('Sending GET request to $uri with headers: $mergedHeaders');
+    final response = await client.get(uri, headers: mergedHeaders);
     log?.fine('Received response from GET $uri: ${response.body}');
     if (response.statusCode != 200) {
       throw A2AException.http(
@@ -59,9 +68,13 @@ class HttpTransport implements Transport {
   }) async {
     final uri = Uri.parse('$url$path');
     log?.fine('Sending POST request to $uri with body: $request');
+    final mergedHeaders = {
+      'Content-Type': 'application/json',
+      ...defaultHeaders,
+    };
     final response = await client.post(
       uri,
-      headers: {'Content-Type': 'application/json'},
+      headers: mergedHeaders,
       body: jsonEncode(request),
     );
     log?.fine('Received response from POST $uri: ${response.body}');

packages/a2a_dart/lib/src/client/sse_transport.dart

@@ -27,16 +27,25 @@ class SseTransport extends HttpTransport {
   /// The [url] is the base URL of the A2A server. An optional [client] can be
   /// provided for testing or to customize the HTTP client. The [log] is an
   /// optional logger.
-  SseTransport({required super.url, super.client, super.log});
+  SseTransport({
+    required super.url,
+    super.client,
+    super.log,
+    super.defaultHeaders,
+  });
 
   @override
   Stream<Map<String, Object?>> sendStream(Map<String, Object?> request) async* {
     final uri = Uri.parse('$url/rpc');
     final body = jsonEncode(request);
     log?.fine('Sending SSE request to $uri with body: $body');
+    final mergedHeaders = {
+      'Content-Type': 'application/json',
+      'Accept': 'text/event-stream',
+      ...defaultHeaders,
+    };
     final httpRequest = http.Request('POST', uri)
-      ..headers['Content-Type'] = 'application/json'
-      ..headers['Accept'] = 'text/event-stream'
+      ..headers.addAll(mergedHeaders)
       ..body = body;
 
     try {

packages/a2a_dart/lib/src/server/create_task_handler.dart

@@ -26,6 +26,9 @@ class CreateTaskHandler implements RequestHandler {
   @override
   String get method => 'create_task';
 
+  @override
+  List<Map<String, List<String>>>? get securityRequirements => null;
+
   @override
   FutureOr<HandlerResult> handle(Map<String, Object?> params) async {
     if (!params.containsKey('message')) {

packages/a2a_dart/lib/src/server/delete_push_config_handler.dart

@@ -17,6 +17,9 @@ class DeletePushConfigHandler implements RequestHandler {
   @override
   String get method => 'tasks/pushNotificationConfig/delete';
 
+  @override
+  List<Map<String, List<String>>>? get securityRequirements => null;
+
   @override
   Future<HandlerResult> handle(Map<String, Object?> params) async {
     final taskId = params['id'] as String?;

packages/a2a_dart/lib/src/server/get_push_config_handler.dart

@@ -18,6 +18,9 @@ class GetPushConfigHandler implements RequestHandler {
   @override
   String get method => 'tasks/pushNotificationConfig/get';
 
+  @override
+  List<Map<String, List<String>>>? get securityRequirements => null;
+
   @override
   Future<HandlerResult> handle(Map<String, Object?> params) async {
     final taskId = params['id'] as String?;

packages/a2a_dart/lib/src/server/get_task_handler.dart

@@ -20,6 +20,9 @@ class GetTaskHandler extends RequestHandler {
   @override
   String get method => 'tasks/get';
 
+  @override
+  List<Map<String, List<String>>>? get securityRequirements => null;
+
   @override
   FutureOr<HandlerResult> handle(Map<String, Object?> params) async {
     final taskId = params['id'] as String?;

packages/a2a_dart/lib/src/server/in_memory_task_manager.dart

@@ -19,7 +19,6 @@ class InMemoryTaskManager implements TaskManager {
   final _uuid = const Uuid();
   final _pushConfigs = <String, Map<String, PushNotificationConfig>>{};
 
-  @override
   @override
   Future<Task> createTask([Message? message]) async {
     final taskId = _uuid.v4();
@@ -35,19 +34,16 @@ class InMemoryTaskManager implements TaskManager {
     return task;
   }
 
-  @override
   @override
   Future<Task?> getTask(String id) async => _tasks[id];
 
-  @override
   @override
   Future<void> updateTask(Task task) async {
     _tasks[task.id] = task.copyWith(
       lastUpdated: DateTime.now().millisecondsSinceEpoch,
     );
   }
 
-  @override
   @override
   Future<Task?> cancelTask(String taskId) async {
     final task = _tasks[taskId];
@@ -62,7 +58,6 @@ class InMemoryTaskManager implements TaskManager {
     return null;
   }
 
-  @override
   @override
   Future<ListTasksResult> listTasks(ListTasksParams params) async {
     var tasks = _tasks.values.toList();

packages/a2a_dart/lib/src/server/io_a2a_server.dart

@@ -8,7 +8,7 @@ import 'dart:io';
 
 import 'package:logging/logging.dart';
 import 'package:shelf/shelf.dart';
-import 'package:shelf/shelf_io.dart' as io;
+import 'package:shelf/shelf_io.dart' as shelf_io;
 import 'package:shelf_router/shelf_router.dart';
 
 import '../core/agent_card.dart';
@@ -44,6 +44,8 @@ class A2AServer {
   final int _requestedPort;
 
   final Map<String, RequestHandler> _handlers = {};
+  final TaskManager taskManager;
+  final Middleware? initialMiddleware;
 
   /// The public agent card for this server.
   ///
@@ -78,12 +80,13 @@ class A2AServer {
   /// ```
   A2AServer(
     List<RequestHandler> handlers,
-    TaskManager taskManager, {
+    this.taskManager, {
     this.host = 'localhost',
     int port = 0,
     Logger? logger,
     this.agentCard,
     this.extendedAgentCard,
+    this.initialMiddleware,
   }) : _requestedPort = port,
        _log = logger {
     for (final handler in handlers) {
@@ -128,6 +131,9 @@ class A2AServer {
       ..get('/.well-known/agent-card.json', _handleAgentCardRequest);
 
     var pipeline = const Pipeline();
+    if (initialMiddleware != null) {
+      pipeline = pipeline.addMiddleware(initialMiddleware!);
+    }
     if (_log != null) {
       pipeline = pipeline.addMiddleware(
         logRequests(
@@ -144,7 +150,7 @@ class A2AServer {
     final handler = pipeline.addHandler(router.call);
 
     _log?.info('Starting A2A server on $host:$_requestedPort...');
-    _server = await io.serve(handler, host, _requestedPort);
+    _server = await shelf_io.serve(handler, host, _requestedPort);
     _log?.info(
       'A2A server started on ${_server!.address.host}:${_server!.port}',
     );
@@ -169,19 +175,23 @@ class A2AServer {
 
   Future<Response> _handleRpcRequest(Request request) async {
     _log?.info('Received request: ${request.method} ${request.requestedUri}');
-    final body = await request.readAsString();
-    _log?.fine('Request body: $body');
 
     Object? id;
-    Map<String, Object?> json;
-    try {
-      json = jsonDecode(body) as Map<String, Object?>;
-      id = json['id'];
-    } on FormatException {
-      return _jsonRpcError(id: null, code: -32700, message: 'Parse error');
+    var json = request.context['a2a_body'] as Map<String, Object?>?;
+
+    if (json == null) {
+      // This should not happen if the security middleware ran first for /rpc requests
+      final body = await request.readAsString();
+      _log?.fine('Request body: $body');
+      try {
+        json = jsonDecode(body) as Map<String, Object?>;
+      } on FormatException {
+        return _jsonRpcError(id: null, code: -32700, message: 'Parse error');
+      }
     }
 
     try {
+      id = json['id'];
       final method = json['method'] as String?;
       final params = json['params'] as Map<String, Object?>?;
 
@@ -198,6 +208,54 @@ class A2AServer {
           responseCode: 404,
         );
       }
+      // Security Check
+      final securityRequirements = handler.securityRequirements;
+      if (securityRequirements != null && securityRequirements.isNotEmpty) {
+        final authContext =
+            request.context['a2a_auth'] as Map<String, dynamic>?;
+
+        if (authContext == null || authContext['isAuthenticated'] != true) {
+          return _jsonRpcError(
+            id: id,
+            code: -32002,
+            message: 'Unauthorized: Missing or failed authentication.',
+            responseCode: 401,
+          );
+        }
+
+        final authenticatedSchemes =
+            authContext['schemes'] as Map<String, List<String>>? ?? {};
+        var authorized = false;
+        for (final requirement in securityRequirements) {
+          var requirementMet = true;
+          for (final schemeName in requirement.keys) {
+            final requiredScopes = requirement[schemeName]!;
+            if (!authenticatedSchemes.containsKey(schemeName)) {
+              requirementMet = false;
+              break;
+            }
+            final grantedScopes = authenticatedSchemes[schemeName]!;
+            if (!requiredScopes.every(grantedScopes.contains)) {
+              requirementMet = false;
+              break;
+            }
+          }
+          if (requirementMet) {
+            authorized = true;
+            break;
+          }
+        }
+
+        if (!authorized) {
+          return _jsonRpcError(
+            id: id,
+            code: -32002,
+            message: 'Unauthorized: Insufficient permissions for this method.',
+            responseCode: 401,
+          );
+        }
+      }
+
       return _executeHandler(handler, params, id);
     } on Exception catch (exception, stackTrace) {
       _log?.severe('Unhandled server exception', exception, stackTrace);

packages/a2a_dart/lib/src/server/list_push_configs_handler.dart

@@ -17,6 +17,9 @@ class ListPushConfigsHandler implements RequestHandler {
   @override
   String get method => 'tasks/pushNotificationConfig/list';
 
+  @override
+  List<Map<String, List<String>>>? get securityRequirements => null;
+
   @override
   Future<HandlerResult> handle(Map<String, Object?> params) async {
     final taskId = params['id'] as String?;

packages/a2a_dart/lib/src/server/list_tasks_handler.dart

@@ -20,6 +20,9 @@ class ListTasksHandler extends RequestHandler {
   @override
   String get method => 'tasks/list';
 
+  @override
+  List<Map<String, List<String>>>? get securityRequirements => null;
+
   @override
   FutureOr<HandlerResult> handle(Map<String, Object?> params) async {
     final listTasksParams = ListTasksParams.fromJson(params);