a2a_server/
middleware.rs

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

a2a_server/
middleware.rs
