use std::pin::Pin;
use std::sync::Arc;
use a2a_protocol_types::proto as apb;
use a2a_protocol_types::proto::convert::ConvertError;
use tokio::sync::mpsc;
use tokio_stream::wrappers::ReceiverStream;
use tonic::{Request, Response, Status};
use super::helpers::{server_error_to_status, validated_metadata};
use super::pb::a2a_service_server::A2aService;
use super::GrpcConfig;
use crate::handler::{RequestHandler, SendMessageResult};
type NativeStream =
Pin<Box<dyn tokio_stream::Stream<Item = Result<apb::StreamResponse, Status>> + Send + 'static>>;
#[allow(clippy::needless_pass_by_value)]
fn bad_request(err: ConvertError) -> Status {
Status::invalid_argument(err.to_string())
}
#[allow(clippy::needless_pass_by_value)]
fn bad_response(err: ConvertError) -> Status {
Status::internal(format!("response not representable in protobuf: {err}"))
}
fn send_result_to_stream(
resp: a2a_protocol_types::responses::SendMessageResponse,
) -> Result<apb::StreamResponse, ConvertError> {
let payload = match resp {
a2a_protocol_types::responses::SendMessageResponse::Task(t) => {
apb::stream_response::Payload::Task(t.try_into()?)
}
a2a_protocol_types::responses::SendMessageResponse::Message(m) => {
apb::stream_response::Payload::Message(m.try_into()?)
}
other => {
return Err(ConvertError {
field: "sendMessageResponse.payload",
reason: format!("unsupported response variant: {other:?}"),
})
}
};
Ok(apb::StreamResponse {
payload: Some(payload),
})
}
fn reader_to_native_stream(
mut reader: crate::streaming::InMemoryQueueReader,
capacity: usize,
) -> NativeStream {
use crate::streaming::EventQueueReader;
let (tx, rx) = mpsc::channel(capacity);
tokio::spawn(async move {
loop {
match reader.read().await {
Some(Ok(event)) => {
let item = apb::StreamResponse::try_from(event).map_err(bad_response);
let is_err = item.is_err();
if tx.send(item).await.is_err() || is_err {
break;
}
}
Some(Err(_)) => {
let _ = tx.send(Err(Status::internal("event queue error"))).await;
break;
}
None => break,
}
}
});
Box::pin(ReceiverStream::new(rx))
}
pub struct A2aServiceImpl {
pub(super) handler: Arc<RequestHandler>,
pub(super) config: GrpcConfig,
}
#[tonic::async_trait]
impl A2aService for A2aServiceImpl {
async fn send_message(
&self,
request: Request<apb::SendMessageRequest>,
) -> Result<Response<apb::SendMessageResponse>, Status> {
let headers = validated_metadata(request.metadata())?;
let params: a2a_protocol_types::params::MessageSendParams =
request.into_inner().try_into().map_err(bad_request)?;
match self
.handler
.on_send_message(params, false, Some(&headers))
.await
{
Ok(SendMessageResult::Response(resp)) => {
Ok(Response::new(resp.try_into().map_err(bad_response)?))
}
Ok(SendMessageResult::Stream(_)) => Err(Status::internal(
"unexpected stream response for unary call",
)),
Err(e) => Err(server_error_to_status(&e)),
}
}
type SendStreamingMessageStream = NativeStream;
async fn send_streaming_message(
&self,
request: Request<apb::SendMessageRequest>,
) -> Result<Response<Self::SendStreamingMessageStream>, Status> {
let headers = validated_metadata(request.metadata())?;
let params: a2a_protocol_types::params::MessageSendParams =
request.into_inner().try_into().map_err(bad_request)?;
match self
.handler
.on_send_message(params, true, Some(&headers))
.await
{
Ok(SendMessageResult::Stream(reader)) => Ok(Response::new(reader_to_native_stream(
reader,
self.config.stream_channel_capacity,
))),
Ok(SendMessageResult::Response(resp)) => {
let payload = send_result_to_stream(resp).map_err(bad_response)?;
let stream = Box::pin(tokio_stream::once(Ok(payload)));
Ok(Response::new(stream as NativeStream))
}
Err(e) => Err(server_error_to_status(&e)),
}
}
async fn get_task(
&self,
request: Request<apb::GetTaskRequest>,
) -> Result<Response<apb::Task>, Status> {
let headers = validated_metadata(request.metadata())?;
let params: a2a_protocol_types::params::TaskQueryParams =
request.into_inner().try_into().map_err(bad_request)?;
match self.handler.on_get_task(params, Some(&headers)).await {
Ok(task) => Ok(Response::new(task.try_into().map_err(bad_response)?)),
Err(e) => Err(server_error_to_status(&e)),
}
}
async fn list_tasks(
&self,
request: Request<apb::ListTasksRequest>,
) -> Result<Response<apb::ListTasksResponse>, Status> {
let headers = validated_metadata(request.metadata())?;
let params: a2a_protocol_types::params::ListTasksParams =
request.into_inner().try_into().map_err(bad_request)?;
match self.handler.on_list_tasks(params, Some(&headers)).await {
Ok(resp) => Ok(Response::new(resp.try_into().map_err(bad_response)?)),
Err(e) => Err(server_error_to_status(&e)),
}
}
async fn cancel_task(
&self,
request: Request<apb::CancelTaskRequest>,
) -> Result<Response<apb::Task>, Status> {
let headers = validated_metadata(request.metadata())?;
let params: a2a_protocol_types::params::CancelTaskParams =
request.into_inner().try_into().map_err(bad_request)?;
match self.handler.on_cancel_task(params, Some(&headers)).await {
Ok(task) => Ok(Response::new(task.try_into().map_err(bad_response)?)),
Err(e) => Err(server_error_to_status(&e)),
}
}
type SubscribeToTaskStream = NativeStream;
async fn subscribe_to_task(
&self,
request: Request<apb::SubscribeToTaskRequest>,
) -> Result<Response<Self::SubscribeToTaskStream>, Status> {
let headers = validated_metadata(request.metadata())?;
let params: a2a_protocol_types::params::TaskIdParams = request.into_inner().into();
match self.handler.on_resubscribe(params, Some(&headers)).await {
Ok(reader) => Ok(Response::new(reader_to_native_stream(
reader,
self.config.stream_channel_capacity,
))),
Err(e) => Err(server_error_to_status(&e)),
}
}
async fn create_task_push_notification_config(
&self,
request: Request<apb::TaskPushNotificationConfig>,
) -> Result<Response<apb::TaskPushNotificationConfig>, Status> {
let headers = validated_metadata(request.metadata())?;
let config: a2a_protocol_types::push::TaskPushNotificationConfig =
request.into_inner().into();
match self
.handler
.on_set_push_config(config, Some(&headers))
.await
{
Ok(cfg) => Ok(Response::new(cfg.into())),
Err(e) => Err(server_error_to_status(&e)),
}
}
async fn get_task_push_notification_config(
&self,
request: Request<apb::GetTaskPushNotificationConfigRequest>,
) -> Result<Response<apb::TaskPushNotificationConfig>, Status> {
let headers = validated_metadata(request.metadata())?;
let params: a2a_protocol_types::params::GetPushConfigParams = request.into_inner().into();
match self
.handler
.on_get_push_config(params, Some(&headers))
.await
{
Ok(cfg) => Ok(Response::new(cfg.into())),
Err(e) => Err(server_error_to_status(&e)),
}
}
async fn list_task_push_notification_configs(
&self,
request: Request<apb::ListTaskPushNotificationConfigsRequest>,
) -> Result<Response<apb::ListTaskPushNotificationConfigsResponse>, Status> {
let headers = validated_metadata(request.metadata())?;
let params: a2a_protocol_types::params::ListPushConfigsParams =
request.into_inner().try_into().map_err(bad_request)?;
match self
.handler
.on_list_push_configs(¶ms.task_id, params.tenant.as_deref(), Some(&headers))
.await
{
Ok(configs) => Ok(Response::new(
apb::ListTaskPushNotificationConfigsResponse {
configs: configs.into_iter().map(Into::into).collect(),
next_page_token: String::new(),
},
)),
Err(e) => Err(server_error_to_status(&e)),
}
}
async fn delete_task_push_notification_config(
&self,
request: Request<apb::DeleteTaskPushNotificationConfigRequest>,
) -> Result<Response<()>, Status> {
let headers = validated_metadata(request.metadata())?;
let params: a2a_protocol_types::params::DeletePushConfigParams =
request.into_inner().into();
match self
.handler
.on_delete_push_config(params, Some(&headers))
.await
{
Ok(()) => Ok(Response::new(())),
Err(e) => Err(server_error_to_status(&e)),
}
}
async fn get_extended_agent_card(
&self,
request: Request<apb::GetExtendedAgentCardRequest>,
) -> Result<Response<apb::AgentCard>, Status> {
let headers = validated_metadata(request.metadata())?;
match self
.handler
.on_get_extended_agent_card(Some(&headers))
.await
{
Ok(card) => Ok(Response::new(card.try_into().map_err(bad_response)?)),
Err(e) => Err(server_error_to_status(&e)),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use a2a_protocol_types::message::{Message, MessageId, MessageRole, Part, PartContent};
use a2a_protocol_types::responses::SendMessageResponse;
use a2a_protocol_types::task::{ContextId, TaskId};
#[test]
fn bad_request_maps_to_invalid_argument() {
let status = bad_request(ConvertError {
field: "message.role",
reason: "unknown Role number 9".into(),
});
assert_eq!(status.code(), tonic::Code::InvalidArgument);
assert!(status.message().contains("message.role"));
}
#[test]
fn bad_response_maps_to_internal() {
let status = bad_response(ConvertError {
field: "task.metadata",
reason: "boom".into(),
});
assert_eq!(status.code(), tonic::Code::Internal);
}
#[test]
fn send_result_message_wraps_into_stream_payload() {
let resp = SendMessageResponse::Message(Message {
id: MessageId("m".into()),
role: MessageRole::Agent,
parts: vec![Part {
content: PartContent::Text("hi".into()),
metadata: None,
filename: None,
media_type: None,
}],
task_id: Some(TaskId("t".into())),
context_id: Some(ContextId("c".into())),
reference_task_ids: None,
extensions: None,
metadata: None,
});
let stream = send_result_to_stream(resp).unwrap();
assert!(matches!(
stream.payload,
Some(apb::stream_response::Payload::Message(_))
));
}
use crate::agent_executor;
use crate::builder::RequestHandlerBuilder;
use a2a_protocol_types::agent_card::{
AgentCapabilities, AgentCard, AgentInterface, AgentSkill,
};
struct NoopExecutor;
agent_executor!(NoopExecutor, |_ctx, _queue| async { Ok(()) });
struct NoopSender;
impl crate::push::PushSender for NoopSender {
fn send<'a>(
&'a self,
_url: &'a str,
_event: &'a a2a_protocol_types::events::StreamResponse,
_config: &'a a2a_protocol_types::push::TaskPushNotificationConfig,
) -> Pin<
Box<
dyn std::future::Future<Output = a2a_protocol_types::error::A2aResult<()>>
+ Send
+ 'a,
>,
> {
Box::pin(async { Ok(()) })
}
fn allows_private_urls(&self) -> bool {
true
}
}
fn test_card() -> AgentCard {
AgentCard {
url: None,
name: "native-grpc-test-agent".into(),
description: "Fixture for the canonical gRPC binding".into(),
version: "1.0.0".into(),
supported_interfaces: vec![AgentInterface {
url: "grpc://localhost:50051".into(),
protocol_binding: "gRPC".into(),
protocol_version: "1.0.0".into(),
tenant: None,
}],
default_input_modes: vec!["text/plain".into()],
default_output_modes: vec!["text/plain".into()],
skills: vec![AgentSkill {
id: "noop".into(),
name: "Noop".into(),
description: "Does nothing".into(),
tags: vec!["test".into()],
examples: None,
input_modes: None,
output_modes: None,
security_requirements: None,
}],
capabilities: AgentCapabilities::none()
.with_extended_agent_card(true)
.with_push_notifications(true),
provider: None,
icon_url: None,
documentation_url: None,
security_schemes: None,
security_requirements: None,
signatures: None,
}
}
fn service() -> A2aServiceImpl {
A2aServiceImpl {
handler: Arc::new(
RequestHandlerBuilder::new(NoopExecutor)
.with_agent_card(test_card())
.with_push_sender(NoopSender)
.allow_unauthenticated_extended_card()
.build()
.expect("default build should succeed"),
),
config: GrpcConfig::default(),
}
}
fn send_request(text: &str) -> apb::SendMessageRequest {
apb::SendMessageRequest {
tenant: String::new(),
message: Some(apb::Message {
message_id: format!("msg-{text}"),
context_id: String::new(),
task_id: String::new(),
role: apb::Role::User as i32,
parts: vec![apb::Part {
metadata: None,
filename: String::new(),
media_type: String::new(),
content: Some(apb::part::Content::Text(text.into())),
}],
metadata: None,
extensions: Vec::new(),
reference_task_ids: Vec::new(),
}),
configuration: None,
metadata: None,
}
}
async fn seed_task(svc: &A2aServiceImpl) -> String {
let resp = svc
.send_message(Request::new(send_request("seed")))
.await
.expect("send_message should succeed")
.into_inner();
match resp.payload {
Some(apb::send_message_response::Payload::Task(t)) => t.id,
other => panic!("expected a Task payload, got {other:?}"),
}
}
#[tokio::test]
async fn send_message_returns_a_populated_payload() {
let svc = service();
let resp = svc
.send_message(Request::new(send_request("hello")))
.await
.expect("send_message should succeed")
.into_inner();
let payload = resp.payload.expect("a default response carries no payload");
match payload {
apb::send_message_response::Payload::Task(t) => {
assert!(!t.id.is_empty(), "the created task must carry an id");
}
apb::send_message_response::Payload::Message(m) => {
assert!(!m.message_id.is_empty(), "the reply must carry an id");
}
}
}
#[tokio::test]
async fn get_task_returns_the_task_that_was_created() {
let svc = service();
let id = seed_task(&svc).await;
let task = svc
.get_task(Request::new(apb::GetTaskRequest {
tenant: String::new(),
id: id.clone(),
history_length: None,
}))
.await
.expect("get_task should find the seeded task")
.into_inner();
assert_eq!(task.id, id, "the id round-trips through the binding");
}
#[tokio::test]
async fn get_task_maps_a_missing_task_to_not_found() {
let svc = service();
let status = svc
.get_task(Request::new(apb::GetTaskRequest {
tenant: String::new(),
id: "no-such-task".into(),
history_length: None,
}))
.await
.expect_err("an unknown id must not resolve");
assert_eq!(status.code(), tonic::Code::NotFound);
}
#[tokio::test]
async fn list_tasks_returns_the_seeded_task() {
let svc = service();
let id = seed_task(&svc).await;
let resp = svc
.list_tasks(Request::new(apb::ListTasksRequest::default()))
.await
.expect("list_tasks should succeed")
.into_inner();
assert!(
resp.tasks.iter().any(|t| t.id == id),
"the seeded task must appear in the listing, got {:?}",
resp.tasks.iter().map(|t| &t.id).collect::<Vec<_>>()
);
}
#[tokio::test]
async fn cancel_task_moves_the_task_out_of_a_running_state() {
let svc = service();
let id = seed_task(&svc).await;
let task = svc
.cancel_task(Request::new(apb::CancelTaskRequest {
tenant: String::new(),
id: id.clone(),
metadata: None,
}))
.await
.expect("cancel_task should succeed")
.into_inner();
assert_eq!(task.id, id);
let status = task.status.expect("a cancelled task carries a status");
assert_eq!(
status.state,
apb::TaskState::Canceled as i32,
"cancel must actually transition the task"
);
}
#[tokio::test]
async fn cancel_task_maps_a_missing_task_to_not_found() {
let svc = service();
let status = svc
.cancel_task(Request::new(apb::CancelTaskRequest {
tenant: String::new(),
id: "no-such-task".into(),
metadata: None,
}))
.await
.expect_err("an unknown id must not resolve");
assert_eq!(status.code(), tonic::Code::NotFound);
}
async fn create_push(
svc: &A2aServiceImpl,
task_id: &str,
url: &str,
) -> apb::TaskPushNotificationConfig {
svc.create_task_push_notification_config(Request::new(apb::TaskPushNotificationConfig {
tenant: String::new(),
id: String::new(),
task_id: task_id.to_owned(),
url: url.to_owned(),
token: String::new(),
authentication: None,
}))
.await
.expect("create_task_push_notification_config should succeed")
.into_inner()
}
async fn list_push(
svc: &A2aServiceImpl,
task_id: &str,
) -> Vec<apb::TaskPushNotificationConfig> {
svc.list_task_push_notification_configs(Request::new(
apb::ListTaskPushNotificationConfigsRequest {
tenant: String::new(),
task_id: task_id.to_owned(),
page_size: 0,
page_token: String::new(),
},
))
.await
.expect("list_task_push_notification_configs should succeed")
.into_inner()
.configs
}
#[tokio::test]
async fn create_push_config_echoes_the_registered_url() {
let svc = service();
let task_id = seed_task(&svc).await;
let created = create_push(&svc, &task_id, "https://example.test/hook").await;
assert_eq!(created.url, "https://example.test/hook");
assert_eq!(created.task_id, task_id);
}
#[tokio::test]
async fn get_push_config_returns_what_create_stored() {
let svc = service();
let task_id = seed_task(&svc).await;
let created = create_push(&svc, &task_id, "https://example.test/get").await;
let fetched = svc
.get_task_push_notification_config(Request::new(
apb::GetTaskPushNotificationConfigRequest {
tenant: String::new(),
task_id: task_id.clone(),
id: created.id.clone(),
},
))
.await
.expect("the config was just created")
.into_inner();
assert_eq!(fetched.id, created.id);
assert_eq!(fetched.url, "https://example.test/get");
}
#[tokio::test]
async fn list_push_configs_includes_the_created_one() {
let svc = service();
let task_id = seed_task(&svc).await;
let created = create_push(&svc, &task_id, "https://example.test/list").await;
let configs = list_push(&svc, &task_id).await;
assert!(
configs.iter().any(|c| c.id == created.id),
"the created config must appear in the listing"
);
}
#[tokio::test]
async fn delete_push_config_removes_it_from_the_listing() {
let svc = service();
let task_id = seed_task(&svc).await;
let created = create_push(&svc, &task_id, "https://example.test/delete").await;
assert_eq!(list_push(&svc, &task_id).await.len(), 1, "precondition");
svc.delete_task_push_notification_config(Request::new(
apb::DeleteTaskPushNotificationConfigRequest {
tenant: String::new(),
task_id: task_id.clone(),
id: created.id.clone(),
},
))
.await
.expect("delete_task_push_notification_config should succeed");
assert!(
list_push(&svc, &task_id).await.is_empty(),
"delete must actually remove the config"
);
}
#[tokio::test]
async fn reader_to_native_stream_stops_after_a_conversion_error() {
use a2a_protocol_types::events::StreamResponse;
use tokio_stream::StreamExt;
use crate::streaming::event_queue::new_in_memory_queue;
use crate::streaming::EventQueueWriter;
fn message(metadata: serde_json::Value) -> Message {
Message {
id: MessageId("m".into()),
role: MessageRole::Agent,
parts: vec![Part {
content: PartContent::Text("hi".into()),
metadata: None,
filename: None,
media_type: None,
}],
task_id: Some(TaskId("t".into())),
context_id: Some(ContextId("c".into())),
reference_task_ids: None,
extensions: None,
metadata: Some(metadata),
}
}
let (writer, reader) = new_in_memory_queue();
writer
.write(StreamResponse::Message(message(serde_json::json!(
"not-an-object"
))))
.await
.expect("write of the unconvertible event");
writer
.write(StreamResponse::Message(message(
serde_json::json!({"ok": true}),
)))
.await
.expect("write of the convertible event");
drop(writer);
let mut stream = reader_to_native_stream(reader, 8);
let first = stream.next().await.expect("the error item is delivered");
let status = first.expect_err("an unconvertible event surfaces as an error");
assert_eq!(status.code(), tonic::Code::Internal);
assert!(
stream.next().await.is_none(),
"the stream must end at the conversion error, not carry on"
);
}
#[tokio::test]
async fn get_extended_agent_card_returns_the_configured_card() {
let svc = service();
let card = svc
.get_extended_agent_card(Request::new(apb::GetExtendedAgentCardRequest {
tenant: String::new(),
}))
.await
.expect("the card is configured and unauthenticated access is opted in")
.into_inner();
assert_eq!(
card.name, "native-grpc-test-agent",
"a default card would carry an empty name"
);
}
}