a2a_server/
middleware.rs

1// Copyright AGNTCY Contributors (https://github.com/agntcy)
2// Copyright A2A Contributors (https://github.com/a2aproject)
3// SPDX-License-Identifier: Apache-2.0
4use a2a::A2AError;
5use async_trait::async_trait;
6use axum::http::HeaderMap;
7use serde_json::Value;
8use std::collections::HashMap;
9
10/// Service parameters — metadata from HTTP headers or gRPC metadata.
11pub type ServiceParams = HashMap<String, Vec<String>>;
12
13/// Build [`ServiceParams`] from an axum [`HeaderMap`].
14///
15/// Header names are normalized to lowercase (axum does this for us), values
16/// that aren't valid ASCII are silently skipped, and multi-valued headers
17/// produce a `Vec<String>` in insertion order.
18pub(crate) fn extract_service_params(headers: &HeaderMap) -> ServiceParams {
19    let mut params = ServiceParams::new();
20    for (name, value) in headers {
21        if let Ok(v) = value.to_str() {
22            params
23                .entry(name.as_str().to_owned())
24                .or_default()
25                .push(v.to_owned());
26        }
27    }
28    params
29}
30
31/// Authenticated user information.
32#[derive(Debug, Clone)]
33pub struct User {
34    pub name: String,
35    pub authenticated: bool,
36    pub attributes: HashMap<String, Value>,
37}
38
39impl User {
40    pub fn authenticated(name: impl Into<String>) -> Self {
41        User {
42            name: name.into(),
43            authenticated: true,
44            attributes: HashMap::new(),
45        }
46    }
47}
48
49/// Context for a single request being processed.
50pub struct CallContext {
51    pub method: String,
52    pub service_params: ServiceParams,
53    pub tenant: Option<String>,
54    pub user: Option<User>,
55}
56
57impl CallContext {
58    pub fn new(method: impl Into<String>, params: ServiceParams) -> Self {
59        CallContext {
60            method: method.into(),
61            service_params: params,
62            tenant: None,
63            user: None,
64        }
65    }
66}
67
68/// Server-side interceptor for modifying requests and responses.
69#[async_trait]
70pub trait CallInterceptor: Send + Sync + 'static {
71    async fn before(&self, ctx: &mut CallContext, request: &Value) -> Result<(), A2AError> {
72        let _ = (ctx, request);
73        Ok(())
74    }
75
76    async fn after(
77        &self,
78        ctx: &CallContext,
79        result: &Result<Value, A2AError>,
80    ) -> Result<(), A2AError> {
81        let _ = (ctx, result);
82        Ok(())
83    }
84}
85
86/// Wraps a [`RequestHandler`](crate::RequestHandler) with interceptors.
87pub struct InterceptedHandler<H> {
88    pub handler: H,
89    pub interceptors: Vec<Box<dyn CallInterceptor>>,
90}
91
92impl<H> InterceptedHandler<H> {
93    pub fn new(handler: H) -> Self {
94        InterceptedHandler {
95            handler,
96            interceptors: Vec::new(),
97        }
98    }
99
100    pub fn with_interceptor(mut self, interceptor: impl CallInterceptor) -> Self {
101        self.interceptors.push(Box::new(interceptor));
102        self
103    }
104}
105
106/// Server-side logging interceptor.
107pub struct LoggingInterceptor;
108
109#[async_trait]
110impl CallInterceptor for LoggingInterceptor {
111    async fn before(&self, ctx: &mut CallContext, _request: &Value) -> Result<(), A2AError> {
112        tracing::info!(method = %ctx.method, "A2A server request");
113        Ok(())
114    }
115
116    async fn after(
117        &self,
118        ctx: &CallContext,
119        result: &Result<Value, A2AError>,
120    ) -> Result<(), A2AError> {
121        match result {
122            Ok(_) => tracing::info!(method = %ctx.method, "A2A server response"),
123            Err(e) => tracing::warn!(method = %ctx.method, error = %e, "A2A server error"),
124        }
125        Ok(())
126    }
127}
128
129#[cfg(test)]
130mod tests {
131    use super::*;
132    use axum::http::HeaderValue;
133
134    #[test]
135    fn test_extract_service_params_empty() {
136        let headers = HeaderMap::new();
137        let params = extract_service_params(&headers);
138        assert!(params.is_empty());
139    }
140
141    #[test]
142    fn test_extract_service_params_basic() {
143        let mut headers = HeaderMap::new();
144        headers.insert(
145            "authorization",
146            HeaderValue::from_static("Bearer token-abc"),
147        );
148        headers.insert("x-tenant-id", HeaderValue::from_static("acme"));
149
150        let params = extract_service_params(&headers);
151
152        assert_eq!(
153            params.get("authorization"),
154            Some(&vec!["Bearer token-abc".to_string()])
155        );
156        assert_eq!(params.get("x-tenant-id"), Some(&vec!["acme".to_string()]));
157    }
158
159    #[test]
160    fn test_extract_service_params_multi_value() {
161        let mut headers = HeaderMap::new();
162        headers.append("x-forwarded-for", HeaderValue::from_static("10.0.0.1"));
163        headers.append("x-forwarded-for", HeaderValue::from_static("10.0.0.2"));
164
165        let params = extract_service_params(&headers);
166
167        assert_eq!(
168            params.get("x-forwarded-for"),
169            Some(&vec!["10.0.0.1".to_string(), "10.0.0.2".to_string()]),
170        );
171    }
172
173    #[test]
174    fn test_extract_service_params_lowercases_names() {
175        let mut headers = HeaderMap::new();
176        // HeaderMap normalizes the name to lowercase regardless of casing here.
177        headers.insert("Authorization", HeaderValue::from_static("Bearer cased"));
178        let params = extract_service_params(&headers);
179        assert_eq!(
180            params.get("authorization"),
181            Some(&vec!["Bearer cased".to_string()])
182        );
183        assert!(!params.contains_key("Authorization"));
184    }
185
186    #[test]
187    fn test_extract_service_params_skips_non_ascii_values() {
188        let mut headers = HeaderMap::new();
189        headers.insert("x-ascii", HeaderValue::from_static("ok"));
190        headers.insert(
191            "x-binary",
192            HeaderValue::from_bytes(&[0xff, 0xfe, 0xfd]).unwrap(),
193        );
194
195        let params = extract_service_params(&headers);
196
197        assert_eq!(params.get("x-ascii"), Some(&vec!["ok".to_string()]));
198        assert!(
199            !params.contains_key("x-binary"),
200            "non-ASCII value should be skipped, not produce an empty entry"
201        );
202    }
203
204    #[test]
205    fn test_user_authenticated() {
206        let user = User::authenticated("alice");
207        assert_eq!(user.name, "alice");
208        assert!(user.authenticated);
209        assert!(user.attributes.is_empty());
210    }
211
212    #[test]
213    fn test_call_context_new() {
214        let params = ServiceParams::new();
215        let ctx = CallContext::new(a2a::jsonrpc::methods::SEND_MESSAGE, params);
216        assert_eq!(ctx.method, a2a::jsonrpc::methods::SEND_MESSAGE);
217        assert!(ctx.tenant.is_none());
218        assert!(ctx.user.is_none());
219    }
220
221    #[test]
222    fn test_intercepted_handler_new() {
223        let handler = "dummy";
224        let ih = InterceptedHandler::new(handler);
225        assert_eq!(ih.handler, "dummy");
226        assert!(ih.interceptors.is_empty());
227    }
228
229    #[test]
230    fn test_intercepted_handler_with_interceptor() {
231        let handler = "dummy";
232        let ih = InterceptedHandler::new(handler).with_interceptor(LoggingInterceptor);
233        assert_eq!(ih.interceptors.len(), 1);
234    }
235
236    struct NoopInterceptor;
237
238    #[async_trait]
239    impl CallInterceptor for NoopInterceptor {}
240
241    #[tokio::test]
242    async fn test_noop_interceptor_before() {
243        let interceptor = NoopInterceptor;
244        let mut ctx = CallContext::new("test", ServiceParams::new());
245        let result = interceptor.before(&mut ctx, &Value::Null).await;
246        assert!(result.is_ok());
247    }
248
249    #[tokio::test]
250    async fn test_noop_interceptor_after() {
251        let interceptor = NoopInterceptor;
252        let ctx = CallContext::new("test", ServiceParams::new());
253        let result = interceptor.after(&ctx, &Ok(Value::Null)).await;
254        assert!(result.is_ok());
255    }
256
257    #[tokio::test]
258    async fn test_logging_interceptor_before() {
259        let interceptor = LoggingInterceptor;
260        let mut ctx = CallContext::new(a2a::jsonrpc::methods::SEND_MESSAGE, ServiceParams::new());
261        let result = interceptor.before(&mut ctx, &Value::Null).await;
262        assert!(result.is_ok());
263    }
264
265    #[tokio::test]
266    async fn test_logging_interceptor_after_ok() {
267        let interceptor = LoggingInterceptor;
268        let ctx = CallContext::new(a2a::jsonrpc::methods::SEND_MESSAGE, ServiceParams::new());
269        let result = interceptor.after(&ctx, &Ok(Value::Null)).await;
270        assert!(result.is_ok());
271    }
272
273    #[tokio::test]
274    async fn test_logging_interceptor_after_err() {
275        let interceptor = LoggingInterceptor;
276        let ctx = CallContext::new(a2a::jsonrpc::methods::SEND_MESSAGE, ServiceParams::new());
277        let err: Result<Value, A2AError> = Err(A2AError::internal("boom"));
278        let result = interceptor.after(&ctx, &err).await;
279        assert!(result.is_ok());
280    }
281}
a2a_server/middleware.rs

a2a_server/
middleware.rs
