use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::time::Duration;
use a2a_protocol_types::proto as apb;
use a2a_protocol_types::proto::convert::ConvertError;
use tokio::sync::mpsc;
use tonic::transport::Channel;
use crate::error::{ClientError, ClientResult};
use crate::streaming::EventStream;
use crate::transport::Transport;
mod proto {
#![allow(
clippy::all,
clippy::pedantic,
clippy::nursery,
missing_docs,
unused_qualifications
)]
tonic::include_proto!("lf.a2a.v1");
}
use proto::a2a_service_client::A2aServiceClient;
#[derive(Debug, Clone)]
pub struct GrpcTransportConfig {
pub timeout: Duration,
pub connect_timeout: Duration,
pub max_message_size: usize,
pub stream_channel_capacity: usize,
}
impl Default for GrpcTransportConfig {
fn default() -> Self {
Self {
timeout: Duration::from_secs(30),
connect_timeout: Duration::from_secs(10),
max_message_size: 4 * 1024 * 1024,
stream_channel_capacity: 64,
}
}
}
impl GrpcTransportConfig {
#[must_use]
pub const fn with_timeout(mut self, timeout: Duration) -> Self {
self.timeout = timeout;
self
}
#[must_use]
pub const fn with_connect_timeout(mut self, timeout: Duration) -> Self {
self.connect_timeout = timeout;
self
}
#[must_use]
pub const fn with_max_message_size(mut self, size: usize) -> Self {
self.max_message_size = size;
self
}
#[must_use]
pub const fn with_stream_channel_capacity(mut self, capacity: usize) -> Self {
self.stream_channel_capacity = capacity;
self
}
}
#[derive(Clone, Debug)]
pub struct GrpcTransport {
inner: Arc<Inner>,
}
#[derive(Debug)]
struct Inner {
channel: Channel,
endpoint: String,
config: GrpcTransportConfig,
}
impl GrpcTransport {
pub async fn connect(endpoint: impl Into<String>) -> ClientResult<Self> {
Self::connect_with_config(endpoint, GrpcTransportConfig::default()).await
}
pub async fn connect_with_config(
endpoint: impl Into<String>,
config: GrpcTransportConfig,
) -> ClientResult<Self> {
let endpoint_str = endpoint.into();
validate_url(&endpoint_str)?;
let channel = tonic::transport::Channel::from_shared(endpoint_str.clone())
.map_err(|e| ClientError::InvalidEndpoint(format!("invalid gRPC endpoint: {e}")))?
.connect_timeout(config.connect_timeout)
.timeout(config.timeout)
.connect()
.await
.map_err(|e| ClientError::Transport(format!("gRPC connect failed: {e}")))?;
Ok(Self {
inner: Arc::new(Inner {
channel,
endpoint: endpoint_str,
config,
}),
})
}
#[must_use]
pub fn endpoint(&self) -> &str {
&self.inner.endpoint
}
fn client(&self) -> A2aServiceClient<Channel> {
A2aServiceClient::new(self.inner.channel.clone())
.max_decoding_message_size(self.inner.config.max_message_size)
.max_encoding_message_size(self.inner.config.max_message_size)
}
fn request<T>(
&self,
message: T,
extra_headers: &HashMap<String, String>,
with_deadline: bool,
) -> ClientResult<tonic::Request<T>> {
let mut req = tonic::Request::new(message);
if with_deadline {
req.set_timeout(self.inner.config.timeout);
}
Self::add_metadata(&mut req, extra_headers)?;
Ok(req)
}
fn add_metadata<T>(
req: &mut tonic::Request<T>,
extra_headers: &HashMap<String, String>,
) -> ClientResult<()> {
let md = req.metadata_mut();
md.insert(
"a2a-version",
a2a_protocol_types::A2A_VERSION
.parse()
.unwrap_or_else(|_| tonic::metadata::MetadataValue::from_static("")),
);
for (k, v) in extra_headers {
let key = k.parse::<tonic::metadata::MetadataKey<_>>().map_err(|e| {
ClientError::Transport(format!("invalid gRPC metadata key {k:?}: {e}"))
})?;
let val = v
.parse::<tonic::metadata::MetadataValue<_>>()
.map_err(|_| {
ClientError::Transport(format!("invalid gRPC metadata value for key {k:?}"))
})?;
md.insert(key, val);
}
Ok(())
}
fn parse_params<T: serde::de::DeserializeOwned>(params: serde_json::Value) -> ClientResult<T> {
serde_json::from_value(params).map_err(ClientError::Serialization)
}
fn to_json<T: serde::Serialize>(value: &T) -> ClientResult<serde_json::Value> {
serde_json::to_value(value).map_err(ClientError::Serialization)
}
fn status_to_error(status: &tonic::Status) -> ClientError {
match status.code() {
tonic::Code::DeadlineExceeded => {
ClientError::Timeout(format!("gRPC deadline exceeded: {}", status.message()))
}
tonic::Code::Cancelled => {
ClientError::Timeout(format!("gRPC request cancelled: {}", status.message()))
}
tonic::Code::Unavailable => {
ClientError::HttpClient(format!("gRPC unavailable: {}", status.message()))
}
tonic::Code::ResourceExhausted => ClientError::UnexpectedStatus {
status: 429,
body: status.message().to_owned(),
retry_after: None,
},
_ => {
use tonic_types::StatusExt as _;
let code = status
.get_details_error_info()
.and_then(|info| a2a_protocol_types::ErrorCode::from_a2a_reason(&info.reason))
.unwrap_or_else(|| grpc_code_to_error_code(status.code()));
let a2a = a2a_protocol_types::A2aError::new(code, status.message().to_owned());
ClientError::Protocol(a2a)
}
}
}
async fn execute_unary(
&self,
method: &str,
params: serde_json::Value,
extra_headers: &HashMap<String, String>,
) -> ClientResult<serde_json::Value> {
trace_info!(
method,
endpoint = %self.inner.endpoint,
"sending gRPC request"
);
let mut client = self.client();
tokio::time::timeout(
self.inner.config.timeout,
self.dispatch_unary(&mut client, method, params, extra_headers),
)
.await
.map_err(|_| {
trace_error!(method, "gRPC request timed out");
ClientError::Timeout("gRPC request timed out".into())
})?
}
#[allow(clippy::too_many_lines)]
async fn dispatch_unary(
&self,
client: &mut A2aServiceClient<Channel>,
method: &str,
params: serde_json::Value,
extra_headers: &HashMap<String, String>,
) -> ClientResult<serde_json::Value> {
match method {
"SendMessage" => {
let p: a2a_protocol_types::params::MessageSendParams = Self::parse_params(params)?;
let req = apb::SendMessageRequest::try_from(p).map_err(convert_error)?;
let resp = client
.send_message(self.request(req, extra_headers, true)?)
.await
.map_err(|s| Self::status_to_error(&s))?;
let domain: a2a_protocol_types::responses::SendMessageResponse =
resp.into_inner().try_into().map_err(convert_error)?;
Self::to_json(&domain)
}
"GetTask" => {
let p: a2a_protocol_types::params::TaskQueryParams = Self::parse_params(params)?;
let req = apb::GetTaskRequest::try_from(p).map_err(convert_error)?;
let resp = client
.get_task(self.request(req, extra_headers, true)?)
.await
.map_err(|s| Self::status_to_error(&s))?;
let domain: a2a_protocol_types::task::Task =
resp.into_inner().try_into().map_err(convert_error)?;
Self::to_json(&domain)
}
"ListTasks" => {
let p: a2a_protocol_types::params::ListTasksParams = Self::parse_params(params)?;
let req = apb::ListTasksRequest::try_from(p).map_err(convert_error)?;
let resp = client
.list_tasks(self.request(req, extra_headers, true)?)
.await
.map_err(|s| Self::status_to_error(&s))?;
let domain: a2a_protocol_types::responses::TaskListResponse =
resp.into_inner().try_into().map_err(convert_error)?;
Self::to_json(&domain)
}
"CancelTask" => {
let p: a2a_protocol_types::params::CancelTaskParams = Self::parse_params(params)?;
let req = apb::CancelTaskRequest::try_from(p).map_err(convert_error)?;
let resp = client
.cancel_task(self.request(req, extra_headers, true)?)
.await
.map_err(|s| Self::status_to_error(&s))?;
let domain: a2a_protocol_types::task::Task =
resp.into_inner().try_into().map_err(convert_error)?;
Self::to_json(&domain)
}
"CreateTaskPushNotificationConfig" => {
let p: a2a_protocol_types::push::TaskPushNotificationConfig =
Self::parse_params(params)?;
let req = apb::TaskPushNotificationConfig::from(p);
let resp = client
.create_task_push_notification_config(self.request(req, extra_headers, true)?)
.await
.map_err(|s| Self::status_to_error(&s))?;
let domain: a2a_protocol_types::push::TaskPushNotificationConfig =
resp.into_inner().into();
Self::to_json(&domain)
}
"GetTaskPushNotificationConfig" => {
let p: a2a_protocol_types::params::GetPushConfigParams =
Self::parse_params(params)?;
let req = apb::GetTaskPushNotificationConfigRequest::from(p);
let resp = client
.get_task_push_notification_config(self.request(req, extra_headers, true)?)
.await
.map_err(|s| Self::status_to_error(&s))?;
let domain: a2a_protocol_types::push::TaskPushNotificationConfig =
resp.into_inner().into();
Self::to_json(&domain)
}
"ListTaskPushNotificationConfigs" => {
let p: a2a_protocol_types::params::ListPushConfigsParams =
Self::parse_params(params)?;
let req = apb::ListTaskPushNotificationConfigsRequest::try_from(p)
.map_err(convert_error)?;
let resp = client
.list_task_push_notification_configs(self.request(req, extra_headers, true)?)
.await
.map_err(|s| Self::status_to_error(&s))?;
let domain: a2a_protocol_types::responses::ListPushConfigsResponse =
resp.into_inner().into();
Self::to_json(&domain)
}
"DeleteTaskPushNotificationConfig" => {
let p: a2a_protocol_types::params::DeletePushConfigParams =
Self::parse_params(params)?;
let req = apb::DeleteTaskPushNotificationConfigRequest::from(p);
client
.delete_task_push_notification_config(self.request(req, extra_headers, true)?)
.await
.map_err(|s| Self::status_to_error(&s))?;
Ok(serde_json::json!({}))
}
"GetExtendedAgentCard" => {
let params = if params.is_null() {
serde_json::json!({})
} else {
params
};
let p: a2a_protocol_types::params::GetExtendedAgentCardParams =
Self::parse_params(params)?;
let req = apb::GetExtendedAgentCardRequest::from(p);
let resp = client
.get_extended_agent_card(self.request(req, extra_headers, true)?)
.await
.map_err(|s| Self::status_to_error(&s))?;
let domain: a2a_protocol_types::agent_card::AgentCard =
resp.into_inner().try_into().map_err(convert_error)?;
Self::to_json(&domain)
}
other => Err(ClientError::Protocol(a2a_protocol_types::A2aError::new(
a2a_protocol_types::ErrorCode::MethodNotFound,
format!("unknown gRPC method: {other}"),
))),
}
}
async fn execute_streaming(
&self,
method: &str,
params: serde_json::Value,
extra_headers: &HashMap<String, String>,
) -> ClientResult<EventStream> {
trace_info!(
method,
endpoint = %self.inner.endpoint,
"opening gRPC stream"
);
let mut client = self.client();
let stream = tokio::time::timeout(self.inner.config.timeout, async {
match method {
"SendStreamingMessage" => {
let p: a2a_protocol_types::params::MessageSendParams =
Self::parse_params(params)?;
let req = apb::SendMessageRequest::try_from(p).map_err(convert_error)?;
client
.send_streaming_message(self.request(req, extra_headers, false)?)
.await
.map(tonic::Response::into_inner)
.map_err(|s| Self::status_to_error(&s))
}
"SubscribeToTask" => {
let p: a2a_protocol_types::params::TaskIdParams = Self::parse_params(params)?;
let req = apb::SubscribeToTaskRequest::from(p);
client
.subscribe_to_task(self.request(req, extra_headers, false)?)
.await
.map(tonic::Response::into_inner)
.map_err(|s| Self::status_to_error(&s))
}
other => Err(ClientError::Protocol(a2a_protocol_types::A2aError::new(
a2a_protocol_types::ErrorCode::MethodNotFound,
format!("unknown streaming gRPC method: {other}"),
))),
}
})
.await
.map_err(|_| {
trace_error!(method, "gRPC stream connect timed out");
ClientError::Timeout("gRPC stream connect timed out".into())
})??;
let cap = self.inner.config.stream_channel_capacity;
let (tx, rx) = mpsc::channel::<crate::streaming::event_stream::BodyChunk>(cap);
let task_handle = tokio::spawn(async move {
grpc_stream_reader_task(stream, tx).await;
});
Ok(
EventStream::with_status(rx, task_handle.abort_handle(), 200)
.with_first_event_timeout(self.inner.config.timeout),
)
}
}
impl Transport for GrpcTransport {
fn send_request<'a>(
&'a self,
method: &'a str,
params: serde_json::Value,
extra_headers: &'a HashMap<String, String>,
) -> Pin<Box<dyn Future<Output = ClientResult<serde_json::Value>> + Send + 'a>> {
Box::pin(self.execute_unary(method, params, extra_headers))
}
fn send_streaming_request<'a>(
&'a self,
method: &'a str,
params: serde_json::Value,
extra_headers: &'a HashMap<String, String>,
) -> Pin<Box<dyn Future<Output = ClientResult<EventStream>> + Send + 'a>> {
Box::pin(self.execute_streaming(method, params, extra_headers))
}
}
async fn grpc_stream_reader_task<S>(
mut stream: S,
tx: mpsc::Sender<crate::streaming::event_stream::BodyChunk>,
) where
S: tonic::codegen::tokio_stream::Stream<Item = Result<apb::StreamResponse, tonic::Status>>
+ Unpin,
{
use tonic::codegen::tokio_stream::StreamExt;
loop {
match stream.next().await {
Some(Ok(pb_event)) => {
let event: a2a_protocol_types::events::StreamResponse =
match pb_event.try_into().map_err(convert_error) {
Ok(e) => e,
Err(err) => {
let _ = tx.send(Err(err)).await;
break;
}
};
let json_str = match serde_json::to_string(&event) {
Ok(s) => s,
Err(e) => {
let _ = tx.send(Err(ClientError::Serialization(e))).await;
break;
}
};
let envelope =
format!("data: {{\"jsonrpc\":\"2.0\",\"id\":null,\"result\":{json_str}}}\n\n");
if tx
.send(Ok(hyper::body::Bytes::from(envelope)))
.await
.is_err()
{
break;
}
}
Some(Err(status)) => {
let _ = tx.send(Err(GrpcTransport::status_to_error(&status))).await;
break;
}
None => break,
}
}
}
#[allow(clippy::needless_pass_by_value)]
fn convert_error(err: ConvertError) -> ClientError {
ClientError::Transport(format!("protobuf conversion failed: {err}"))
}
fn validate_url(url: &str) -> ClientResult<()> {
if url.is_empty() {
return Err(ClientError::InvalidEndpoint("URL must not be empty".into()));
}
if !url.starts_with("http://") && !url.starts_with("https://") {
return Err(ClientError::InvalidEndpoint(format!(
"URL must start with http:// or https://: {url}"
)));
}
Ok(())
}
const fn grpc_code_to_error_code(code: tonic::Code) -> a2a_protocol_types::ErrorCode {
match code {
tonic::Code::NotFound => a2a_protocol_types::ErrorCode::TaskNotFound,
tonic::Code::InvalidArgument
| tonic::Code::Unauthenticated
| tonic::Code::PermissionDenied
| tonic::Code::ResourceExhausted => a2a_protocol_types::ErrorCode::InvalidParams,
tonic::Code::Unimplemented => a2a_protocol_types::ErrorCode::MethodNotFound,
tonic::Code::FailedPrecondition => a2a_protocol_types::ErrorCode::TaskNotCancelable,
_ => a2a_protocol_types::ErrorCode::InternalError,
}
}
#[cfg(test)]
mod tests {
use super::*;
use a2a_protocol_types::events::TaskStatusUpdateEvent;
use a2a_protocol_types::task::{ContextId, TaskId, TaskState, TaskStatus};
#[test]
fn validate_url_rejects_empty() {
assert!(validate_url("").is_err());
}
#[test]
fn validate_url_rejects_non_http() {
assert!(validate_url("ftp://example.com").is_err());
}
#[test]
fn validate_url_accepts_http() {
assert!(validate_url("http://localhost:50051").is_ok());
}
#[test]
fn config_default_timeout() {
let cfg = GrpcTransportConfig::default();
assert_eq!(cfg.timeout, Duration::from_secs(30));
}
#[test]
fn config_builder() {
let cfg = GrpcTransportConfig::default()
.with_timeout(Duration::from_secs(60))
.with_max_message_size(8 * 1024 * 1024)
.with_stream_channel_capacity(128);
assert_eq!(cfg.timeout, Duration::from_secs(60));
assert_eq!(cfg.max_message_size, 8 * 1024 * 1024);
assert_eq!(cfg.stream_channel_capacity, 128);
}
#[test]
fn convert_error_maps_to_non_retryable_transport() {
let err = convert_error(ConvertError {
field: "part.raw",
reason: "invalid base64".into(),
});
assert!(
matches!(err, ClientError::Transport(_)),
"conversion failures must be non-retryable: {err:?}"
);
assert!(!err.is_retryable());
}
#[test]
fn grpc_code_not_found_maps_to_task_not_found() {
assert_eq!(
grpc_code_to_error_code(tonic::Code::NotFound),
a2a_protocol_types::ErrorCode::TaskNotFound,
);
}
#[test]
fn grpc_code_invalid_argument_maps_to_invalid_params() {
assert_eq!(
grpc_code_to_error_code(tonic::Code::InvalidArgument),
a2a_protocol_types::ErrorCode::InvalidParams,
);
}
#[test]
fn grpc_code_unauthenticated_maps_to_invalid_params() {
assert_eq!(
grpc_code_to_error_code(tonic::Code::Unauthenticated),
a2a_protocol_types::ErrorCode::InvalidParams,
);
}
#[test]
fn grpc_code_permission_denied_maps_to_invalid_params() {
assert_eq!(
grpc_code_to_error_code(tonic::Code::PermissionDenied),
a2a_protocol_types::ErrorCode::InvalidParams,
);
}
#[test]
fn grpc_code_resource_exhausted_maps_to_invalid_params() {
assert_eq!(
grpc_code_to_error_code(tonic::Code::ResourceExhausted),
a2a_protocol_types::ErrorCode::InvalidParams,
);
}
#[test]
fn grpc_code_unimplemented_maps_to_method_not_found() {
assert_eq!(
grpc_code_to_error_code(tonic::Code::Unimplemented),
a2a_protocol_types::ErrorCode::MethodNotFound,
);
}
#[test]
fn grpc_code_failed_precondition_maps_to_task_not_cancelable() {
assert_eq!(
grpc_code_to_error_code(tonic::Code::FailedPrecondition),
a2a_protocol_types::ErrorCode::TaskNotCancelable,
);
}
#[test]
fn grpc_code_deadline_exceeded_maps_to_internal() {
assert_eq!(
grpc_code_to_error_code(tonic::Code::DeadlineExceeded),
a2a_protocol_types::ErrorCode::InternalError,
);
}
#[test]
fn grpc_code_cancelled_maps_to_internal() {
assert_eq!(
grpc_code_to_error_code(tonic::Code::Cancelled),
a2a_protocol_types::ErrorCode::InternalError,
);
}
#[test]
fn grpc_code_unknown_maps_to_internal() {
assert_eq!(
grpc_code_to_error_code(tonic::Code::Unknown),
a2a_protocol_types::ErrorCode::InternalError,
);
}
#[test]
fn add_metadata_injects_a2a_version() {
let mut req = tonic::Request::new(());
let headers = HashMap::new();
GrpcTransport::add_metadata(&mut req, &headers).expect("valid headers");
let md = req.metadata();
let version_value = md
.get("a2a-version")
.expect("a2a-version header should be present");
assert_eq!(
version_value.to_str().unwrap(),
a2a_protocol_types::A2A_VERSION,
);
}
#[test]
fn add_metadata_injects_extra_headers() {
let mut req = tonic::Request::new(());
let mut headers = HashMap::new();
headers.insert("x-custom".to_string(), "value123".to_string());
GrpcTransport::add_metadata(&mut req, &headers).expect("valid headers");
let md = req.metadata();
assert_eq!(md.get("x-custom").unwrap().to_str().unwrap(), "value123",);
}
#[test]
fn add_metadata_fails_closed_on_invalid_header() {
let mut req = tonic::Request::new(());
let mut headers = HashMap::new();
headers.insert("authorization".to_string(), "Bearer bad\nvalue".to_string());
let result = GrpcTransport::add_metadata(&mut req, &headers);
assert!(
matches!(result, Err(ClientError::Transport(_))),
"invalid metadata must fail closed, got: {result:?}"
);
if let Err(ClientError::Transport(msg)) = result {
assert!(!msg.contains("Bearer bad"), "value leaked in error: {msg}");
}
}
#[test]
fn resource_exhausted_maps_to_retryable_429() {
let status = tonic::Status::resource_exhausted("slow down");
let err = GrpcTransport::status_to_error(&status);
assert!(
matches!(err, ClientError::UnexpectedStatus { status: 429, .. }),
"ResourceExhausted should map to 429, got {err:?}"
);
assert!(
err.is_retryable(),
"gRPC ResourceExhausted must be retryable"
);
}
#[test]
fn status_to_error_deadline_exceeded_is_timeout() {
let status = tonic::Status::deadline_exceeded("test deadline");
let err = GrpcTransport::status_to_error(&status);
assert!(
matches!(err, ClientError::Timeout(_)),
"DeadlineExceeded should map to Timeout, got: {err:?}"
);
}
#[test]
fn status_to_error_cancelled_is_timeout() {
let status = tonic::Status::cancelled("test cancel");
let err = GrpcTransport::status_to_error(&status);
assert!(
matches!(err, ClientError::Timeout(_)),
"Cancelled should map to Timeout, got: {err:?}"
);
}
#[test]
fn status_to_error_unavailable_is_http_client() {
let status = tonic::Status::unavailable("test unavailable");
let err = GrpcTransport::status_to_error(&status);
assert!(
matches!(err, ClientError::HttpClient(_)),
"Unavailable should map to HttpClient, got: {err:?}"
);
}
#[test]
fn status_to_error_other_is_protocol() {
let status = tonic::Status::internal("test internal");
let err = GrpcTransport::status_to_error(&status);
assert!(
matches!(err, ClientError::Protocol(_)),
"other codes should map to Protocol, got: {err:?}"
);
}
#[test]
fn status_to_error_prefers_error_info_reason() {
use tonic_types::StatusExt as _;
let mut details = tonic_types::ErrorDetails::new();
details.set_error_info(
"TASK_NOT_CANCELABLE",
"a2a-protocol.org",
std::collections::HashMap::<String, String>::new(),
);
let status = tonic::Status::with_error_details(
tonic::Code::FailedPrecondition,
"task done",
details,
);
let err = GrpcTransport::status_to_error(&status);
match err {
ClientError::Protocol(a2a) => assert_eq!(
a2a.code,
a2a_protocol_types::ErrorCode::TaskNotCancelable,
"ErrorInfo reason must resolve the exact A2A code"
),
other => panic!("expected Protocol error, got: {other:?}"),
}
}
#[test]
fn status_to_error_unknown_reason_falls_back_to_code() {
use tonic_types::StatusExt as _;
let mut details = tonic_types::ErrorDetails::new();
details.set_error_info(
"SOMETHING_NOVEL",
"a2a-protocol.org",
std::collections::HashMap::<String, String>::new(),
);
let status = tonic::Status::with_error_details(tonic::Code::NotFound, "missing", details);
let err = GrpcTransport::status_to_error(&status);
match err {
ClientError::Protocol(a2a) => assert_eq!(
a2a.code,
a2a_protocol_types::ErrorCode::TaskNotFound,
"unknown reason must fall back to code-based mapping"
),
other => panic!("expected Protocol error, got: {other:?}"),
}
}
fn status_update_event() -> apb::StreamResponse {
let event = TaskStatusUpdateEvent {
task_id: TaskId("t-1".into()),
context_id: ContextId("c-1".into()),
status: TaskStatus {
state: TaskState::Working,
message: None,
timestamp: None,
},
metadata: None,
};
apb::StreamResponse {
payload: Some(apb::stream_response::Payload::StatusUpdate(
event.try_into().unwrap(),
)),
}
}
#[tokio::test]
async fn grpc_stream_reader_task_forwards_typed_event_as_sse() {
let payloads = vec![Ok(status_update_event())];
let stream = tonic::codegen::tokio_stream::iter(payloads);
let (tx, mut rx) = mpsc::channel::<crate::streaming::event_stream::BodyChunk>(8);
grpc_stream_reader_task(stream, tx).await;
let first = rx.recv().await.expect("expected one chunk");
let bytes = first.expect("expected Ok chunk");
let text = std::str::from_utf8(&bytes).expect("utf8");
assert!(
text.starts_with("data: "),
"chunk must be SSE-framed: {text}"
);
assert!(
text.contains("\"jsonrpc\":\"2.0\""),
"chunk must be JSON-RPC envelope: {text}"
);
assert!(
text.contains("\"statusUpdate\""),
"typed event must serialize as the domain union: {text}"
);
assert!(
text.contains("TASK_STATE_WORKING"),
"state must use canonical wire encoding: {text}"
);
assert!(rx.recv().await.is_none());
}
#[tokio::test]
async fn grpc_stream_reader_task_forwards_multiple_payloads() {
let payloads = vec![
Ok(status_update_event()),
Ok(status_update_event()),
Ok(status_update_event()),
];
let stream = tonic::codegen::tokio_stream::iter(payloads);
let (tx, mut rx) = mpsc::channel::<crate::streaming::event_stream::BodyChunk>(8);
grpc_stream_reader_task(stream, tx).await;
let mut received = 0;
while let Some(item) = rx.recv().await {
assert!(item.is_ok());
received += 1;
}
assert_eq!(received, 3, "all three payloads must be forwarded");
}
#[tokio::test]
async fn grpc_stream_reader_task_maps_status_error_to_protocol_error() {
let payloads: Vec<Result<apb::StreamResponse, tonic::Status>> =
vec![Err(tonic::Status::not_found("missing"))];
let stream = tonic::codegen::tokio_stream::iter(payloads);
let (tx, mut rx) = mpsc::channel::<crate::streaming::event_stream::BodyChunk>(8);
grpc_stream_reader_task(stream, tx).await;
let chunk = rx.recv().await.expect("expected an error chunk");
match chunk {
Err(ClientError::Protocol(a2a)) => {
assert_eq!(a2a.code, a2a_protocol_types::ErrorCode::TaskNotFound);
assert!(a2a.message.contains("missing"));
}
other => panic!("expected Protocol(TaskNotFound), got {other:?}"),
}
}
#[tokio::test]
async fn grpc_stream_reader_task_rejects_empty_payload() {
let payloads = vec![Ok(apb::StreamResponse { payload: None })];
let stream = tonic::codegen::tokio_stream::iter(payloads);
let (tx, mut rx) = mpsc::channel::<crate::streaming::event_stream::BodyChunk>(8);
grpc_stream_reader_task(stream, tx).await;
let chunk = rx.recv().await.expect("expected an error chunk");
match chunk {
Err(ClientError::Transport(msg)) => {
assert!(
msg.contains("streamResponse.payload"),
"msg should name the field: {msg}"
);
}
other => panic!("expected Transport error, got {other:?}"),
}
}
#[tokio::test]
async fn grpc_transport_endpoint_returns_input_url() {
let endpoint_str = "http://localhost:50055".to_string();
let channel = tonic::transport::Channel::from_shared(endpoint_str.clone())
.expect("valid endpoint")
.connect_lazy();
let transport = GrpcTransport {
inner: Arc::new(Inner {
channel,
endpoint: endpoint_str.clone(),
config: GrpcTransportConfig::default(),
}),
};
assert_eq!(transport.endpoint(), endpoint_str);
}
#[tokio::test]
async fn grpc_transport_endpoint_preserves_distinct_urls() {
let a = "http://example.com:1234".to_string();
let b = "https://other.test:9000".to_string();
let mk = |s: String| {
let ch = tonic::transport::Channel::from_shared(s.clone())
.unwrap()
.connect_lazy();
GrpcTransport {
inner: Arc::new(Inner {
channel: ch,
endpoint: s,
config: GrpcTransportConfig::default(),
}),
}
};
let ta = mk(a.clone());
let tb = mk(b.clone());
assert_eq!(ta.endpoint(), a);
assert_eq!(tb.endpoint(), b);
assert_ne!(ta.endpoint(), tb.endpoint());
}
}