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 serde_json::Value;
6use std::collections::HashMap;
7
8/// Service parameters — metadata from HTTP headers or gRPC metadata.
9pub type ServiceParams = HashMap<String, Vec<String>>;
10
11/// Authenticated user information.
12#[derive(Debug, Clone)]
13pub struct User {
14    pub name: String,
15    pub authenticated: bool,
16    pub attributes: HashMap<String, Value>,
17}
18
19impl User {
20    pub fn authenticated(name: impl Into<String>) -> Self {
21        User {
22            name: name.into(),
23            authenticated: true,
24            attributes: HashMap::new(),
25        }
26    }
27}
28
29/// Context for a single request being processed.
30pub struct CallContext {
31    pub method: String,
32    pub service_params: ServiceParams,
33    pub tenant: Option<String>,
34    pub user: Option<User>,
35}
36
37impl CallContext {
38    pub fn new(method: impl Into<String>, params: ServiceParams) -> Self {
39        CallContext {
40            method: method.into(),
41            service_params: params,
42            tenant: None,
43            user: None,
44        }
45    }
46}
47
48/// Server-side interceptor for modifying requests and responses.
49#[async_trait]
50pub trait CallInterceptor: Send + Sync + 'static {
51    async fn before(&self, ctx: &mut CallContext, request: &Value) -> Result<(), A2AError> {
52        let _ = (ctx, request);
53        Ok(())
54    }
55
56    async fn after(
57        &self,
58        ctx: &CallContext,
59        result: &Result<Value, A2AError>,
60    ) -> Result<(), A2AError> {
61        let _ = (ctx, result);
62        Ok(())
63    }
64}
65
66/// Wraps a [`RequestHandler`](crate::RequestHandler) with interceptors.
67pub struct InterceptedHandler<H> {
68    pub handler: H,
69    pub interceptors: Vec<Box<dyn CallInterceptor>>,
70}
71
72impl<H> InterceptedHandler<H> {
73    pub fn new(handler: H) -> Self {
74        InterceptedHandler {
75            handler,
76            interceptors: Vec::new(),
77        }
78    }
79
80    pub fn with_interceptor(mut self, interceptor: impl CallInterceptor) -> Self {
81        self.interceptors.push(Box::new(interceptor));
82        self
83    }
84}
85
86/// Server-side logging interceptor.
87pub struct LoggingInterceptor;
88
89#[async_trait]
90impl CallInterceptor for LoggingInterceptor {
91    async fn before(&self, ctx: &mut CallContext, _request: &Value) -> Result<(), A2AError> {
92        tracing::info!(method = %ctx.method, "A2A server request");
93        Ok(())
94    }
95
96    async fn after(
97        &self,
98        ctx: &CallContext,
99        result: &Result<Value, A2AError>,
100    ) -> Result<(), A2AError> {
101        match result {
102            Ok(_) => tracing::info!(method = %ctx.method, "A2A server response"),
103            Err(e) => tracing::warn!(method = %ctx.method, error = %e, "A2A server error"),
104        }
105        Ok(())
106    }
107}
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112
113    #[test]
114    fn test_user_authenticated() {
115        let user = User::authenticated("alice");
116        assert_eq!(user.name, "alice");
117        assert!(user.authenticated);
118        assert!(user.attributes.is_empty());
119    }
120
121    #[test]
122    fn test_call_context_new() {
123        let params = ServiceParams::new();
124        let ctx = CallContext::new(a2a::jsonrpc::methods::SEND_MESSAGE, params);
125        assert_eq!(ctx.method, a2a::jsonrpc::methods::SEND_MESSAGE);
126        assert!(ctx.tenant.is_none());
127        assert!(ctx.user.is_none());
128    }
129
130    #[test]
131    fn test_intercepted_handler_new() {
132        let handler = "dummy";
133        let ih = InterceptedHandler::new(handler);
134        assert_eq!(ih.handler, "dummy");
135        assert!(ih.interceptors.is_empty());
136    }
137
138    #[test]
139    fn test_intercepted_handler_with_interceptor() {
140        let handler = "dummy";
141        let ih = InterceptedHandler::new(handler).with_interceptor(LoggingInterceptor);
142        assert_eq!(ih.interceptors.len(), 1);
143    }
144
145    struct NoopInterceptor;
146
147    #[async_trait]
148    impl CallInterceptor for NoopInterceptor {}
149
150    #[tokio::test]
151    async fn test_noop_interceptor_before() {
152        let interceptor = NoopInterceptor;
153        let mut ctx = CallContext::new("test", ServiceParams::new());
154        let result = interceptor.before(&mut ctx, &Value::Null).await;
155        assert!(result.is_ok());
156    }
157
158    #[tokio::test]
159    async fn test_noop_interceptor_after() {
160        let interceptor = NoopInterceptor;
161        let ctx = CallContext::new("test", ServiceParams::new());
162        let result = interceptor.after(&ctx, &Ok(Value::Null)).await;
163        assert!(result.is_ok());
164    }
165
166    #[tokio::test]
167    async fn test_logging_interceptor_before() {
168        let interceptor = LoggingInterceptor;
169        let mut ctx = CallContext::new(a2a::jsonrpc::methods::SEND_MESSAGE, ServiceParams::new());
170        let result = interceptor.before(&mut ctx, &Value::Null).await;
171        assert!(result.is_ok());
172    }
173
174    #[tokio::test]
175    async fn test_logging_interceptor_after_ok() {
176        let interceptor = LoggingInterceptor;
177        let ctx = CallContext::new(a2a::jsonrpc::methods::SEND_MESSAGE, ServiceParams::new());
178        let result = interceptor.after(&ctx, &Ok(Value::Null)).await;
179        assert!(result.is_ok());
180    }
181
182    #[tokio::test]
183    async fn test_logging_interceptor_after_err() {
184        let interceptor = LoggingInterceptor;
185        let ctx = CallContext::new(a2a::jsonrpc::methods::SEND_MESSAGE, ServiceParams::new());
186        let err: Result<Value, A2AError> = Err(A2AError::internal("boom"));
187        let result = interceptor.after(&ctx, &err).await;
188        assert!(result.is_ok());
189    }
190}
a2a_server/middleware.rs

a2a_server/
middleware.rs
