a2a_protocol_server/dispatch/grpc/
native.rs1use std::pin::Pin;
14use std::sync::Arc;
15
16use a2a_protocol_types::proto as apb;
17use a2a_protocol_types::proto::convert::ConvertError;
18use tokio::sync::mpsc;
19use tokio_stream::wrappers::ReceiverStream;
20use tonic::{Request, Response, Status};
21
22use super::helpers::{server_error_to_status, validated_metadata};
23use super::pb::a2a_service_server::A2aService;
24use super::GrpcConfig;
25use crate::handler::{RequestHandler, SendMessageResult};
26
27type NativeStream =
29 Pin<Box<dyn tokio_stream::Stream<Item = Result<apb::StreamResponse, Status>> + Send + 'static>>;
30
31#[allow(clippy::needless_pass_by_value)]
33fn bad_request(err: ConvertError) -> Status {
34 Status::invalid_argument(err.to_string())
35}
36
37#[allow(clippy::needless_pass_by_value)]
40fn bad_response(err: ConvertError) -> Status {
41 Status::internal(format!("response not representable in protobuf: {err}"))
42}
43
44fn send_result_to_stream(
46 resp: a2a_protocol_types::responses::SendMessageResponse,
47) -> Result<apb::StreamResponse, ConvertError> {
48 let payload = match resp {
49 a2a_protocol_types::responses::SendMessageResponse::Task(t) => {
50 apb::stream_response::Payload::Task(t.try_into()?)
51 }
52 a2a_protocol_types::responses::SendMessageResponse::Message(m) => {
53 apb::stream_response::Payload::Message(m.try_into()?)
54 }
55 other => {
56 return Err(ConvertError {
57 field: "sendMessageResponse.payload",
58 reason: format!("unsupported response variant: {other:?}"),
59 })
60 }
61 };
62 Ok(apb::StreamResponse {
63 payload: Some(payload),
64 })
65}
66
67fn reader_to_native_stream(
69 mut reader: crate::streaming::InMemoryQueueReader,
70 capacity: usize,
71) -> NativeStream {
72 use crate::streaming::EventQueueReader;
73 let (tx, rx) = mpsc::channel(capacity);
74 tokio::spawn(async move {
75 loop {
76 match reader.read().await {
77 Some(Ok(event)) => {
78 let item = apb::StreamResponse::try_from(event).map_err(bad_response);
79 let is_err = item.is_err();
80 if tx.send(item).await.is_err() || is_err {
81 break;
82 }
83 }
84 Some(Err(_)) => {
85 let _ = tx.send(Err(Status::internal("event queue error"))).await;
86 break;
87 }
88 None => break,
89 }
90 }
91 });
92 Box::pin(ReceiverStream::new(rx))
93}
94
95pub struct A2aServiceImpl {
100 pub(super) handler: Arc<RequestHandler>,
101 pub(super) config: GrpcConfig,
102}
103
104#[tonic::async_trait]
105impl A2aService for A2aServiceImpl {
106 async fn send_message(
109 &self,
110 request: Request<apb::SendMessageRequest>,
111 ) -> Result<Response<apb::SendMessageResponse>, Status> {
112 let headers = validated_metadata(request.metadata())?;
113 let params: a2a_protocol_types::params::MessageSendParams =
114 request.into_inner().try_into().map_err(bad_request)?;
115 match self
116 .handler
117 .on_send_message(params, false, Some(&headers))
118 .await
119 {
120 Ok(SendMessageResult::Response(resp)) => {
121 Ok(Response::new(resp.try_into().map_err(bad_response)?))
122 }
123 Ok(SendMessageResult::Stream(_)) => Err(Status::internal(
124 "unexpected stream response for unary call",
125 )),
126 Err(e) => Err(server_error_to_status(&e)),
127 }
128 }
129
130 type SendStreamingMessageStream = NativeStream;
131
132 async fn send_streaming_message(
133 &self,
134 request: Request<apb::SendMessageRequest>,
135 ) -> Result<Response<Self::SendStreamingMessageStream>, Status> {
136 let headers = validated_metadata(request.metadata())?;
137 let params: a2a_protocol_types::params::MessageSendParams =
138 request.into_inner().try_into().map_err(bad_request)?;
139 match self
140 .handler
141 .on_send_message(params, true, Some(&headers))
142 .await
143 {
144 Ok(SendMessageResult::Stream(reader)) => Ok(Response::new(reader_to_native_stream(
145 reader,
146 self.config.stream_channel_capacity,
147 ))),
148 Ok(SendMessageResult::Response(resp)) => {
149 let payload = send_result_to_stream(resp).map_err(bad_response)?;
151 let stream = Box::pin(tokio_stream::once(Ok(payload)));
152 Ok(Response::new(stream as NativeStream))
153 }
154 Err(e) => Err(server_error_to_status(&e)),
155 }
156 }
157
158 async fn get_task(
161 &self,
162 request: Request<apb::GetTaskRequest>,
163 ) -> Result<Response<apb::Task>, Status> {
164 let headers = validated_metadata(request.metadata())?;
165 let params: a2a_protocol_types::params::TaskQueryParams =
166 request.into_inner().try_into().map_err(bad_request)?;
167 match self.handler.on_get_task(params, Some(&headers)).await {
168 Ok(task) => Ok(Response::new(task.try_into().map_err(bad_response)?)),
169 Err(e) => Err(server_error_to_status(&e)),
170 }
171 }
172
173 async fn list_tasks(
174 &self,
175 request: Request<apb::ListTasksRequest>,
176 ) -> Result<Response<apb::ListTasksResponse>, Status> {
177 let headers = validated_metadata(request.metadata())?;
178 let params: a2a_protocol_types::params::ListTasksParams =
179 request.into_inner().try_into().map_err(bad_request)?;
180 match self.handler.on_list_tasks(params, Some(&headers)).await {
181 Ok(resp) => Ok(Response::new(resp.try_into().map_err(bad_response)?)),
182 Err(e) => Err(server_error_to_status(&e)),
183 }
184 }
185
186 async fn cancel_task(
187 &self,
188 request: Request<apb::CancelTaskRequest>,
189 ) -> Result<Response<apb::Task>, Status> {
190 let headers = validated_metadata(request.metadata())?;
191 let params: a2a_protocol_types::params::CancelTaskParams =
192 request.into_inner().try_into().map_err(bad_request)?;
193 match self.handler.on_cancel_task(params, Some(&headers)).await {
194 Ok(task) => Ok(Response::new(task.try_into().map_err(bad_response)?)),
195 Err(e) => Err(server_error_to_status(&e)),
196 }
197 }
198
199 type SubscribeToTaskStream = NativeStream;
200
201 async fn subscribe_to_task(
202 &self,
203 request: Request<apb::SubscribeToTaskRequest>,
204 ) -> Result<Response<Self::SubscribeToTaskStream>, Status> {
205 let headers = validated_metadata(request.metadata())?;
206 let params: a2a_protocol_types::params::TaskIdParams = request.into_inner().into();
207 match self.handler.on_resubscribe(params, Some(&headers)).await {
208 Ok(reader) => Ok(Response::new(reader_to_native_stream(
209 reader,
210 self.config.stream_channel_capacity,
211 ))),
212 Err(e) => Err(server_error_to_status(&e)),
213 }
214 }
215
216 async fn create_task_push_notification_config(
219 &self,
220 request: Request<apb::TaskPushNotificationConfig>,
221 ) -> Result<Response<apb::TaskPushNotificationConfig>, Status> {
222 let headers = validated_metadata(request.metadata())?;
223 let config: a2a_protocol_types::push::TaskPushNotificationConfig =
224 request.into_inner().into();
225 match self
226 .handler
227 .on_set_push_config(config, Some(&headers))
228 .await
229 {
230 Ok(cfg) => Ok(Response::new(cfg.into())),
231 Err(e) => Err(server_error_to_status(&e)),
232 }
233 }
234
235 async fn get_task_push_notification_config(
236 &self,
237 request: Request<apb::GetTaskPushNotificationConfigRequest>,
238 ) -> Result<Response<apb::TaskPushNotificationConfig>, Status> {
239 let headers = validated_metadata(request.metadata())?;
240 let params: a2a_protocol_types::params::GetPushConfigParams = request.into_inner().into();
241 match self
242 .handler
243 .on_get_push_config(params, Some(&headers))
244 .await
245 {
246 Ok(cfg) => Ok(Response::new(cfg.into())),
247 Err(e) => Err(server_error_to_status(&e)),
248 }
249 }
250
251 async fn list_task_push_notification_configs(
252 &self,
253 request: Request<apb::ListTaskPushNotificationConfigsRequest>,
254 ) -> Result<Response<apb::ListTaskPushNotificationConfigsResponse>, Status> {
255 let headers = validated_metadata(request.metadata())?;
256 let params: a2a_protocol_types::params::ListPushConfigsParams =
257 request.into_inner().try_into().map_err(bad_request)?;
258 match self
259 .handler
260 .on_list_push_configs(¶ms.task_id, params.tenant.as_deref(), Some(&headers))
261 .await
262 {
263 Ok(configs) => Ok(Response::new(
264 apb::ListTaskPushNotificationConfigsResponse {
265 configs: configs.into_iter().map(Into::into).collect(),
266 next_page_token: String::new(),
267 },
268 )),
269 Err(e) => Err(server_error_to_status(&e)),
270 }
271 }
272
273 async fn delete_task_push_notification_config(
274 &self,
275 request: Request<apb::DeleteTaskPushNotificationConfigRequest>,
276 ) -> Result<Response<()>, Status> {
277 let headers = validated_metadata(request.metadata())?;
278 let params: a2a_protocol_types::params::DeletePushConfigParams =
279 request.into_inner().into();
280 match self
281 .handler
282 .on_delete_push_config(params, Some(&headers))
283 .await
284 {
285 Ok(()) => Ok(Response::new(())),
286 Err(e) => Err(server_error_to_status(&e)),
287 }
288 }
289
290 async fn get_extended_agent_card(
293 &self,
294 request: Request<apb::GetExtendedAgentCardRequest>,
295 ) -> Result<Response<apb::AgentCard>, Status> {
296 let headers = validated_metadata(request.metadata())?;
297 match self
298 .handler
299 .on_get_extended_agent_card(Some(&headers))
300 .await
301 {
302 Ok(card) => Ok(Response::new(card.try_into().map_err(bad_response)?)),
303 Err(e) => Err(server_error_to_status(&e)),
304 }
305 }
306}
307
308#[cfg(test)]
309mod tests {
310 use super::*;
311 use a2a_protocol_types::message::{Message, MessageId, MessageRole, Part, PartContent};
312 use a2a_protocol_types::responses::SendMessageResponse;
313 use a2a_protocol_types::task::{ContextId, TaskId};
314
315 #[test]
316 fn bad_request_maps_to_invalid_argument() {
317 let status = bad_request(ConvertError {
318 field: "message.role",
319 reason: "unknown Role number 9".into(),
320 });
321 assert_eq!(status.code(), tonic::Code::InvalidArgument);
322 assert!(status.message().contains("message.role"));
323 }
324
325 #[test]
326 fn bad_response_maps_to_internal() {
327 let status = bad_response(ConvertError {
328 field: "task.metadata",
329 reason: "boom".into(),
330 });
331 assert_eq!(status.code(), tonic::Code::Internal);
332 }
333
334 #[test]
335 fn send_result_message_wraps_into_stream_payload() {
336 let resp = SendMessageResponse::Message(Message {
337 id: MessageId("m".into()),
338 role: MessageRole::Agent,
339 parts: vec![Part {
340 content: PartContent::Text("hi".into()),
341 metadata: None,
342 filename: None,
343 media_type: None,
344 }],
345 task_id: Some(TaskId("t".into())),
346 context_id: Some(ContextId("c".into())),
347 reference_task_ids: None,
348 extensions: None,
349 metadata: None,
350 });
351 let stream = send_result_to_stream(resp).unwrap();
352 assert!(matches!(
353 stream.payload,
354 Some(apb::stream_response::Payload::Message(_))
355 ));
356 }
357}