Skip to main content

a2a_client/
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;
6
7use crate::transport::ServiceParams;
8
9/// Interceptor for modifying requests and responses at the client level.
10///
11/// Interceptors are called in order for `before`, and in reverse order for `after`.
12#[async_trait]
13pub trait CallInterceptor: Send + Sync {
14    /// Called before sending a request. Can modify params (e.g., add auth headers).
15    async fn before(&self, method: &str, params: &mut ServiceParams) -> Result<(), A2AError> {
16        let _ = (method, params);
17        Ok(())
18    }
19
20    /// Called after receiving a response.
21    async fn after(&self, method: &str, result: &Result<(), A2AError>) -> Result<(), A2AError> {
22        let _ = (method, result);
23        Ok(())
24    }
25}
26
27/// Logging interceptor using `tracing`.
28pub struct LoggingInterceptor;
29
30#[async_trait]
31impl CallInterceptor for LoggingInterceptor {
32    async fn before(&self, method: &str, _params: &mut ServiceParams) -> Result<(), A2AError> {
33        tracing::info!(method = method, "A2A client request");
34        Ok(())
35    }
36
37    async fn after(&self, method: &str, result: &Result<(), A2AError>) -> Result<(), A2AError> {
38        match result {
39            Ok(()) => tracing::info!(method = method, "A2A client response"),
40            Err(e) => tracing::warn!(method = method, error = %e, "A2A client error"),
41        }
42        Ok(())
43    }
44}
45
46#[cfg(test)]
47mod tests {
48    use super::*;
49
50    struct NoopInterceptor;
51
52    #[async_trait]
53    impl CallInterceptor for NoopInterceptor {}
54
55    #[tokio::test]
56    async fn test_default_before() {
57        let interceptor = NoopInterceptor;
58        let mut params = ServiceParams::new();
59        let result = interceptor.before("test", &mut params).await;
60        assert!(result.is_ok());
61    }
62
63    #[tokio::test]
64    async fn test_default_after() {
65        let interceptor = NoopInterceptor;
66        let result = interceptor.after("test", &Ok(())).await;
67        assert!(result.is_ok());
68    }
69
70    #[tokio::test]
71    async fn test_default_after_with_error() {
72        let interceptor = NoopInterceptor;
73        let err = Err(A2AError::internal("fail"));
74        let result = interceptor.after("test", &err).await;
75        assert!(result.is_ok());
76    }
77
78    #[tokio::test]
79    async fn test_logging_interceptor_before() {
80        let interceptor = LoggingInterceptor;
81        let mut params = ServiceParams::new();
82        let result = interceptor
83            .before(a2a::jsonrpc::methods::SEND_MESSAGE, &mut params)
84            .await;
85        assert!(result.is_ok());
86    }
87
88    #[tokio::test]
89    async fn test_logging_interceptor_after_ok() {
90        let interceptor = LoggingInterceptor;
91        let result = interceptor
92            .after(a2a::jsonrpc::methods::SEND_MESSAGE, &Ok(()))
93            .await;
94        assert!(result.is_ok());
95    }
96
97    #[tokio::test]
98    async fn test_logging_interceptor_after_err() {
99        let interceptor = LoggingInterceptor;
100        let err = Err(A2AError::internal("boom"));
101        let result = interceptor
102            .after(a2a::jsonrpc::methods::SEND_MESSAGE, &err)
103            .await;
104        assert!(result.is_ok());
105    }
106}